done tool state and context

packages/ai/src/tools/communication-tools/done-tool/done-tool-execute.ts

@@ -1,18 +1,43 @@
+import { updateMessageEntries } from '@buster/database';
 import { wrapTraced } from 'braintrust';
-import type { DoneToolContext, DoneToolInput, DoneToolOutput } from './done-tool';
+import { createRawToolResultEntry } from '../../shared/create-raw-llm-tool-result-entry';
+import {
+  DONE_TOOL_NAME,
+  type DoneToolContext,
+  type DoneToolInput,
+  type DoneToolOutput,
+  type DoneToolState,
+} from './done-tool';
 
 // Process done tool execution with todo management
-async function processDone(): Promise<DoneToolOutput> {
-  return {
+async function processDone(toolCallId: string, messageId: string): Promise<DoneToolOutput> {
+  const output: DoneToolOutput = {
     success: true,
   };
+
+  const rawToolResultEntry = createRawToolResultEntry(toolCallId, DONE_TOOL_NAME, output);
+
+  try {
+    await updateMessageEntries({
+      messageId,
+      rawLlmMessages: [rawToolResultEntry],
+    });
+  } catch (error) {
+    console.error('[done-tool] Error updating message entries:', error);
+  }
+
+  return output;
 }
 
 // Factory function that creates the execute function with proper context typing
-export function createDoneToolExecute() {
+export function createDoneToolExecute(state: DoneToolState, context: DoneToolContext) {
   return wrapTraced(
     async (_input: DoneToolInput): Promise<DoneToolOutput> => {
-      return processDone();
+      if (!state.toolCallId) {
+        throw new Error('Tool call ID is required');
+      }
+
+      return processDone(state.toolCallId, context.messageId);
     },
     { name: 'Done Tool' }
   );

packages/ai/src/tools/communication-tools/done-tool/done-tool.ts

@@ -53,7 +53,7 @@ export function createDoneTool(context: DoneToolContext) {
     finalResponse: undefined,
   };
 
-  const execute = createDoneToolExecute();
+  const execute = createDoneToolExecute(state, context);
   const onInputStart = createDoneToolStart(state, context);
   const onInputDelta = createDoneToolDelta(state, context);
   const onInputAvailable = createDoneToolFinish(state, context);

packages/ai/src/tools/shared/create-raw-llm-tool-result-entry.ts

@@ -0,0 +1,22 @@
+import type { ModelMessage } from 'ai';
+
+export function createRawToolResultEntry<T>(
+  toolCallId: string,
+  toolName: string,
+  output: T
+): ModelMessage {
+  return {
+    role: 'tool',
+    content: [
+      {
+        type: 'tool-result',
+        toolCallId,
+        toolName,
+        output: {
+          type: 'json',
+          value: JSON.stringify(output),
+        },
+      },
+    ],
+  };
+}