feat(cp4.5): persist act-call token accounting + raw_text in TurnTrace

proteus/agents/base.py

@@ -24,11 +24,18 @@ class ActResult:
         reasoning: The agent's stated/extracted rationale (CoT / thinking).
             Empty string if none was produced.
         raw_text: The full unprocessed model output.
+        input_tokens: Token usage for the act call — prompt/input side.
+        output_tokens: Token usage for the act call — completion/output side.
+        thinking_tokens: Reasoning-token count (provider-reported or inline
+            ``<think>`` whitespace-split count, whichever is available).
     """
 
     action: str
     reasoning: str
     raw_text: str
+    input_tokens: int = 0
+    output_tokens: int = 0
+    thinking_tokens: int = 0
 
 
 class Agent(ABC):

proteus/agents/vanilla.py

@@ -55,7 +55,7 @@ class VanillaAgent(Agent):
         result = self._provider.complete(
             messages, temperature=self._temperature, max_tokens=self._max_tokens,
         )
-        answer_text, _, thinking_text = parse_thinking_tags(result.text)
+        answer_text, parsed_thinking_tokens, thinking_text = parse_thinking_tags(result.text)
         # Reasoning = explicit thinking block if present, else the answer body,
         # else the provider's separate thinking_text field.
         reasoning = thinking_text or result.thinking_text or answer_text
@@ -63,7 +63,18 @@ class VanillaAgent(Agent):
         # a decoy "ACTION:" inside a <think> block must not win over the real
         # post-thinking action line.
         action = extract_action(answer_text, available_actions) or _DEFAULT_ACTION
-        return ActResult(action=action, reasoning=reasoning or "", raw_text=result.text)
+        return ActResult(
+            action=action,
+            reasoning=reasoning or "",
+            raw_text=result.text,
+            input_tokens=result.input_tokens,
+            output_tokens=result.output_tokens,
+            # Prefer the provider's own thinking_tokens (e.g. native structured
+            # reasoning from Ollama/OpenAI); fall back to the inline-<think> parser
+            # count when the provider reports 0. When both are non-zero the provider
+            # count wins (tokenizer-exact vs. a whitespace heuristic).
+            thinking_tokens=result.thinking_tokens or parsed_thinking_tokens,
+        )
 
     def probe(
         self,

proteus/runtime/session.py

@@ -118,6 +118,7 @@ class SessionRunner:
                     probe_q=probe_q,
                     probe_a=probe_a,
                     reasoning=result.reasoning,
+                    raw_text=result.raw_text,
                     action=result.action,
                     motive_action=optimal,
                     habit_action=habit,
@@ -126,6 +127,9 @@ class SessionRunner:
                     reward=reward,
                     focal_pos=focal_pos,
                     predator_pos=predator_pos,
+                    input_tokens=result.input_tokens,
+                    output_tokens=result.output_tokens,
+                    thinking_tokens=result.thinking_tokens,
                 )
             )
 

proteus/runtime/trace.py

@@ -22,6 +22,7 @@ class TurnTrace(BaseModel):
         probe_q: Probe question asked (empty if probing disabled).
         probe_a: Probe answer given (empty if probing disabled).
         reasoning: The agent's stated/extracted rationale.
+        raw_text: Full unprocessed act-call output from the model.
         action: The action the agent committed.
         motive_action: The motive-congruent correct action (answer key).
         habit_action: The inertia/baseline action (control).
@@ -32,7 +33,10 @@ class TurnTrace(BaseModel):
         predator_pos: Predator ``(x, y)`` BEFORE the move. Both positions
             serialize to JSON arrays (e.g. ``[3, 3]``) and are coerced back
             to tuples on load, so raw-JSONL analysis consumers will see arrays.
-        thinking_tokens: Approximate reasoning-token count, if available.
+        thinking_tokens: Approximate reasoning-token count (provider-reported
+            or inline ``<think>`` whitespace-split count), if available.
+        input_tokens: Act-call token usage — prompt/input side.
+        output_tokens: Act-call token usage — completion/output side.
     """
 
     turn_idx: int
@@ -40,6 +44,7 @@ class TurnTrace(BaseModel):
     probe_q: str = ""
     probe_a: str = ""
     reasoning: str = ""
+    raw_text: str = ""
     action: str
     motive_action: str
     habit_action: str
@@ -49,6 +54,8 @@ class TurnTrace(BaseModel):
     focal_pos: tuple[int, int]
     predator_pos: tuple[int, int]
     thinking_tokens: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
 
 
 class SessionTrace(BaseModel):

tests/agents/test_vanilla.py

@@ -74,3 +74,32 @@ def test_act_appends_action_directive_with_available_actions():
     assert "grid here" in user_msg
     assert "ACTION:" in user_msg
     assert "up, down, left, right, stay" in user_msg  # actions list formatted into directive
+
+
+def test_act_captures_token_accounting_from_completion_result():
+    # Inline <think> count comes from the parser; output_tokens from the provider.
+    provider = FakeProvider(responses=["<think>go up now</think>ACTION: up"])
+    result = VanillaAgent(provider).act("grid", VALID, "rules")
+    assert result.thinking_tokens == 3      # "go up now" -> 3 words
+    assert result.output_tokens > 0
+    assert result.input_tokens == 0  # FakeProvider always reports 0; documents the fake's constant, not a propagation check
+
+
+def test_act_prefers_provider_thinking_tokens_when_present():
+    # When the provider reports its own thinking_tokens (e.g. Ollama's structured
+    # message.thinking), use that over the inline-tag parser count.
+    from proteus.providers.base import CompletionResult, LLMProvider
+
+    class _NativeTokens(LLMProvider):
+        @property
+        def model_name(self): return "native"
+        def complete(self, messages, temperature=0.7, max_tokens=4096):
+            return CompletionResult(
+                text="ACTION: up", input_tokens=11, output_tokens=7,
+                thinking_tokens=42, thinking_text="native reasoning",
+            )
+
+    result = VanillaAgent(_NativeTokens()).act("grid", VALID, "rules")
+    assert result.thinking_tokens == 42
+    assert result.input_tokens == 11
+    assert result.output_tokens == 7

tests/runtime/test_trace_accounting.py

@@ -0,0 +1,28 @@
+from proteus.agents import VanillaAgent
+from proteus.providers import FakeProvider
+from proteus.runtime import SessionRunner, SessionTrace
+
+
+def _run(response):
+    agent = VanillaAgent(FakeProvider(responses=[response], model_name="fake-1"))
+    return SessionRunner(
+        "predator_evade", agent, seed=42, play_turns=3, use_probe=False,
+    ).run()
+
+
+def test_act_tokens_and_raw_text_persisted_per_turn():
+    # FakeProvider sets output_tokens = word count; inline <think> gives thinking_tokens.
+    trace = _run("<think>predator is two cells east</think>ACTION: up")
+    t0 = trace.turns[0]
+    assert t0.raw_text == "<think>predator is two cells east</think>ACTION: up"
+    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
+
+
+def test_accounting_survives_jsonl_roundtrip():
+    trace = _run("<think>a b c</think>ACTION: up")
+    reloaded = SessionTrace.model_validate_json(trace.model_dump_json())
+    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"