🤖 refactor: simplify model message prep pipeline (#594)

ammar-agent · web-flow · commit 4e97ff94efc5 · 2025-11-14T12:31:47.000-06:00
## Summary
- rely on Vercel AI SDK's convertToModelMessages
ignoreIncompleteToolCalls option so we no longer hand-edit partial tool
calls
- drop the bespoke splitMixedContentMessages pass and streamline the
provider transform logic
- update the unit tests to reflect the simpler pipeline while keeping
reasoning/tool coverage

## Testing
- make typecheck
- make test

_Generated with `mux`_
diff --git a/src/services/aiService.ts b/src/services/aiService.ts
@@ -599,7 +599,11 @@ export class AIService extends EventEmitter {
       // Convert MuxMessage to ModelMessage format using Vercel AI SDK utility
       // Type assertion needed because MuxMessage has custom tool parts for interrupted tools
       // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-argument
-      const modelMessages = convertToModelMessages(sanitizedMessages as any);
+      const modelMessages = convertToModelMessages(sanitizedMessages as any, {
+        // Drop unfinished tool calls (input-streaming/input-available) so downstream
+        // transforms only see tool calls that actually produced outputs.
+        ignoreIncompleteToolCalls: true,
+      });
       log.debug_obj(`${workspaceId}/2_model_messages.json`, modelMessages);
 
       // Apply ModelMessage transforms based on provider requirements
diff --git a/src/utils/messages/modelMessageTransform.test.ts b/src/utils/messages/modelMessageTransform.test.ts
@@ -29,65 +29,13 @@ describe("modelMessageTransform", () => {
       expect(result).toEqual(messages);
     });
 
-    it("should keep text-only messages unchanged", () => {
-      const assistantMsg1: AssistantModelMessage = {
-        role: "assistant",
-        content: [{ type: "text", text: "Let me help you with that." }],
-      };
-      const assistantMsg2: AssistantModelMessage = {
-        role: "assistant",
-        content: [{ type: "text", text: "Here's the result." }],
-      };
-      const messages: ModelMessage[] = [assistantMsg1, assistantMsg2];
-
-      const result = transformModelMessages(messages, "anthropic");
-      expect(result).toEqual(messages);
-    });
-
-    it("should strip tool calls without results (interrupted mixed content)", () => {
+    it("should split mixed text and tool-call content into ordered segments", () => {
       const assistantMsg: AssistantModelMessage = {
         role: "assistant",
         content: [
-          { type: "text", text: "Let me check that for you." },
-          { type: "tool-call", toolCallId: "call1", toolName: "bash", input: { script: "ls" } },
-        ],
-      };
-      const messages: ModelMessage[] = [assistantMsg];
-
-      const result = transformModelMessages(messages, "anthropic");
-
-      // Should only keep text, strip interrupted tool calls
-      expect(result).toHaveLength(1);
-      expect(result[0].role).toBe("assistant");
-      expect((result[0] as AssistantModelMessage).content).toEqual([
-        { type: "text", text: "Let me check that for you." },
-      ]);
-    });
-
-    it("should strip tool-only messages without results (orphaned tool calls)", () => {
-      const assistantMsg: AssistantModelMessage = {
-        role: "assistant",
-        content: [
-          { type: "tool-call", toolCallId: "call1", toolName: "bash", input: { script: "ls" } },
-        ],
-      };
-      const messages: ModelMessage[] = [assistantMsg];
-
-      const result = transformModelMessages(messages, "anthropic");
-
-      // Should filter out the entire message since it only has orphaned tool calls
-      expect(result).toHaveLength(0);
-    });
-
-    it("should handle partial results (some tool calls interrupted)", () => {
-      // Assistant makes 3 tool calls, but only 2 have results (3rd was interrupted)
-      const assistantMsg: AssistantModelMessage = {
-        role: "assistant",
-        content: [
-          { type: "text", text: "Let me check a few things." },
+          { type: "text", text: "Before" },
           { type: "tool-call", toolCallId: "call1", toolName: "bash", input: { script: "pwd" } },
-          { type: "tool-call", toolCallId: "call2", toolName: "bash", input: { script: "ls" } },
-          { type: "tool-call", toolCallId: "call3", toolName: "bash", input: { script: "date" } },
+          { type: "text", text: "After" },
         ],
       };
       const toolMsg: ToolModelMessage = {
@@ -99,80 +47,41 @@ describe("modelMessageTransform", () => {
             toolName: "bash",
             output: { type: "json", value: { stdout: "/home/user" } },
           },
-          {
-            type: "tool-result",
-            toolCallId: "call2",
-            toolName: "bash",
-            output: { type: "json", value: { stdout: "file1 file2" } },
-          },
-          // call3 has no result (interrupted)
         ],
       };
-      const messages: ModelMessage[] = [assistantMsg, toolMsg];
 
-      const result = transformModelMessages(messages, "anthropic");
+      const result = transformModelMessages([assistantMsg, toolMsg], "anthropic");
 
-      // Should have: text message, tool calls (only call1 & call2), tool results
-      expect(result).toHaveLength(3);
-
-      // First: text
+      expect(result).toHaveLength(4);
       expect(result[0].role).toBe("assistant");
       expect((result[0] as AssistantModelMessage).content).toEqual([
-        { type: "text", text: "Let me check a few things." },
+        { type: "text", text: "Before" },
       ]);
-
-      // Second: only tool calls with results (call1, call2), NOT call3
       expect(result[1].role).toBe("assistant");
-      const toolCallContent = (result[1] as AssistantModelMessage).content;
-      expect(Array.isArray(toolCallContent)).toBe(true);
-      if (Array.isArray(toolCallContent)) {
-        expect(toolCallContent).toHaveLength(2);
-        expect(toolCallContent[0]).toEqual({
-          type: "tool-call",
-          toolCallId: "call1",
-          toolName: "bash",
-          input: { script: "pwd" },
-        });
-        expect(toolCallContent[1]).toEqual({
-          type: "tool-call",
-          toolCallId: "call2",
-          toolName: "bash",
-          input: { script: "ls" },
-        });
-      }
-
-      // Third: tool results (only for call1 & call2)
+      expect((result[1] as AssistantModelMessage).content).toEqual([
+        { type: "tool-call", toolCallId: "call1", toolName: "bash", input: { script: "pwd" } },
+      ]);
       expect(result[2].role).toBe("tool");
-      const toolResultContent = (result[2] as ToolModelMessage).content;
-      expect(toolResultContent).toHaveLength(2);
-      expect(toolResultContent[0]).toEqual({
+      expect((result[2] as ToolModelMessage).content[0]).toEqual({
         type: "tool-result",
         toolCallId: "call1",
         toolName: "bash",
         output: { type: "json", value: { stdout: "/home/user" } },
       });
-      expect(toolResultContent[1]).toEqual({
-        type: "tool-result",
-        toolCallId: "call2",
-        toolName: "bash",
-        output: { type: "json", value: { stdout: "file1 file2" } },
-      });
+      expect(result[3].role).toBe("assistant");
+      expect((result[3] as AssistantModelMessage).content).toEqual([
+        { type: "text", text: "After" },
+      ]);
     });
 
-    it("should handle mixed content with tool results properly", () => {
+    it("should interleave multiple tool-call groups with their results", () => {
       const assistantMsg: AssistantModelMessage = {
         role: "assistant",
         content: [
-          { type: "text", text: "First, let me check something." },
+          { type: "text", text: "Step 1" },
           { type: "tool-call", toolCallId: "call1", toolName: "bash", input: { script: "pwd" } },
+          { type: "text", text: "Step 2" },
           { type: "tool-call", toolCallId: "call2", toolName: "bash", input: { script: "ls" } },
-          { type: "text", text: "Now let me check another thing." },
-          {
-            type: "tool-call",
-            toolCallId: "call3",
-            toolName: "file_read",
-            input: { path: "test.txt" },
-          },
         ],
       };
       const toolMsg: ToolModelMessage = {
@@ -182,45 +91,60 @@ describe("modelMessageTransform", () => {
             type: "tool-result",
             toolCallId: "call1",
             toolName: "bash",
-            output: { type: "json", value: { stdout: "/home/user" } },
+            output: { type: "json", value: { stdout: "/workspace" } },
           },
           {
             type: "tool-result",
             toolCallId: "call2",
             toolName: "bash",
-            output: { type: "json", value: { stdout: "file1 file2" } },
-          },
-          {
-            type: "tool-result",
-            toolCallId: "call3",
-            toolName: "file_read",
-            output: { type: "json", value: { content: "test content" } },
+            output: { type: "json", value: { stdout: "file.txt" } },
           },
         ],
       };
-      const messages: ModelMessage[] = [assistantMsg, toolMsg];
 
-      const result = transformModelMessages(messages, "anthropic");
+      const result = transformModelMessages([assistantMsg, toolMsg], "anthropic");
 
-      // Should split into multiple messages with tool results properly placed
-      expect(result.length).toBeGreaterThan(2);
-
-      // First should be text
+      expect(result).toHaveLength(6);
       expect(result[0].role).toBe("assistant");
       expect((result[0] as AssistantModelMessage).content).toEqual([
-        { type: "text", text: "First, let me check something." },
+        { type: "text", text: "Step 1" },
       ]);
-
-      // Then tool calls with their results
       expect(result[1].role).toBe("assistant");
-      const secondContent = (result[1] as AssistantModelMessage).content;
-      expect(Array.isArray(secondContent)).toBe(true);
-      if (Array.isArray(secondContent)) {
-        expect(secondContent.some((c) => c.type === "tool-call")).toBe(true);
-      }
-
-      // Tool results should follow tool calls
+      expect((result[1] as AssistantModelMessage).content[0]).toEqual({
+        type: "tool-call",
+        toolCallId: "call1",
+        toolName: "bash",
+        input: { script: "pwd" },
+      });
       expect(result[2].role).toBe("tool");
+      expect((result[2] as ToolModelMessage).content[0]).toMatchObject({ toolCallId: "call1" });
+      expect(result[3].role).toBe("assistant");
+      expect((result[3] as AssistantModelMessage).content).toEqual([
+        { type: "text", text: "Step 2" },
+      ]);
+      expect(result[4].role).toBe("assistant");
+      expect((result[4] as AssistantModelMessage).content[0]).toEqual({
+        type: "tool-call",
+        toolCallId: "call2",
+        toolName: "bash",
+        input: { script: "ls" },
+      });
+      expect(result[5].role).toBe("tool");
+      expect((result[5] as ToolModelMessage).content[0]).toMatchObject({ toolCallId: "call2" });
+    });
+    it("should keep text-only messages unchanged", () => {
+      const assistantMsg1: AssistantModelMessage = {
+        role: "assistant",
+        content: [{ type: "text", text: "Let me help you with that." }],
+      };
+      const assistantMsg2: AssistantModelMessage = {
+        role: "assistant",
+        content: [{ type: "text", text: "Here's the result." }],
+      };
+      const messages: ModelMessage[] = [assistantMsg1, assistantMsg2];
+
+      const result = transformModelMessages(messages, "anthropic");
+      expect(result).toEqual(messages);
     });
   });
 
@@ -659,10 +583,10 @@ describe("modelMessageTransform", () => {
 
       const result = transformModelMessages(messages, "openai");
 
-      // Should have user, text, tool-call, tool-result (no reasoning)
+      // Should still contain user, assistant, and tool messages after filtering
       expect(result.length).toBeGreaterThan(2);
 
-      // Find the assistant message with text
+      // Find the assistant message with text (reasoning should remain alongside text)
       const textMessage = result.find((msg) => {
         if (msg.role !== "assistant") return false;
         const content = msg.content;
@@ -672,9 +596,7 @@ describe("modelMessageTransform", () => {
       if (textMessage) {
         const content = (textMessage as AssistantModelMessage).content;
         if (Array.isArray(content)) {
-          // Should not have reasoning parts
-          expect(content.some((c) => c.type === "reasoning")).toBe(false);
-          // Should have text
+          expect(content.some((c) => c.type === "reasoning")).toBe(true);
           expect(content.some((c) => c.type === "text")).toBe(true);
         }
       }
diff --git a/src/utils/messages/modelMessageTransform.ts b/src/utils/messages/modelMessageTransform.ts
