Fix stack 06: per-turn Codex usage + session_id fallback parity

src/agent_runtime_kit/adapters/antigravity.py

@@ -325,6 +325,10 @@ async def _invoke(
                     )
                 )
 
+        # Fall back to the caller's conversation id when the SDK does not echo one,
+        # so a resumed task always returns a usable session_id (matches Claude).
+        session_id = session_id or _conversation_id(task)
+
         process_metadata = (
             self._process_reuse_metadata(process_reused) if self._reuse_process else None
         )

src/agent_runtime_kit/adapters/codex.py

@@ -292,7 +292,7 @@ async def _run_codex(
             task,
             raw_result,
             model=model,
-            session_id=_thread_id(thread),
+            session_id=_thread_id(thread) or _task_session_id(task),
             process_metadata=(
                 self._process_reuse_metadata(process_reused) if self._reuse_process else None
             ),
@@ -571,16 +571,25 @@ def _tool_arguments(value: Any) -> Mapping[str, Any]:
 
 
 def _codex_usage(value: Any) -> Usage:
-    total = field_value(value, "total")
-    if total is None and isinstance(value, Mapping):
-        total = value.get("total", value)
-    input_tokens = optional_int(field_value(total, "input_tokens"))
-    output_tokens = optional_int(field_value(total, "output_tokens"))
-    cached = optional_int(field_value(total, "cached_input_tokens"))
-    total_tokens = optional_int(field_value(total, "total_tokens"))
-    # OpenAI reports cached input inside input_tokens, so never add cached on top.
+    # Prefer the per-turn breakdown ('last') over the thread-cumulative 'total':
+    # on a resumed thread 'total' re-reports every prior turn's tokens, inflating
+    # this turn's usage and cost accounting. Fall back to 'total' when 'last' is
+    # absent (older SDKs / dict-based fakes).
+    breakdown = field_value(value, "last")
+    if breakdown is None:
+        breakdown = field_value(value, "total")
+    if breakdown is None and isinstance(value, Mapping):
+        breakdown = value.get("total", value)
+    input_tokens = optional_int(field_value(breakdown, "input_tokens"))
+    output_tokens = optional_int(field_value(breakdown, "output_tokens"))
+    cached = optional_int(field_value(breakdown, "cached_input_tokens"))
+    total_tokens = optional_int(field_value(breakdown, "total_tokens"))
+    # OpenAI reports cached input INSIDE input_tokens, while the Usage contract
+    # (and the Antigravity adapter) excludes cache reads from input_tokens and
+    # reports them separately. Subtract rather than double-count across the two
+    # fields; the raw value still backs the vendor-style total fallback.
     return Usage(
-        input_tokens=input_tokens,
+        input_tokens=max(input_tokens - cached, 0),
         output_tokens=output_tokens,
         cache_read_tokens=cached,
         total_tokens=total_tokens or input_tokens + output_tokens,
@@ -644,6 +653,14 @@ def _codex_client_key(
     )
 
 
+def _task_session_id(task: AgentTask) -> str | None:
+    """Conversation id the caller supplied, for result session_id fallback."""
+
+    if task.resume_from is not None:
+        return task.resume_from.session_id
+    return task.session_id
+
+
 def _thread_id(thread: Any) -> str | None:
     value = getattr(thread, "id", None)
     return str(value) if value else None

tests/test_antigravity_adapter.py

@@ -214,6 +214,43 @@ async def test_antigravity_runtime_runs_with_injected_sdk(tmp_path: Path) -> Non
     assert FakeAgent.last_config.kwargs["api_key"] == "key"
 
 
+@pytest.mark.asyncio
+async def test_antigravity_session_id_falls_back_to_task(tmp_path: Path) -> None:
+    class NoIdResponse:
+        def __init__(self, prompt: str) -> None:
+            self.chunks = _chunks(prompt)
+            self.usage_metadata = FakeUsage()
+
+        async def structured_output(self) -> None:
+            return None
+
+    class NoIdAgent:
+        def __init__(self, config: FakeConfig) -> None:
+            self.conversation_id = None
+
+        async def __aenter__(self) -> NoIdAgent:
+            return self
+
+        async def __aexit__(self, *args: object) -> None:
+            return None
+
+        async def chat(self, prompt: str) -> NoIdResponse:
+            return NoIdResponse(prompt)
+
+    runtime = AntigravityAgentRuntime(
+        api_key="key",
+        data_dir=tmp_path,
+        agent_cls=NoIdAgent,
+        config_cls=FakeConfig,
+        types_module=FakeTypes,
+        policy_module=FakePolicy,
+    )
+
+    result = await runtime.run(AgentTask(goal="x", session_id="conv-77"))
+
+    assert result.session_id == "conv-77"
+
+
 @pytest.mark.asyncio
 async def test_antigravity_max_tokens_stop_reason_fails(tmp_path: Path) -> None:
     class TruncatedResponse:

tests/test_codex_adapter.py

@@ -148,7 +148,9 @@ async def test_codex_runtime_runs_with_injected_sdk() -> None:
 
     assert result.output == '{"ok": true}'
     assert result.parsed_output == {"ok": True}
-    assert result.usage.input_tokens == 4
+    # 4 raw input tokens minus 1 cached: input_tokens excludes cache reads.
+    assert result.usage.input_tokens == 3
+    assert result.usage.cache_read_tokens == 1
     assert result.session_id == "thread-new"
     assert sink.events[-1]["name"] == "agent.task.completed"
 
@@ -472,6 +474,56 @@ async def test_codex_sandbox_mapping(filesystem: FilesystemAccess, expected: str
     assert FakeThread.last_run_kwargs["sandbox"] == expected
 
 
+@pytest.mark.asyncio
+async def test_codex_usage_prefers_per_turn_over_cumulative() -> None:
+    # A resumed thread reports cumulative 'total'; the per-turn 'last' is what this
+    # turn actually cost and must be preferred so usage/cost are not over-counted.
+    run_result = {
+        "status": "completed",
+        "final_response": "ok",
+        "usage": {
+            "total": {"input_tokens": 1000, "output_tokens": 2000, "total_tokens": 3000},
+            "last": {"input_tokens": 4, "output_tokens": 6, "total_tokens": 10},
+        },
+    }
+    runtime = make_runtime(run_result)
+
+    result = await runtime.run(AgentTask(goal="x"))
+
+    assert result.usage.input_tokens == 4
+    assert result.usage.output_tokens == 6
+    assert result.usage.total_tokens == 10
+
+
+@pytest.mark.asyncio
+async def test_codex_session_id_falls_back_to_task_when_sdk_omits_it() -> None:
+    # Thread with a falsy id -> result session_id should still reflect the resume id.
+    run_result = {"status": "completed", "final_response": "ok"}
+
+    def codex_factory(*, config: FakeCodexConfig) -> FakeCodex:
+        codex = FakeCodex(config, run_result)
+        original_resume = codex.thread_resume
+
+        async def resume_without_id(thread_id: str, **kwargs: Any) -> FakeThread:
+            thread = await original_resume(thread_id, **kwargs)
+            thread.id = ""
+            return thread
+
+        codex.thread_resume = resume_without_id  # type: ignore[method-assign]
+        return codex
+
+    runtime = CodexAgentRuntime(
+        codex_cls=codex_factory,
+        config_cls=FakeCodexConfig,
+        sandbox_cls=FakeSandbox,
+        approval_mode_cls=FakeApprovalMode,
+    )
+
+    result = await runtime.run(AgentTask(goal="x", session_id="thread-123"))
+
+    assert result.session_id == "thread-123"
+
+
 @pytest.mark.asyncio
 async def test_codex_empty_output_with_tool_calls_succeeds() -> None:
     # Empty final_response is fine when the turn did real tool work; only a
@@ -675,7 +727,7 @@ async def test_codex_rejects_network() -> None:
 
 
 @pytest.mark.asyncio
-async def test_codex_usage_fallback_excludes_cached() -> None:
+async def test_codex_usage_excludes_cached_from_input_tokens() -> None:
     run_result = {
         "status": "completed",
         "final_response": "done",
@@ -685,10 +737,31 @@ async def test_codex_usage_fallback_excludes_cached() -> None:
 
     result = await runtime.run(AgentTask(goal="x"))
 
-    # Fallback total is input + output only (cached is already inside input_tokens).
+    # OpenAI counts cached input inside input_tokens; the Usage contract reports
+    # cache reads separately, so the 4 cached tokens must not appear in both fields.
+    assert result.usage.input_tokens == 6
+    assert result.usage.cache_read_tokens == 4
+    # Fallback total keeps vendor semantics: raw input (incl. cached) + output.
     assert result.usage.total_tokens == 15
 
 
+@pytest.mark.asyncio
+async def test_codex_usage_clamps_when_cached_exceeds_input() -> None:
+    # Defensive clamp: a vendor bug reporting cached > input must not produce a
+    # negative input_tokens.
+    run_result = {
+        "status": "completed",
+        "final_response": "done",
+        "usage": {"total": {"input_tokens": 2, "output_tokens": 5, "cached_input_tokens": 7}},
+    }
+    runtime = make_runtime(run_result)
+
+    result = await runtime.run(AgentTask(goal="x"))
+
+    assert result.usage.input_tokens == 0
+    assert result.usage.cache_read_tokens == 7
+
+
 def test_codex_availability_uses_injected_sdk() -> None:
     runtime = CodexAgentRuntime(codex_cls=FakeCodex, config_cls=FakeCodexConfig)