🤖 refactor: Shared tool result processing for live & historical

ammar-agent · ammar-agent · commit 01b72d5b6258 · 2025-10-28T18:16:52.000Z
## Problem
Status (and TODOs) weren't persisting across page reloads/workspace switches
because loadHistoricalMessages() didn't process tool results to reconstruct
derived state.

## Solution
Extracted processToolResult() as single source of truth for updating derived
state from tool results. Called from:
1. handleToolCallEnd() - live streaming events
2. loadHistoricalMessages() - historical message loading

This ensures agentStatus and currentTodos are reconstructed uniformly whether
processing live events or historical snapshots.

## Implementation
- Extracted 23-line processToolResult() method
- Updated handleToolCallEnd() to call shared method (-12 LoC)
- Updated loadHistoricalMessages() to process tool parts (+6 LoC)
- Added 3 tests for historical message reconstruction

## Benefits
- Single source of truth for tool result processing
- No special reconstruction logic needed
- Easier to extend with new derived state
- Leverages existing tool part architecture

## Testing
- Added 3 new tests for historical message scenarios
- All 842 tests pass
- Typecheck and lint pass

Net: ~17 LoC
diff --git a/src/utils/messages/StreamingMessageAggregator.status.test.ts b/src/utils/messages/StreamingMessageAggregator.status.test.ts
@@ -376,4 +376,122 @@ describe("StreamingMessageAggregator - Agent Status", () => {
     expect(status?.emoji).toBe("🔍");
     expect(status?.message).toBe("Analyzing code");
   });
+
+  it("should reconstruct agentStatus when loading historical messages", () => {
+    const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z");
+
+    // Create historical messages with a completed status_set tool call
+    const historicalMessages = [
+      {
+        id: "msg1",
+        role: "user" as const,
+        parts: [{ type: "text" as const, text: "Hello" }],
+        metadata: { timestamp: Date.now(), historySequence: 1 },
+      },
+      {
+        id: "msg2",
+        role: "assistant" as const,
+        parts: [
+          { type: "text" as const, text: "Working on it..." },
+          {
+            type: "dynamic-tool" as const,
+            toolCallId: "tool1",
+            toolName: "status_set",
+            state: "output-available" as const,
+            input: { emoji: "🔍", message: "Analyzing code" },
+            output: { success: true, emoji: "🔍", message: "Analyzing code" },
+            timestamp: Date.now(),
+          },
+        ],
+        metadata: { timestamp: Date.now(), historySequence: 2 },
+      },
+    ];
+
+    // Load historical messages
+    aggregator.loadHistoricalMessages(historicalMessages);
+
+    // Status should be reconstructed from the historical tool call
+    const status = aggregator.getAgentStatus();
+    expect(status).toBeDefined();
+    expect(status?.emoji).toBe("🔍");
+    expect(status?.message).toBe("Analyzing code");
+  });
+
+  it("should use most recent status_set when loading multiple historical messages", () => {
+    const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z");
+
+    // Create historical messages with multiple status_set calls
+    const historicalMessages = [
+      {
+        id: "msg1",
+        role: "assistant" as const,
+        parts: [
+          {
+            type: "dynamic-tool" as const,
+            toolCallId: "tool1",
+            toolName: "status_set",
+            state: "output-available" as const,
+            input: { emoji: "🔍", message: "First status" },
+            output: { success: true, emoji: "🔍", message: "First status" },
+            timestamp: Date.now(),
+          },
+        ],
+        metadata: { timestamp: Date.now(), historySequence: 1 },
+      },
+      {
+        id: "msg2",
+        role: "assistant" as const,
+        parts: [
+          {
+            type: "dynamic-tool" as const,
+            toolCallId: "tool2",
+            toolName: "status_set",
+            state: "output-available" as const,
+            input: { emoji: "📝", message: "Second status" },
+            output: { success: true, emoji: "📝", message: "Second status" },
+            timestamp: Date.now(),
+          },
+        ],
+        metadata: { timestamp: Date.now(), historySequence: 2 },
+      },
+    ];
+
+    // Load historical messages
+    aggregator.loadHistoricalMessages(historicalMessages);
+
+    // Should use the most recent (last processed) status
+    const status = aggregator.getAgentStatus();
+    expect(status?.emoji).toBe("📝");
+    expect(status?.message).toBe("Second status");
+  });
+
+  it("should not reconstruct status from failed status_set in historical messages", () => {
+    const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z");
+
+    // Create historical message with failed status_set
+    const historicalMessages = [
+      {
+        id: "msg1",
+        role: "assistant" as const,
+        parts: [
+          {
+            type: "dynamic-tool" as const,
+            toolCallId: "tool1",
+            toolName: "status_set",
+            state: "output-available" as const,
+            input: { emoji: "not-emoji", message: "test" },
+            output: { success: false, error: "emoji must be a single emoji character" },
+            timestamp: Date.now(),
+          },
+        ],
+        metadata: { timestamp: Date.now(), historySequence: 1 },
+      },
+    ];
+
+    // Load historical messages
+    aggregator.loadHistoricalMessages(historicalMessages);
+
+    // Status should remain undefined (failed validation)
+    expect(aggregator.getAgentStatus()).toBeUndefined();
+  });
 });
