streaming fix on non done tool responses

packages/ai/src/tools/communication-tools/message-user-clarifying-question/message-user-clarifying-question-execute.ts

@@ -1,4 +1,4 @@
-import { updateMessageEntries } from '@buster/database/queries';
+import { updateMessage, updateMessageEntries } from '@buster/database/queries';
 import { wrapTraced } from 'braintrust';
 import { cleanupState } from '../../shared/cleanup-state';
 import { createRawToolResultEntry } from '../../shared/create-raw-llm-tool-result-entry';
@@ -39,6 +39,11 @@ async function processMessageUserClarifyingQuestion(
       messageId,
       rawLlmMessages,
     });
+
+    // Mark the message as completed
+    await updateMessage(messageId, {
+      isCompleted: true,
+    });
   } catch (error) {
     console.error('[message-user-clarifying-question] Error updating message entries:', error);
   }
@@ -53,15 +58,18 @@ export function createMessageUserClarifyingQuestionExecute(
 ) {
   return wrapTraced(
     async (
-      _input: MessageUserClarifyingQuestionInput
+      _input: MessageUserClarifyingQuestionInput,
+      options?: { toolCallId?: string }
     ): Promise<MessageUserClarifyingQuestionOutput> => {
-      if (!state.toolCallId) {
+      // Use toolCallId from state if available, otherwise from options
+      const toolCallId = state.toolCallId || options?.toolCallId;
+      if (!toolCallId) {
         throw new Error('Tool call ID is required');
       }
 
       const result = await processMessageUserClarifyingQuestion(
         state,
-        state.toolCallId,
+        toolCallId,
         context.messageId
       );
       cleanupState(state);

packages/ai/src/tools/communication-tools/message-user-clarifying-question/message-user-clarifying-question-start.ts

@@ -66,7 +66,6 @@ export function createMessageUserClarifyingQuestionStart(
         }
 
         await updateMessage(context.messageId, {
-          isCompleted: true,
           finalReasoningMessage: `Reasoned for ${timeString}`,
         });
       }

packages/ai/src/tools/communication-tools/respond-without-asset-creation/respond-without-asset-creation-execute.ts

@@ -1,4 +1,4 @@
-import { updateMessageEntries } from '@buster/database/queries';
+import { updateMessage, updateMessageEntries } from '@buster/database/queries';
 import { wrapTraced } from 'braintrust';
 import { cleanupState } from '../../shared/cleanup-state';
 import { createRawToolResultEntry } from '../../shared/create-raw-llm-tool-result-entry';
@@ -38,6 +38,11 @@ async function processRespondWithoutAssetCreation(
       messageId,
       rawLlmMessages,
     });
+
+    // Mark the message as completed
+    await updateMessage(messageId, {
+      isCompleted: true,
+    });
   } catch (error) {
     console.error('[respond-without-asset-creation] Error updating message entries:', error);
   }
@@ -52,17 +57,16 @@ export function createRespondWithoutAssetCreationExecute(
 ) {
   return wrapTraced(
     async (
-      _input: RespondWithoutAssetCreationInput
+      _input: RespondWithoutAssetCreationInput,
+      options?: { toolCallId?: string }
     ): Promise<RespondWithoutAssetCreationOutput> => {
-      if (!state.toolCallId) {
+      // Use toolCallId from state if available, otherwise from options
+      const toolCallId = state.toolCallId || options?.toolCallId;
+      if (!toolCallId) {
         throw new Error('Tool call ID is required');
       }
 
-      const result = await processRespondWithoutAssetCreation(
-        state,
-        state.toolCallId,
-        context.messageId
-      );
+      const result = await processRespondWithoutAssetCreation(state, toolCallId, context.messageId);
       cleanupState(state);
       return result;
     },

packages/ai/src/tools/communication-tools/respond-without-asset-creation/respond-without-asset-creation-start.ts

@@ -64,7 +64,6 @@ export function createRespondWithoutAssetCreationStart(
         }
 
         await updateMessage(context.messageId, {
-          isCompleted: true,
           finalReasoningMessage: `Reasoned for ${timeString}`,
         });
       }

packages/ai/src/tools/communication-tools/respond-without-asset-creation/respond-without-asset-creation.int.test.ts

@@ -11,6 +11,12 @@ vi.mock('braintrust', () => ({
   wrapTraced: (fn: unknown) => fn,
 }));
 
+// Mock database operations for tests
+vi.mock('@buster/database/queries', () => ({
+  updateMessage: vi.fn().mockResolvedValue({ success: true }),
+  updateMessageEntries: vi.fn().mockResolvedValue({ success: true }),
+}));
+
 describe('Respond Without Asset Creation Tool Integration Tests', () => {
   describe('Tool Creation and Configuration', () => {
     test('should create tool with minimal context', () => {
@@ -130,7 +136,7 @@ describe('Respond Without Asset Creation Tool Integration Tests', () => {
         } as ToolCallOptions);
 
         expect(result).toBeDefined();
-        expect(result).toEqual({});
+        expect(result).toEqual({ success: true });
       }
     });
 
@@ -152,7 +158,7 @@ describe('Respond Without Asset Creation Tool Integration Tests', () => {
         } as ToolCallOptions);
 
         expect(result).toBeDefined();
-        expect(result).toEqual({});
+        expect(result).toEqual({ success: true });
       }
     });
 
@@ -177,7 +183,7 @@ Escaped: \n \t \\ \" \'
         } as ToolCallOptions);
 
         expect(result).toBeDefined();
-        expect(result).toEqual({});
+        expect(result).toEqual({ success: true });
       }
     });
 
@@ -198,7 +204,7 @@ Escaped: \n \t \\ \" \'
         } as ToolCallOptions);
 
         expect(result).toBeDefined();
-        expect(result).toEqual({});
+        expect(result).toEqual({ success: true });
       }
     });
   });
@@ -238,28 +244,31 @@ Escaped: \n \t \\ \" \'
           tool2.execute(input2, { toolCallId: 'tc-2', messages: [] } as ToolCallOptions),
         ]);
 
-        expect(result1).toEqual({});
-        expect(result2).toEqual({});
+        expect(result1).toEqual({ success: true });
+        expect(result2).toEqual({ success: true });
       }
     });
   });
 
   describe('Output Schema Validation', () => {
-    test('should have empty output schema', () => {
+    test('should validate output schema with success field', () => {
       const outputSchema = RespondWithoutAssetCreationOutputSchema;
 
-      const emptyOutput = {};
-      expect(() => outputSchema.parse(emptyOutput)).not.toThrow();
+      const validOutput = { success: true };
+      expect(() => outputSchema.parse(validOutput)).not.toThrow();
     });
 
-    test('should accept empty output only', () => {
+    test('should require success field in output', () => {
       const outputSchema = RespondWithoutAssetCreationOutputSchema;
 
-      // The output schema is z.object({}) which accepts empty objects
-      // Additional properties are allowed in Zod by default unless strict() is used
+      // The output schema requires a success boolean field
+      const validOutput = { success: true };
+      const result = outputSchema.parse(validOutput);
+      expect(result).toEqual({ success: true });
+
+      // Should throw for empty object
       const emptyOutput = {};
-      const result = outputSchema.parse(emptyOutput);
-      expect(result).toEqual({});
+      expect(() => outputSchema.parse(emptyOutput)).toThrow();
     });
   });