feat(cp4.5): preserve probe reasoning + raw_text + token accounting via ProbeResult

proteus/agents/__init__.py

@@ -1,6 +1,6 @@
 """proteus.agents — slim LLM agent abstraction (no forfeit/stake/risk)."""
 
-from proteus.agents.base import Agent, ActResult
+from proteus.agents.base import Agent, ActResult, ProbeResult
 from proteus.agents.vanilla import VanillaAgent
 
-__all__ = ["Agent", "ActResult", "VanillaAgent"]
+__all__ = ["Agent", "ActResult", "ProbeResult", "VanillaAgent"]

proteus/agents/base.py

@@ -38,6 +38,30 @@ class ActResult:
     thinking_tokens: int = 0
 
 
+@dataclass(frozen=True)
+class ProbeResult:
+    """An agent's answer to a side-channel probe question for one turn.
+
+    Immutable: written verbatim into the session trace as a 1st-class
+    measurement target (the probe is scored offline later).
+
+    Attributes:
+        answer: The think-stripped probe answer.
+        reasoning: The probe's stated/extracted rationale (CoT / thinking).
+        raw_text: The full unprocessed model output for the probe.
+        input_tokens: Prompt token usage for the probe call.
+        output_tokens: Completion token usage for the probe call.
+        thinking_tokens: Reasoning-token count for the probe call.
+    """
+
+    answer: str
+    reasoning: str = ""
+    raw_text: str = ""
+    input_tokens: int = 0
+    output_tokens: int = 0
+    thinking_tokens: int = 0
+
+
 class Agent(ABC):
     """Abstract LLM agent that plays a motive_grid scenario."""
 
@@ -71,7 +95,7 @@ class Agent(ABC):
         observation: str,
         question: str,
         system_prompt: str,
-    ) -> str:
+    ) -> ProbeResult:
         """Answer a side-channel comprehension probe (does not change state).
 
         Args:
@@ -80,7 +104,8 @@ class Agent(ABC):
             system_prompt: Rules + handover framing for the session.
 
         Returns:
-            Free-text probe answer.
+            A :class:`ProbeResult` containing the think-stripped answer,
+            reasoning, raw model output, and token accounting.
         """
 
     @abstractmethod

proteus/agents/vanilla.py

@@ -8,7 +8,7 @@ as the agent's reasoning even when the action text is terse.
 
 from __future__ import annotations
 
-from proteus.agents.base import Agent, ActResult
+from proteus.agents.base import Agent, ActResult, ProbeResult
 from proteus.agents.parsing import extract_action
 from proteus.providers.base import LLMProvider
 from proteus.providers.thinking_utils import parse_thinking_tags
@@ -81,7 +81,12 @@ class VanillaAgent(Agent):
         observation: str,
         question: str,
         system_prompt: str,
