From 282e98e3c718f1a276f0e43ea34bcbc10c3688f1 Mon Sep 17 00:00:00 2001
From: Shuxin Lin <linshuhsin@gmail.com>
Date: Fri, 31 Jul 2026 10:26:02 -0400
Subject: [PATCH] Persist Stirrup assistant reasoning in trajectories

Signed-off-by: Shuxin Lin <linshuhsin@gmail.com>
---
 src/agent/stirrup_agent/tests/test_runner.py | 16 +++++++++++++
 src/agent/stirrup_agent/trajectory.py        | 24 ++++++++++++++++++--
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/src/agent/stirrup_agent/tests/test_runner.py b/src/agent/stirrup_agent/tests/test_runner.py
index c64430c2..c41e2c5b 100644
--- a/src/agent/stirrup_agent/tests/test_runner.py
+++ b/src/agent/stirrup_agent/tests/test_runner.py
@@ -51,9 +51,16 @@ class _TC:
     tool_call_id: str
 
 
+@dataclass
+class _Reasoning:
+    content: str
+    signature: str | None = None
+
+
 @dataclass
 class _Assistant:
     content: str
+    reasoning: _Reasoning | None = None
     tool_calls: list = field(default_factory=list)
     token_usage: _Usage = field(default_factory=_Usage)
     request_start_time: float | None = None
@@ -192,6 +199,10 @@ def test_build_trajectory_maps_turns_calls_and_outputs():
         [
             _Assistant(
                 content="let me check work orders",
+                reasoning=_Reasoning(
+                    content="I should query the work-order tool.",
+                    signature="sig-1",
+                ),
                 tool_calls=[_TC("wo__get_work_order", '{"asset": "CWC04013"}', "t1")],
                 token_usage=_Usage(input=20, answer=8, reasoning=2),
                 request_start_time=1.0,
@@ -217,6 +228,11 @@ def test_build_trajectory_maps_turns_calls_and_outputs():
     assert len(traj.turns) == 2
     assert traj.total_input_tokens == 25
     assert traj.total_output_tokens == 16  # (8+2) + (6+0)
+    assert traj.turns[0].reasoning == {
+        "signature": "sig-1",
+        "content": "I should query the work-order tool.",
+    }
+    assert traj.turns[1].reasoning is None
 
     call = traj.all_tool_calls[0]
     assert call.name == "wo__get_work_order"
diff --git a/src/agent/stirrup_agent/trajectory.py b/src/agent/stirrup_agent/trajectory.py
index ba8b9bc3..0c3541dd 100644
--- a/src/agent/stirrup_agent/trajectory.py
+++ b/src/agent/stirrup_agent/trajectory.py
@@ -8,7 +8,8 @@
 
 Mapping:
   * each ``AssistantMessage`` -> one :class:`~agent.models.TurnRecord`
-    (its ``content`` text, ``tool_calls``, ``token_usage``, request timing);
+    (its ``content`` text, ``reasoning``, ``tool_calls``, ``token_usage``,
+    request timing);
   * each ``ToolMessage`` -> the ``output`` of the matching :class:`ToolCall`,
     joined by ``tool_call_id``.
 
@@ -20,6 +21,7 @@
 from __future__ import annotations
 
 import json
+from dataclasses import dataclass
 from typing import Any, Iterable
 
 from ..models import ToolCall, Trajectory, TurnRecord
@@ -31,6 +33,13 @@
 _WEB_TOOL_NAMES = {"web_search", "web_fetch"}
 
 
+@dataclass
+class StirrupTurnRecord(TurnRecord):
+    """A shared turn record plus Stirrup's optional reasoning payload."""
+
+    reasoning: dict[str, Any] | None = None
+
+
 def classify_tool(tool_name: str, domain_servers: set[str]) -> str:
     """Label a Stirrup tool call ``"domain"`` / ``"code"`` / ``"other"``.
 
@@ -79,6 +88,16 @@ def _parse_arguments(arguments: Any) -> dict:
     return {}
 
 
+def _reasoning_payload(reasoning: Any) -> dict[str, Any] | None:
+    """Return Stirrup ``AssistantMessage.reasoning`` in JSON-friendly form."""
+    if reasoning is None:
+        return None
+    return {
+        "signature": getattr(reasoning, "signature", None),
+        "content": getattr(reasoning, "content", "") or "",
+    }
+
+
 def _ms(start: float | None, end: float | None) -> float | None:
     if start is None or end is None:
         return None
@@ -123,7 +142,7 @@ def build_trajectory(history: Iterable[Any]) -> Trajectory:
             out_tok = getattr(usage, "output", 0) if usage else 0
 
             trajectory.turns.append(
-                TurnRecord(
+                StirrupTurnRecord(
                     index=turn_index,
                     text=_content_text(getattr(msg, "content", "")),
                     tool_calls=tool_calls,
@@ -133,6 +152,7 @@ def build_trajectory(history: Iterable[Any]) -> Trajectory:
                         getattr(msg, "request_start_time", None),
                         getattr(msg, "request_end_time", None),
                     ),
+                    reasoning=_reasoning_payload(getattr(msg, "reasoning", None)),
                 )
             )
             turn_index += 1