diff --git a/src/agent_runtime_kit/adapters/antigravity.py b/src/agent_runtime_kit/adapters/antigravity.py index d8b2d96..589e66d 100644 --- a/src/agent_runtime_kit/adapters/antigravity.py +++ b/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 ) diff --git a/src/agent_runtime_kit/adapters/codex.py b/src/agent_runtime_kit/adapters/codex.py index e728990..c1bf150 100644 --- a/src/agent_runtime_kit/adapters/codex.py +++ b/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 diff --git a/tests/test_antigravity_adapter.py b/tests/test_antigravity_adapter.py index 2425842..4b3069b 100644 --- a/tests/test_antigravity_adapter.py +++ b/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: diff --git a/tests/test_codex_adapter.py b/tests/test_codex_adapter.py index 9af9475..7d8ddd8 100644 --- a/tests/test_codex_adapter.py +++ b/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)