diff --git a/packages/ai/src/tools/communication-tools/done-tool/done-tool-execute.ts b/packages/ai/src/tools/communication-tools/done-tool/done-tool-execute.ts index 1f6204969..95353d52a 100644 --- a/packages/ai/src/tools/communication-tools/done-tool/done-tool-execute.ts +++ b/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 { - return { +async function processDone(toolCallId: string, messageId: string): Promise { + 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 => { - return processDone(); + if (!state.toolCallId) { + throw new Error('Tool call ID is required'); + } + + return processDone(state.toolCallId, context.messageId); }, { name: 'Done Tool' } ); diff --git a/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts b/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts index 1e23d5564..f75c56e04 100644 --- a/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts +++ b/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); diff --git a/packages/ai/src/tools/shared/create-raw-llm-tool-result-entry.ts b/packages/ai/src/tools/shared/create-raw-llm-tool-result-entry.ts new file mode 100644 index 000000000..fe2c8119a --- /dev/null +++ b/packages/ai/src/tools/shared/create-raw-llm-tool-result-entry.ts @@ -0,0 +1,22 @@ +import type { ModelMessage } from 'ai'; + +export function createRawToolResultEntry( + toolCallId: string, + toolName: string, + output: T +): ModelMessage { + return { + role: 'tool', + content: [ + { + type: 'tool-result', + toolCallId, + toolName, + output: { + type: 'json', + value: JSON.stringify(output), + }, + }, + ], + }; +}