fix: file update tool results

Author: charleslien

backend/src/run-agent-step.ts

@@ -36,7 +36,10 @@ import { countTokensJson } from './util/token-counter'
 import { getRequestContext } from './websockets/request-context'
 
 import type { AgentResponseTrace } from '@codebuff/bigquery'
-import type { CodebuffToolMessage } from '@codebuff/common/tools/list'
+import type {
+  CodebuffToolMessage,
+  CodebuffToolOutput,
+} from '@codebuff/common/tools/list'
 import type { AgentTemplate } from '@codebuff/common/types/agent-template'
 import type {
   AssistantMessage,
@@ -195,7 +198,7 @@ export const runAgentStep = async (
             ),
           },
         },
-      ],
+      ] satisfies CodebuffToolOutput<'file_updates'>,
     })
   }
 

backend/src/util/messages.ts

@@ -297,36 +297,71 @@ export function getPreviouslyReadFiles(messages: Message[]): {
   content: string
   referencedBy?: Record<string, string[]>
 }[] {
-  return buildArray(
-    messages
-      .filter(
-        (
-          m,
-        ): m is ToolMessage & {
-          content: { toolName: 'read_files' }
-        } =>
-          m.role === 'tool' &&
-          (m.content.toolName === 'read_files' ||
-            m.content.toolName === 'find_files' ||
-            m.content.toolName === 'file_updates'),
-      )
-      .map((m) => {
-        try {
-          return (
-            m as CodebuffToolMessage<'read_files'>
-          ).content.output[0].value.map((file) => {
-            if ('contentOmittedForLength' in file) {
-              return undefined
-            }
-            return file
-          })
-        } catch (error) {
-          logger.error(
-            { error: errorToObject(error), m },
-            'Error parsing read_files output from message',
-          )
-          return []
+  const files: ReturnType<typeof getPreviouslyReadFiles> = []
+  for (const message of messages) {
+    if (message.role !== 'tool') continue
+    if (message.content.toolName === 'read_files') {
+      try {
+        files.push(
+          ...(
+            message as CodebuffToolMessage<'read_files'>
+          ).content.output[0].value.filter(
+            (
+              file,
+            ): file is typeof file & { contentOmittedForLength: undefined } =>
+              !('contentOmittedForLength' in file),
+          ),
+        )
+      } catch (error) {
+        logger.error(
+          { error: errorToObject(error), message },
+          'Error parsing read_files output from message',
+        )
+      }
+    }
+
+    if (message.content.toolName === 'find_files') {
+      try {
+        const v = (message as CodebuffToolMessage<'find_files'>).content
+          .output[0].value
+        if ('message' in v) {
+          continue
         }
-      }),
-  )
+        files.push(
+          ...v.filter(
+            (
+              file,
+            ): file is typeof file & { contentOmittedForLength: undefined } =>
+              !('contentOmittedForLength' in file),
+          ),
+        )
+      } catch (error) {
+        logger.error(
+          { error: errorToObject(error), message },
+          'Error parsing find_files output from message',
+        )
+      }
+    }
+
+    if (message.content.toolName === 'file_updates') {
+      try {
+        files.push(
+          ...(
+            message as CodebuffToolMessage<'file_updates'>
+          ).content.output[0].value.files.filter(
+            (
+              file,
+            ): file is typeof file & { contentOmittedForLength: undefined } =>
+              !('contentOmittedForLength' in file),
+          ),
+        )
+      } catch (error) {
+        logger.error(
+          { error: errorToObject(error), message },
+          'Error parsing find_files output from message',
+        )
+      }
+    }
+  }
+  return files
 }

common/src/tools/constants.ts

@@ -71,3 +71,8 @@ export type $ToolParams<T extends ToolName = ToolName> = {
   parameters: z.ZodType
   outputs: z.ZodType<ToolResultOutput[]>
 }
+
+export type $ToolResults = {
+  toolName: string
+  outputs: $ToolParams['outputs']
+}

common/src/tools/list.ts

@@ -7,6 +7,7 @@ import { browserLogsParams } from './params/tool/browser-logs'
 import { codeSearchParams } from './params/tool/code-search'
 import { createPlanParams } from './params/tool/create-plan'
 import { endTurnParams } from './params/tool/end-turn'
+import { fileUpdatesResultSchema } from './params/tool/file-updates'
 import { findFilesParams } from './params/tool/find-files'
 import { readDocsParams } from './params/tool/read-docs'
 import { readFilesParams } from './params/tool/read-files'
@@ -23,7 +24,12 @@ import { updateSubgoalParams } from './params/tool/update-subgoal'
 import { webSearchParams } from './params/tool/web-search'
 import { writeFileParams } from './params/tool/write-file'
 
-import type { ToolName, $ToolParams, PublishedToolName } from './constants'
+import type {
+  $ToolParams,
+  $ToolResults,
+  PublishedToolName,
+  ToolName,
+} from './constants'
 import type { ToolMessage } from '../types/messages/codebuff-message'
 import type {
   ToolCallPart,
@@ -56,6 +62,11 @@ export const $toolParams = {
   [K in ToolName]: $ToolParams<K>
 }
 
+export const additionalToolResultSchemas = {
+  file_updates: fileUpdatesResultSchema,
+} satisfies Record<string, $ToolResults>
+type ResultOnlyToolName = keyof typeof additionalToolResultSchemas
+
 // Tool call from LLM
 export type CodebuffToolCall<T extends ToolName = ToolName> = {
   [K in ToolName]: {
@@ -64,19 +75,33 @@ export type CodebuffToolCall<T extends ToolName = ToolName> = {
   } & Omit<ToolCallPart, 'type'>
 }[T]
 
-export type CodebuffToolOutput<T extends ToolName = ToolName> = {
-  [K in ToolName]: z.infer<(typeof $toolParams)[K]['outputs']>
-}[T]
-export type CodebuffToolResult<T extends ToolName = ToolName> = {
-  [K in ToolName]: {
+export type CodebuffToolOutput<
+  T extends ToolName | ResultOnlyToolName = ToolName,
+> = T extends ToolName
+  ? {
+      [K in ToolName]: z.infer<(typeof $toolParams)[K]['outputs']>
+    }[T]
+  : T extends ResultOnlyToolName
+    ? {
+        [K in ResultOnlyToolName]: z.infer<
+          (typeof additionalToolResultSchemas)[K]['outputs']
+        >
+      }[T]
+    : never
+export type CodebuffToolResult<
+  T extends ToolName | ResultOnlyToolName = ToolName,
+> = {
+  [K in ToolName | ResultOnlyToolName]: {
     toolName: K
     output: CodebuffToolOutput<K>
   } & Omit<ToolResultPart, 'type'>
 }[T]
 
-export type CodebuffToolMessage<T extends ToolName = ToolName> = ToolMessage &
+export type CodebuffToolMessage<
+  T extends ToolName | ResultOnlyToolName = ToolName,
+> = ToolMessage &
   {
-    [K in ToolName]: {
+    [K in ToolName | ResultOnlyToolName]: {
       toolName: K
       content: {
         output: CodebuffToolOutput<K>
@@ -126,4 +151,4 @@ export type ClientToolCall<T extends ClientToolName = ClientToolName> = z.infer<
   typeof clientToolCallSchema
 > & { toolName: T } & Omit<ToolCallPart, 'type'>
 
-export type PublishedClientToolName = ClientToolName & PublishedToolName
+export type PublishedClientToolName = ClientToolName & PublishedToolName

common/src/tools/params/tool/file-updates.ts

@@ -0,0 +1,23 @@
+import z from 'zod/v4'
+
+import type { $ToolResults } from '../../../tools/constants'
+
+const toolName = 'file_updates'
+export const fileUpdatesResultSchema = {
+  toolName: toolName,
+  outputs: z.tuple([
+    z.object({
+      type: z.literal('json'),
+      value: z.object({
+        message: z.string().optional(),
+        files: z.array(
+          z.object({
+            path: z.string(),
+            content: z.string(),
+            referencedBy: z.record(z.string(), z.array(z.string())),
+          }),
+        ),
+      }),
+    }),
+  ]),
+} as const satisfies $ToolResults