Merge PR #1859: reviewed focused maintenance fix

agent/cli/ui/rail.py

@@ -44,6 +44,7 @@ class RailStep:
     title: str
     tool: str = ""
     args: Any = None
+    call_id: str = ""
     status: StepStatus = "active"
     lines: list[str] = field(default_factory=list)
     started_at: float = field(default_factory=time.monotonic)
@@ -346,8 +347,9 @@ class RailRunDashboard:
             self.thinking_active = False
             tool = str(data.get("tool") or "tool")
             args = data.get("arguments") or {}
+            call_id = str(data.get("call_id") or "")
             self.iterations += 1
-            step = RailStep(title=_tool_title(tool, args), tool=tool, args=args)
+            step = RailStep(title=_tool_title(tool, args), tool=tool, args=args, call_id=call_id)
             detail = _initial_detail(tool, args)
             if detail:
                 step.lines.append(detail)
@@ -358,7 +360,7 @@ class RailRunDashboard:
 
         if event_type == "tool_progress":
             tool = str(data.get("tool") or "")
-            step = self._active_step(tool)
+            step = self._active_step(tool, call_id=str(data.get("call_id") or ""))
             pieces: list[str] = []
             stage = _collapse_ws(data.get("stage"))
             current = data.get("current")
@@ -381,14 +383,14 @@ class RailRunDashboard:
 
         if event_type == "tool_heartbeat":
             tool = str(data.get("tool") or "")
-            self._active_step(tool)
+            self._active_step(tool, call_id=str(data.get("call_id") or ""))
             self.refresh()
             return
 
         if event_type == "tool_result":
             tool = str(data.get("tool") or "")
             status = str(data.get("status") or "ok")
-            step = self._active_step(tool)
+            step = self._active_step(tool, call_id=str(data.get("call_id") or ""))
             step.status = "done" if status == "ok" else "error"
             step.duration_s = float(data.get("elapsed_ms") or 0.0) / 1000
             for line in _result_summary(tool, status, data.get("preview")):
@@ -458,12 +460,23 @@ class RailRunDashboard:
         lines = list(step.lines)
         return lines[-7:]
 
-    def _active_step(self, tool: str) -> RailStep:
+    def _active_step(self, tool: str, call_id: str = "") -> RailStep:
         raw_tool = str(tool or "")
+        raw_call_id = str(call_id or "")
+        # Concurrent calls to the same tool (the loop runs consecutive
+        # readonly tool calls in parallel) can have more than one step
+        # "active" at once with the same tool name. Matching on the call's
+        # own id -- always present on tool_call/tool_progress/tool_heartbeat/
+        # tool_result payloads -- keeps each event paired with the step it
+        # actually belongs to, instead of guessing by tool name and recency.
+        if raw_call_id:
+            for step in self.steps:
+                if step.call_id == raw_call_id:
+                    return step
         for step in reversed(self.steps):
             if step.status == "active" and (not raw_tool or step.tool == raw_tool):
                 return step
-        step = RailStep(title=_tool_title(raw_tool or "tool", None), tool=raw_tool)
+        step = RailStep(title=_tool_title(raw_tool or "tool", None), tool=raw_tool, call_id=raw_call_id)
         self.steps.append(step)
         self.steps = self.steps[-10:]
         return step

agent/tests/test_cli_rail_dashboard.py

@@ -0,0 +1,70 @@
+"""Tests for the CLI rail dashboard's tool_call/tool_result pairing."""
+
+from __future__ import annotations
+
+from typing import Iterator
+
+import pytest
+
+from cli.ui.rail import RailRunDashboard
+
+
+@pytest.fixture()
+def dashboard() -> Iterator[RailRunDashboard]:
+    dash = RailRunDashboard("test prompt", max_iter=10)
+    try:
+        yield dash
+    finally:
+        dash.close()
+
+
+class TestConcurrentToolCallPairing:
+    """Two concurrent calls to the same tool must not have their results swapped.
+
+    The agent loop runs consecutive readonly tool calls in parallel, so more
+    than one step for the same tool name can be "active" at once. Matching a
+    tool_result to a step by tool name and recency alone pairs the first
+    arriving result with the most recently started step instead of the one
+    it actually belongs to.
+    """
+
+    def test_results_pair_with_their_own_call_via_call_id(self, dashboard: RailRunDashboard) -> None:
+        dashboard.handle_event(
+            "tool_call", {"tool": "get_market_data", "arguments": {"symbol": "AAPL"}, "call_id": "call-1"}
+        )
+        dashboard.handle_event(
+            "tool_call", {"tool": "get_market_data", "arguments": {"symbol": "MSFT"}, "call_id": "call-2"}
+        )
+
+        dashboard.handle_event(
+            "tool_result",
+            {"tool": "get_market_data", "status": "ok", "preview": '{"summary": "AAPL result"}', "call_id": "call-1"},
+        )
+        dashboard.handle_event(
+            "tool_result",
+            {"tool": "get_market_data", "status": "ok", "preview": '{"summary": "MSFT result"}', "call_id": "call-2"},
+        )
+
+        aapl_step = next(s for s in dashboard.steps if s.args == {"symbol": "AAPL"})
+        msft_step = next(s for s in dashboard.steps if s.args == {"symbol": "MSFT"})
+        assert any("AAPL result" in line for line in aapl_step.lines)
+        assert not any("MSFT result" in line for line in aapl_step.lines)
+        assert any("MSFT result" in line for line in msft_step.lines)
+        assert not any("AAPL result" in line for line in msft_step.lines)
+
+    def test_progress_event_updates_its_own_call_via_call_id(self, dashboard: RailRunDashboard) -> None:
+        dashboard.handle_event(
+            "tool_call", {"tool": "backtest", "arguments": {"symbol": "AAPL"}, "call_id": "call-1"}
+        )
+        dashboard.handle_event(
+            "tool_call", {"tool": "backtest", "arguments": {"symbol": "MSFT"}, "call_id": "call-2"}
+        )
+
+        dashboard.handle_event(
+            "tool_progress", {"tool": "backtest", "message": "halfway for MSFT", "call_id": "call-2"}
+        )
+
+        aapl_step = next(s for s in dashboard.steps if s.args == {"symbol": "AAPL"})
+        msft_step = next(s for s in dashboard.steps if s.args == {"symbol": "MSFT"})
+        assert not any("halfway for MSFT" in line for line in aapl_step.lines)
+        assert any("halfway for MSFT" in line for line in msft_step.lines)