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 b29211d38..fbd7062b7 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,4 +1,8 @@ -import { updateMessage, updateMessageEntries, waitForPendingUpdates } from '@buster/database/queries'; +import { + updateMessage, + updateMessageEntries, + waitForPendingUpdates, +} from '@buster/database/queries'; import { wrapTraced } from 'braintrust'; import { cleanupState } from '../../shared/cleanup-state'; import { createRawToolResultEntry } from '../../shared/create-raw-llm-tool-result-entry'; @@ -57,10 +61,10 @@ export function createDoneToolExecute(context: DoneToolContext, state: DoneToolS } const result = await processDone(state, state.toolCallId, context.messageId, context); - + // Wait for all pending updates from delta/finish to complete before returning await waitForPendingUpdates(context.messageId); - + cleanupState(state); return result; },