-    ) -> str:
+    ) -> ProbeResult:
+        """Answer a side-channel comprehension probe and return a ProbeResult.
+
+        Mirrors act()'s token accounting logic: provider-reported thinking_tokens
+        are preferred; the inline-<think> parser count is the fallback.
+        """
         messages = [
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": f"{observation}\n\n{question}"},
@@ -89,8 +94,18 @@ class VanillaAgent(Agent):
         result = self._provider.complete(
             messages, temperature=self._temperature, max_tokens=self._max_tokens,
         )
-        answer_text, _, _ = parse_thinking_tags(result.text)
-        return answer_text
+        answer_text, parsed_thinking_tokens, thinking_text = parse_thinking_tags(result.text)
+        reasoning = thinking_text or result.thinking_text or ""
+        return ProbeResult(
+            answer=answer_text,
+            reasoning=reasoning,
+            raw_text=result.text,
+            input_tokens=result.input_tokens,
+            output_tokens=result.output_tokens,
+            # Provider count preferred; inline-<think> parser count is the fallback
+            # (mirrors act()).
+            thinking_tokens=result.thinking_tokens or parsed_thinking_tokens,
+        )
 
     def reset(self) -> None:
         # Vanilla carries no cross-turn state.

proteus/runtime/session.py

@@ -93,10 +93,17 @@ class SessionRunner:
         for turn_idx in range(1, self._play_turns + 1):
             observation = self._observation(turn_idx)
 
-            probe_q = probe_a = ""
+            probe_q = probe_a = probe_reasoning = probe_raw_text = ""
+            probe_input_tokens = probe_output_tokens = probe_thinking_tokens = 0
             if self._use_probe:
                 probe_q = _PROBE_QUESTION
-                probe_a = self._agent.probe(observation, probe_q, system_prompt)
+                probe = self._agent.probe(observation, probe_q, system_prompt)
+                probe_a = probe.answer
+                probe_reasoning = probe.reasoning
+                probe_raw_text = probe.raw_text
+                probe_input_tokens = probe.input_tokens
+                probe_output_tokens = probe.output_tokens
+                probe_thinking_tokens = probe.thinking_tokens
 
             # Pre-move answer keys + positions.
             optimal = self._scenario.optimal_action(self._game)
@@ -117,6 +124,11 @@ class SessionRunner:
                     observation=observation,
                     probe_q=probe_q,
                     probe_a=probe_a,
+                    probe_reasoning=probe_reasoning,
+                    probe_raw_text=probe_raw_text,
+                    probe_input_tokens=probe_input_tokens,
+                    probe_output_tokens=probe_output_tokens,
+                    probe_thinking_tokens=probe_thinking_tokens,
                     reasoning=result.reasoning,
                     raw_text=result.raw_text,
                     action=result.action,

proteus/runtime/trace.py

@@ -21,6 +21,12 @@ class TurnTrace(BaseModel):
         observation: The text observation shown to the agent this turn.
         probe_q: Probe question asked (empty if probing disabled).
         probe_a: Probe answer given (empty if probing disabled).
+        probe_reasoning: The probe's stated/extracted rationale (CoT / thinking).
+        probe_raw_text: Full unprocessed probe-call output from the model.
+        probe_input_tokens: Probe-call token usage — prompt/input side.
+        probe_output_tokens: Probe-call token usage — completion/output side.
+        probe_thinking_tokens: Reasoning-token count for the probe call
+            (provider-reported or inline ``<think>`` whitespace-split count).
         reasoning: The agent's stated/extracted rationale.
         raw_text: Full unprocessed act-call output from the model.
         action: The action the agent committed.
@@ -43,6 +49,11 @@ class TurnTrace(BaseModel):
     observation: str
     probe_q: str = ""
     probe_a: str = ""
+    probe_reasoning: str = ""
+    probe_raw_text: str = ""
+    probe_input_tokens: int = 0
+    probe_output_tokens: int = 0
+    probe_thinking_tokens: int = 0
     reasoning: str = ""
     raw_text: str = ""
     action: str

tests/agents/test_vanilla.py

@@ -26,8 +26,8 @@ def test_act_falls_back_to_stay_when_unparseable():
 def test_probe_returns_text_and_sends_question():
     provider = FakeProvider(responses=["the predator is east"])
     agent = VanillaAgent(provider)
-    answer = agent.probe("grid", "where is the predator?", "rules")
-    assert answer == "the predator is east"
+    result = agent.probe("grid", "where is the predator?", "rules")
+    assert result.answer == "the predator is east"
     # the probe question reached the provider
     assert any("where is the predator?" in m["content"] for m in provider.calls[-1])
 
@@ -103,3 +103,15 @@ def test_act_prefers_provider_thinking_tokens_when_present():
     assert result.thinking_tokens == 42
     assert result.input_tokens == 11
     assert result.output_tokens == 7
+
+
+def test_probe_returns_probe_result_with_reasoning_and_tokens():
+    from proteus.agents import ProbeResult
+    provider = FakeProvider(responses=["<think>predator is two cells east</think>go up"])
+    result = VanillaAgent(provider).probe("grid", "where is the predator?", "rules")
+    assert isinstance(result, ProbeResult)
+    assert result.answer == "go up"                 # think-stripped answer
+    assert "predator is two cells east" in result.reasoning
+    assert result.raw_text == "<think>predator is two cells east</think>go up"
+    assert result.thinking_tokens == 5              # 5-word think block (parser fallback)
+    assert result.output_tokens > 0

tests/runtime/test_trace_accounting.py

@@ -18,6 +18,9 @@ def test_act_tokens_and_raw_text_persisted_per_turn():
     assert t0.output_tokens > 0                     # propagated from CompletionResult
     assert t0.input_tokens == 0                     # FakeProvider reports 0; output_tokens below is the live propagation check
     assert t0.thinking_tokens == 5                  # whitespace-split of the think block
+    # use_probe=False: probe fields stay at their defaults.
+    assert t0.probe_raw_text == ""
+    assert t0.probe_thinking_tokens == 0
 
 
 def test_accounting_survives_jsonl_roundtrip():
@@ -26,3 +29,29 @@ def test_accounting_survives_jsonl_roundtrip():
     assert reloaded.model_dump() == trace.model_dump()
     assert reloaded.turns[0].thinking_tokens == 3
     assert reloaded.turns[0].raw_text == "<think>a b c</think>ACTION: up"
+
+
+def test_probe_accounting_persisted_when_enabled():
+    # Distinct probe vs act responses (different <think> word counts) so a
+    # probe/act field cross-wiring bug would be visible: turn 1's probe consumes
+    # responses[0] (3-word think) and turn 1's act consumes responses[1] (5-word think).
+    agent = VanillaAgent(FakeProvider(
+        responses=[
+            "<think>it is east</think>predator is east",          # probe: 3 think-words
+            "<think>move up and away now</think>I should go up\nACTION: up",  # act: 5 think-words
+        ],
+        model_name="fake-1",
+    ))
+    trace = SessionRunner(
+        "predator_evade", agent, seed=42, play_turns=3, use_probe=True,
+    ).run()
+    t0 = trace.turns[0]
+    assert t0.probe_a                               # answer recorded
+    assert "it is east" in t0.probe_reasoning       # probe reasoning preserved
+    assert t0.probe_raw_text                        # raw probe output preserved
+    assert t0.probe_thinking_tokens == 3            # from the PROBE response
+    assert t0.thinking_tokens == 5                  # from the ACT response — NOT cross-wired
+    assert t0.probe_output_tokens > 0
+    # survives JSONL round-trip
+    reloaded = SessionTrace.model_validate_json(trace.model_dump_json())
+    assert reloaded.model_dump() == trace.model_dump()