diff --git a/src/utils/messages/StreamingMessageAggregator.ts b/src/utils/messages/StreamingMessageAggregator.ts
@@ -199,7 +199,18 @@ export class StreamingMessageAggregator {
   loadHistoricalMessages(messages: CmuxMessage[]): void {
     for (const message of messages) {
       this.messages.set(message.id, message);
+
+      // Process completed tool calls to reconstruct derived state (todos, agentStatus)
+      // This ensures state persists across page reloads and workspace switches
+      if (message.role === "assistant") {
+        for (const part of message.parts) {
+          if (isDynamicToolPart(part) && part.state === "output-available") {
+            this.processToolResult(part.toolName, part.input, part.output);
+          }
+        }
+      }
     }
+
     this.invalidateCache();
   }
 
@@ -491,6 +502,31 @@ export class StreamingMessageAggregator {
     // Tool deltas are for display - args are in dynamic-tool part
   }
 
+
+  /**
+   * Process a completed tool call's result to update derived state.
+   * Called for both live tool-call-end events and historical tool parts.
+   *
+   * This is the single source of truth for updating state from tool results,
+   * ensuring consistency whether processing live events or historical messages.
+   */
+  private processToolResult(toolName: string, input: unknown, output: unknown): void {
+    // Update TODO state if this was a successful todo_write
+    if (toolName === "todo_write" && hasSuccessResult(output)) {
+      const args = input as { todos: TodoItem[] };
+      // Only update if todos actually changed (prevents flickering from reference changes)
+      if (!this.todosEqual(this.currentTodos, args.todos)) {
+        this.currentTodos = args.todos;
+      }
+    }
+
+    // Update agent status if this was a successful status_set
+    if (toolName === "status_set" && hasSuccessResult(output)) {
+      const args = input as { emoji: string; message: string };
+      this.agentStatus = { emoji: args.emoji, message: args.message };
+    }
+  }
+
   handleToolCallEnd(data: ToolCallEndEvent): void {
     const message = this.messages.get(data.messageId);
     if (message) {
@@ -505,20 +541,8 @@ export class StreamingMessageAggregator {
         (toolPart as DynamicToolPartAvailable).state = "output-available";
         (toolPart as DynamicToolPartAvailable).output = data.result;
 
-        // Update TODO state if this was a successful todo_write
-        if (data.toolName === "todo_write" && hasSuccessResult(data.result)) {
-          const args = toolPart.input as { todos: TodoItem[] };
-          // Only update if todos actually changed (prevents flickering from reference changes)
-          if (!this.todosEqual(this.currentTodos, args.todos)) {
-            this.currentTodos = args.todos;
-          }
-        }
-
-        // Update agent status if this was a successful status_set
-        if (data.toolName === "status_set" && hasSuccessResult(data.result)) {
-          const args = toolPart.input as { emoji: string; message: string };
-          this.agentStatus = { emoji: args.emoji, message: args.message };
-        }
+        // Process tool result to update derived state (todos, agentStatus, etc.)
+        this.processToolResult(data.toolName, toolPart.input, data.result);
       }
       this.invalidateCache();
     }
