Store the proposed files per runId

Author: jahooma

packages/agent-runtime/src/__tests__/propose-tools.test.ts

@@ -19,6 +19,7 @@ import {
   clearAgentGeneratorCache,
   runProgrammaticStep,
 } from '../run-programmatic-step'
+import { clearAllProposedContent } from '../tools/handlers/tool/proposed-content-store'
 import { mockFileContext } from './test-utils'
 import * as toolExecutor from '../tools/tool-executor'
 
@@ -268,6 +269,7 @@ console.log(add(1, 2));
   afterEach(() => {
     mock.restore()
     clearAgentGeneratorCache({ logger })
+    clearAllProposedContent()
   })
 
   describe('propose_str_replace', () => {

packages/agent-runtime/src/run-programmatic-step.ts

@@ -2,6 +2,7 @@ import { getErrorObject } from '@codebuff/common/util/error'
 import { assistantMessage } from '@codebuff/common/util/messages'
 import { cloneDeep } from 'lodash'
 
+import { clearProposedContentForRun } from './tools/handlers/tool/proposed-content-store'
 import { executeToolCall } from './tools/tool-executor'
 import { parseTextWithToolCalls } from './util/parse-tool-calls-from-text'
 
@@ -38,6 +39,7 @@ export const runIdToStepAll: Set<string> = new Set()
 // Function to clear the generator cache for testing purposes
 export function clearAgentGeneratorCache(params: { logger: Logger }) {
   for (const key in runIdToGenerator) {
+    clearProposedContentForRun(key)
     delete runIdToGenerator[key]
   }
   runIdToStepAll.clear()
@@ -190,7 +192,6 @@ export async function runProgrammaticStep(
     fileChangeErrors: [],
     fileChanges: [],
     firstFileProcessed: false,
-    proposedContentByPath: {},
   }
   const agentContext = cloneDeep(agentState.agentContext)
   const sendSubagentChunk = (data: {
@@ -378,6 +379,7 @@ export async function runProgrammaticStep(
     if (endTurn) {
       delete runIdToGenerator[agentState.runId]
       runIdToStepAll.delete(agentState.runId)
+      clearProposedContentForRun(agentState.runId)
     }
   }
 }

packages/agent-runtime/src/tools/handlers/tool/propose-str-replace.ts

@@ -1,7 +1,10 @@
 import { processStrReplace } from '../../../process-str-replace'
+import {
+  getProposedContent,
+  setProposedContent,
+} from './proposed-content-store'
 
 import type { CodebuffToolHandlerFunction } from '../handler-function-type'
-import type { FileProcessingState } from './write-file'
 import type {
   CodebuffToolCall,
   CodebuffToolOutput,
@@ -16,9 +19,9 @@ export const handleProposeStrReplace = (async (
     previousToolCallFinished: Promise<void>
     toolCall: CodebuffToolCall<'propose_str_replace'>
 
-    fileProcessingState: FileProcessingState
     logger: Logger
     agentState: AgentState
+    runId: string
 
     requestOptionalFile: RequestOptionalFileFn
   } & ParamsExcluding<RequestOptionalFileFn, 'filePath'>,
@@ -27,16 +30,16 @@ export const handleProposeStrReplace = (async (
     previousToolCallFinished,
     toolCall,
 
-    fileProcessingState,
     logger,
+    runId,
 
     requestOptionalFile,
   } = params
   const { path, replacements } = toolCall.input
 
-  // Get content from proposed state first, then fall back to disk
+  // Get content from proposed state first (by runId), then fall back to disk
   const getProposedOrDiskContent = async (): Promise<string | null> => {
-    const proposedContent = fileProcessingState.proposedContentByPath[path]
+    const proposedContent = getProposedContent(runId, path)
     if (proposedContent !== undefined) {
       return proposedContent
     }
@@ -59,9 +62,13 @@ export const handleProposeStrReplace = (async (
     }
   })
 
-  // Store the proposed content for future propose calls on the same file
-  fileProcessingState.proposedContentByPath[path] = strReplaceResultPromise.then(
-    (result) => ('content' in result ? result.content : null),
+  // Store the proposed content for future propose calls on the same file (by runId)
+  setProposedContent(
+    runId,
+    path,
+    strReplaceResultPromise.then((result) =>
+      'content' in result ? result.content : null,
+    ),
   )
 
   await previousToolCallFinished

packages/agent-runtime/src/tools/handlers/tool/propose-write-file.ts

@@ -1,7 +1,11 @@
 import { createPatch } from 'diff'
 
+import {
+  getProposedContent,
+  setProposedContent,
+} from './proposed-content-store'
+
 import type { CodebuffToolHandlerFunction } from '../handler-function-type'
-import type { FileProcessingState } from './write-file'
 import type {
   CodebuffToolCall,
   CodebuffToolOutput,
@@ -20,24 +24,24 @@ export const handleProposeWriteFile = (async (
     previousToolCallFinished: Promise<void>
     toolCall: CodebuffToolCall<'propose_write_file'>
 
-    fileProcessingState: FileProcessingState
     logger: Logger
+    runId: string
 
     requestOptionalFile: RequestOptionalFileFn
   } & ParamsExcluding<RequestOptionalFileFn, 'filePath'>,
 ): Promise<{ output: CodebuffToolOutput<'propose_write_file'> }> => {
   const {
     previousToolCallFinished,
     toolCall,
-    fileProcessingState,
     logger,
+    runId,
     requestOptionalFile,
   } = params
   const { path, content } = toolCall.input
 
-  // Get content from proposed state first, then fall back to disk
+  // Get content from proposed state first (by runId), then fall back to disk
   const getProposedOrDiskContent = async (): Promise<string | null> => {
-    const proposedContent = fileProcessingState.proposedContentByPath[path]
+    const proposedContent = getProposedContent(runId, path)
     if (proposedContent !== undefined) {
       return proposedContent
     }
@@ -49,8 +53,8 @@ export const handleProposeWriteFile = (async (
   // Normalize content (remove leading newline if present)
   const newContent = content.startsWith('\n') ? content.slice(1) : content
 
-  // Store the proposed content for future propose calls on the same file
-  fileProcessingState.proposedContentByPath[path] = Promise.resolve(newContent)
+  // Store the proposed content for future propose calls on the same file (by runId)
+  setProposedContent(runId, path, Promise.resolve(newContent))
 
   await previousToolCallFinished
 

packages/agent-runtime/src/tools/handlers/tool/proposed-content-store.ts

@@ -0,0 +1,64 @@
+/**
+ * Store for proposed file content by runId.
+ * This allows propose_str_replace and propose_write_file tools to
+ * track proposed changes within an agent run, isolated by runId.
+ */
+
+/** Map of runId -> path -> Promise<content | null> */
+const proposedContentByRunId = new Map<
+  string,
+  Record<string, Promise<string | null>>
+>()
+
+/**
+ * Get the proposed content map for a specific runId.
+ * Creates an empty record if none exists.
+ */
+export function getProposedContentForRun(
+  runId: string,
+): Record<string, Promise<string | null>> {
+  let contentByPath = proposedContentByRunId.get(runId)
+  if (!contentByPath) {
+    contentByPath = {}
+    proposedContentByRunId.set(runId, contentByPath)
+  }
+  return contentByPath
+}
+
+/**
+ * Get proposed content for a specific file in a run.
+ */
+export function getProposedContent(
+  runId: string,
+  path: string,
+): Promise<string | null> | undefined {
+  const contentByPath = proposedContentByRunId.get(runId)
+  return contentByPath?.[path]
+}
+
+/**
+ * Set proposed content for a specific file in a run.
+ */
+export function setProposedContent(
+  runId: string,
+  path: string,
+  content: Promise<string | null>,
+): void {
+  const contentByPath = getProposedContentForRun(runId)
+  contentByPath[path] = content
+}
+
+/**
+ * Clear all proposed content for a specific runId.
+ * Should be called when an agent run completes.
+ */
+export function clearProposedContentForRun(runId: string): void {
+  proposedContentByRunId.delete(runId)
+}
+
+/**
+ * Clear all proposed content (for testing purposes).
+ */
+export function clearAllProposedContent(): void {
+  proposedContentByRunId.clear()
+}

packages/agent-runtime/src/tools/handlers/tool/write-file.ts

@@ -37,8 +37,6 @@ export type FileProcessingState = {
   fileChangeErrors: Extract<FileProcessing, { error: string }>[]
   fileChanges: Exclude<FileProcessing, { error: string }>[]
   firstFileProcessed: boolean
-  /** Tracks proposed file contents for propose_* tools (not written to disk) */
-  proposedContentByPath: Record<string, Promise<string | null>>
 }
 
 export function getFileProcessingValues(
@@ -50,7 +48,6 @@ export function getFileProcessingValues(
     fileChangeErrors: [],
     fileChanges: [],
     firstFileProcessed: false,
-    proposedContentByPath: {},
   }
   for (const [key, value] of Object.entries(state)) {
     const typedKey = key as keyof typeof fileProcessingValues

packages/agent-runtime/src/tools/stream-parser.ts

@@ -96,7 +96,6 @@ export async function processStream(
     fileChangeErrors: [],
     fileChanges: [],
     firstFileProcessed: false,
-    proposedContentByPath: {},
   }
 
   function toolCallback<T extends ToolName>(toolName: T) {