From bd0b75706655573f92e7c0287483f1e120a19838 Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Tue, 1 Sep 2026 07:38:18 +0800 Subject: [PATCH 1/3] fix(types): allow model identity in usage metadata --- agent_core/loop_types.py | 23 ++++++++++++++++++++--- agent_core/runtime/loop/_response.py | 7 ++++--- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/agent_core/loop_types.py b/agent_core/loop_types.py index e9a5fee..5c94410 100644 --- a/agent_core/loop_types.py +++ b/agent_core/loop_types.py @@ -8,7 +8,7 @@ from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, field from functools import partial -from typing import Any, Literal, Protocol, cast, runtime_checkable +from typing import Any, Literal, Protocol, TypedDict, cast, runtime_checkable from agent_core.messages import Message @@ -18,6 +18,23 @@ WALL_DEADLINE_MONOTONIC_KEY = "wall_deadline_monotonic" +class UsageMetadata(TypedDict, total=False): + """Normalized observer-facing usage and model attribution.""" + + provider: str + model: str + prompt_tokens: int + completion_tokens: int + input_tokens: int + output_tokens: int + cache_read_tokens: int + cache_write_tokens: int + cached_tokens: int + cache_creation_tokens: int + reasoning_tokens: int + estimated: bool + + def deadline_remaining_s(metadata: Mapping[str, Any] | None) -> float | None: """Return seconds to a structural lease or absolute soft deadline. @@ -135,7 +152,7 @@ class TurnContext: thinking: str tool_calls: list[dict[str, Any]] messages: list[Message] - usage: dict[str, int] | None + usage: UsageMetadata | None metadata: dict[str, Any] # Reasoning recovered from tags leaked into visible content. leaked_reasoning: str = "" @@ -209,7 +226,7 @@ class LLMAttemptContext: recovery_action: str = "" duration_ms: int = 0 ttft_ms: int | None = None - usage: dict[str, int] | None = None + usage: dict[str, Any] | None = None finish_reason: str = "" visible_chars: int = 0 reasoning_chars: int = 0 diff --git a/agent_core/runtime/loop/_response.py b/agent_core/runtime/loop/_response.py index 133dc48..30a4ca9 100644 --- a/agent_core/runtime/loop/_response.py +++ b/agent_core/runtime/loop/_response.py @@ -6,6 +6,7 @@ from typing import Any from agent_core.llm import LLMResponse +from agent_core.loop_types import UsageMetadata from agent_core.messages import Message logger = logging.getLogger(__name__) @@ -118,7 +119,7 @@ def _pick_int(*candidates: Any) -> int: return 0 -def extract_usage(response: Any) -> dict[str, int | str] | None: +def extract_usage(response: Any) -> UsageMetadata | None: """Extract token usage from an LLM response, normalized to OpenAI shape. Returns a dict with keys ``provider`` / ``model`` / ``prompt_tokens`` / @@ -186,7 +187,7 @@ def extract_usage(response: Any) -> dict[str, int | str] | None: # the pre-split cache fields. cache_read = int(usage.get("cached_tokens", 0) or 0) cache_write = int(usage.get("cache_creation_tokens", 0) or 0) - out_dict: dict[str, int | str] = { + out_dict: UsageMetadata = { "provider": provider, "model": response.model or "", "prompt_tokens": inp, @@ -228,7 +229,7 @@ def _build( cached: int, cache_create: int, reasoning: int, - ) -> dict[str, int | str]: + ) -> UsageMetadata: # ``cached`` carries cache READ; ``cache_create`` carries cache # WRITE. The legacy ``cached_tokens`` / ``cache_creation_tokens`` # keys are kept as a derived sum and an alias respectively so From 1324c6e8b3ab264d1267f03aca6f3a02fd6e3d2b Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Tue, 1 Sep 2026 08:38:51 +0800 Subject: [PATCH 2/3] fix(types): enforce normalized usage contract --- agent_core/loop_types.py | 22 ++++++++++++++++------ agent_core/runtime/loop/_response.py | 2 +- tests/test_loop_types.py | 26 ++++++++++++++++++++++++++ 3 files changed, 43 insertions(+), 7 deletions(-) diff --git a/agent_core/loop_types.py b/agent_core/loop_types.py index 5c94410..e63b411 100644 --- a/agent_core/loop_types.py +++ b/agent_core/loop_types.py @@ -18,21 +18,30 @@ WALL_DEADLINE_MONOTONIC_KEY = "wall_deadline_monotonic" -class UsageMetadata(TypedDict, total=False): - """Normalized observer-facing usage and model attribution.""" +class _OptionalUsageMetadata(TypedDict, total=False): + """Host-specific usage aliases and provenance.""" + + input_tokens: int + output_tokens: int + estimated: bool + + +class UsageMetadata(_OptionalUsageMetadata): + """Normalized observer-facing usage and model attribution. + + The normalized fields are required so observers can index them directly. + Host-specific aliases inherited above remain optional. + """ provider: str model: str prompt_tokens: int completion_tokens: int - input_tokens: int - output_tokens: int cache_read_tokens: int cache_write_tokens: int cached_tokens: int cache_creation_tokens: int reasoning_tokens: int - estimated: bool def deadline_remaining_s(metadata: Mapping[str, Any] | None) -> float | None: @@ -226,7 +235,7 @@ class LLMAttemptContext: recovery_action: str = "" duration_ms: int = 0 ttft_ms: int | None = None - usage: dict[str, Any] | None = None + usage: UsageMetadata | None = None finish_reason: str = "" visible_chars: int = 0 reasoning_chars: int = 0 @@ -648,6 +657,7 @@ async def notify_tool_result( "ToolCallIntervention", "ToolResult", "TurnContext", + "UsageMetadata", "deadline_remaining_s", "drain_background_observers", "merge_interventions", diff --git a/agent_core/runtime/loop/_response.py b/agent_core/runtime/loop/_response.py index 30a4ca9..6d2e14d 100644 --- a/agent_core/runtime/loop/_response.py +++ b/agent_core/runtime/loop/_response.py @@ -196,6 +196,7 @@ def extract_usage(response: Any) -> UsageMetadata | None: "cache_write_tokens": cache_write, "cached_tokens": cache_read + cache_write, "cache_creation_tokens": cache_write, + "reasoning_tokens": int(usage.get("reasoning_tokens", 0) or 0), } # Reasoning/thinking tokens (Anthropic extended thinking / OpenAI # reasoning models). They are part of completion_tokens but surfaced @@ -204,7 +205,6 @@ def extract_usage(response: Any) -> UsageMetadata | None: # below must return the SAME key set, or a consumer indexing # ``usage["reasoning_tokens"]`` works on one response object and # raises KeyError on the other. - out_dict["reasoning_tokens"] = int(usage.get("reasoning_tokens", 0) or 0) return out_dict rmd = getattr(response, "response_metadata", None) or {} diff --git a/tests/test_loop_types.py b/tests/test_loop_types.py index 4ad1979..6d3f130 100644 --- a/tests/test_loop_types.py +++ b/tests/test_loop_types.py @@ -8,6 +8,7 @@ from __future__ import annotations import asyncio +from typing import get_type_hints import pytest @@ -19,6 +20,31 @@ merge_interventions = lt.merge_interventions +def test_usage_metadata_contract_marks_normalized_keys_required(): + assert lt.UsageMetadata.__required_keys__ == frozenset( + { + "provider", + "model", + "prompt_tokens", + "completion_tokens", + "cache_read_tokens", + "cache_write_tokens", + "cached_tokens", + "cache_creation_tokens", + "reasoning_tokens", + } + ) + assert lt.UsageMetadata.__optional_keys__ == frozenset( + {"input_tokens", "output_tokens", "estimated"} + ) + assert "UsageMetadata" in lt.__all__ + + +def test_attempt_context_reuses_normalized_usage_contract(): + usage_type = get_type_hints(lt.LLMAttemptContext)["usage"] + assert usage_type == lt.UsageMetadata | None + + def _tool_result(name: str = "t", result: str = "r"): return ToolResult( name=name, From b047aa0f1f6f351f361da4e5bafa5993a589897c Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Tue, 1 Sep 2026 08:59:53 +0800 Subject: [PATCH 3/3] review: keep observer usage fields open, pin the key-set invariant MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-up on the normalized usage contract. Narrowing TurnContext.usage / LLMAttemptContext.usage to the UsageMetadata TypedDict broke the producers it was meant to serve. Products construct those contexts, and their shapes are ones a TypedDict rejects in strict *and* standard mode: partial literals in test doubles, alias-only budget shapes, and the dict(event["usage"]) re-wrap the ApodexHarness loop stamps provider/model onto before building LLMAttemptContext. Widen both fields to Mapping[str, Any] | None — Mapping, not dict[str, Any], because it is covariant and so also accepts a UsageMetadata, which a dict field would reject. extract_usage keeps the precise return type; consumers wanting the shape annotate their own parameter. A TypedDict validates nothing at runtime, so the key-set invariant the moved comment used to guard was left unenforced: pin each of the three branches against UsageMetadata.__required_keys__ instead. Also: - coerce model to str on the legacy path, symmetric with provider — it is read off an untyped gateway metadata dict and UsageMetadata declares it str; - publish UsageMetadataExtras and add total_tokens, so hosts can describe the aliases they stamp (ApodexHarness native clients carry total_tokens); - reverse the field-annotation test to assert the mapping stays open, covering TurnContext as well as LLMAttemptContext; - move the reasoning_tokens comment onto its key; - document the contract in docs/llm-runtime-boundary.md. Validation: ruff, pyright 0 errors, 285 passed. Nine real downstream call sites from FrontierAgent and ApodexHarness reproduced as a pyright probe: 0 errors under standard, down from 6. Co-Authored-By: Claude Opus 5 (1M context) --- agent_core/loop_types.py | 52 ++++++++++++++---- agent_core/runtime/loop/_response.py | 19 ++++--- docs/llm-runtime-boundary.md | 39 ++++++++++++++ tests/test_llm_runtime_extract_usage.py | 71 +++++++++++++++++++++++++ tests/test_loop_types.py | 29 +++++++--- 5 files changed, 186 insertions(+), 24 deletions(-) diff --git a/agent_core/loop_types.py b/agent_core/loop_types.py index e63b411..db1f64e 100644 --- a/agent_core/loop_types.py +++ b/agent_core/loop_types.py @@ -18,19 +18,47 @@ WALL_DEADLINE_MONOTONIC_KEY = "wall_deadline_monotonic" -class _OptionalUsageMetadata(TypedDict, total=False): - """Host-specific usage aliases and provenance.""" +class UsageMetadataExtras(TypedDict, total=False): + """Host-specific usage aliases and provenance. + + ``extract_usage`` never emits these; they exist so a host that stamps + its own aliases onto a usage mapping can still describe the result as + a :class:`UsageMetadata`. ApodexHarness' budget observer reads the + ``input_tokens`` / ``output_tokens`` aliases, and its native clients + carry ``total_tokens`` on ``LLMResponse.usage``. + + Optional keys are NOT directly indexable under a type checker + (``reportTypedDictNotRequiredAccess``) — read them with ``.get(key, 0)``. + """ input_tokens: int output_tokens: int + total_tokens: int estimated: bool -class UsageMetadata(_OptionalUsageMetadata): - """Normalized observer-facing usage and model attribution. - - The normalized fields are required so observers can index them directly. - Host-specific aliases inherited above remain optional. +class UsageMetadata(UsageMetadataExtras): + """What :func:`agent_core.runtime.loop.extract_usage` returns. + + The normalized fields below are required: every non-``None`` return + from ``extract_usage`` carries all of them, zero-filled when the + provider reported nothing, so a consumer can index them without + probing which response shape it was handed. + ``tests/test_llm_runtime_extract_usage.py`` pins that at runtime + against ``__required_keys__``. + + This is a *producer* type, deliberately narrower than the + ``usage`` fields on :class:`TurnContext` / :class:`LLMAttemptContext`. + Those stay ``Mapping[str, Any] | None`` because hosts build usage + mappings of their own — partial dicts in test doubles, alias-only + shapes, ``dict(event["usage"])`` re-wraps out of a ``dict[str, Any]`` + attempt event — and neither ``dict[str, int]`` nor ``dict[str, Any]`` + is assignable to a TypedDict. A consumer that wants the precise shape + annotates its own parameter as ``UsageMetadata``; the loop does not + force it on producers. Extra keys beyond those declared here are + likewise a host's business: a TypedDict cannot express "open" on + Python 3.12 (PEP 728 lands later), which is the other reason the + dataclass fields stay a plain mapping. """ provider: str @@ -161,7 +189,11 @@ class TurnContext: thinking: str tool_calls: list[dict[str, Any]] messages: list[Message] - usage: UsageMetadata | None + # Read-only mapping, not ``UsageMetadata``, on purpose: producers are + # hosts, and a TypedDict rejects the ``dict[str, int]`` / ``dict[str, Any]`` + # shapes they build. ``UsageMetadata`` documents what ``extract_usage`` + # puts here; annotate a consumer's own parameter with it for precision. + usage: Mapping[str, Any] | None metadata: dict[str, Any] # Reasoning recovered from tags leaked into visible content. leaked_reasoning: str = "" @@ -235,7 +267,8 @@ class LLMAttemptContext: recovery_action: str = "" duration_ms: int = 0 ttft_ms: int | None = None - usage: UsageMetadata | None = None + # See ``TurnContext.usage`` for why this is a plain mapping. + usage: Mapping[str, Any] | None = None finish_reason: str = "" visible_chars: int = 0 reasoning_chars: int = 0 @@ -658,6 +691,7 @@ async def notify_tool_result( "ToolResult", "TurnContext", "UsageMetadata", + "UsageMetadataExtras", "deadline_remaining_s", "drain_background_observers", "merge_interventions", diff --git a/agent_core/runtime/loop/_response.py b/agent_core/runtime/loop/_response.py index 6d2e14d..826c210 100644 --- a/agent_core/runtime/loop/_response.py +++ b/agent_core/runtime/loop/_response.py @@ -196,15 +196,13 @@ def extract_usage(response: Any) -> UsageMetadata | None: "cache_write_tokens": cache_write, "cached_tokens": cache_read + cache_write, "cache_creation_tokens": cache_write, + # Reasoning/thinking tokens (Anthropic extended thinking / + # OpenAI reasoning models). Part of completion_tokens, but + # surfaced separately for cost / analysis; the client's usage + # dict carries them. Present even as 0 — see the key-set + # invariant on the ``UsageMetadata`` return type. "reasoning_tokens": int(usage.get("reasoning_tokens", 0) or 0), } - # Reasoning/thinking tokens (Anthropic extended thinking / OpenAI - # reasoning models). They are part of completion_tokens but surfaced - # separately for cost / analysis; the client's usage dict carries them. - # Always present, including as 0: this branch and the legacy shapes - # below must return the SAME key set, or a consumer indexing - # ``usage["reasoning_tokens"]`` works on one response object and - # raises KeyError on the other. return out_dict rmd = getattr(response, "response_metadata", None) or {} @@ -215,7 +213,12 @@ def extract_usage(response: Any) -> UsageMetadata | None: # own ``model_name`` lands on non-streaming responses but not on # streamed usage chunks). Falling through to it keeps streaming usage # attribution alive. - model = ( + # ``str(...)`` for the same reason ``provider`` gets it: these come off + # an untyped provider/gateway metadata dict, and ``UsageMetadata`` declares + # ``model`` as ``str``. A gateway echoing a non-string here — a nested + # dict, say — would otherwise put that object behind a field consumers + # format as text. + model = str( rmd.get("model_name") or rmd.get("model") or rmd.get("model_actually_used") diff --git a/docs/llm-runtime-boundary.md b/docs/llm-runtime-boundary.md index 724982f..2fd0438 100644 --- a/docs/llm-runtime-boundary.md +++ b/docs/llm-runtime-boundary.md @@ -33,3 +33,42 @@ translate semantic retry intent into provider-specific request fields. Model profiles, provider client construction, tool parsing/execution, and the agent loop remain product-owned in this phase. + +## The usage contract + +`extract_usage` returns a `UsageMetadata` (`agent_core.loop_types`). Every +non-`None` return carries all nine normalized keys — `provider`, `model`, +`prompt_tokens`, `completion_tokens`, `cache_read_tokens`, +`cache_write_tokens`, `cached_tokens`, `cache_creation_tokens`, +`reasoning_tokens` — zero-filled when the provider reported nothing, whichever +of the three response shapes it parsed. That is a runtime invariant, not just a +declaration: a TypedDict validates nothing at import time, so +`tests/test_llm_runtime_extract_usage.py` pins each branch's key set against +`UsageMetadata.__required_keys__`. Consumers can index those nine directly. + +The `usage` fields on `TurnContext` and `LLMAttemptContext` are deliberately +*wider* — `Mapping[str, Any] | None`, not the TypedDict. Products, not +AgentCore, construct those contexts, and the shapes they hand over are ones a +TypedDict rejects outright: + +- partial literals in test doubles — `usage={"prompt_tokens": 182_000}`; +- alias-only shapes — `usage={"input_tokens": …, "output_tokens": …}`, which + ApodexHarness' budget observer reads; +- `dict(event["usage"])` re-wraps out of a `dict[str, Any]` attempt event, + which the ApodexHarness loop stamps `provider` / `model` onto before + constructing `LLMAttemptContext`. + +Neither `dict[str, int]` nor `dict[str, Any]` is assignable to a TypedDict, in +strict *or* standard mode, so narrowing these fields would break the product +loops the normalized contract exists to serve. `Mapping` rather than +`dict[str, Any]` because it is covariant in its value type: it accepts all of +the above *and* a `UsageMetadata`, which a `dict[str, Any]` field would reject. +A consumer wanting the precise shape annotates its own parameter +`UsageMetadata`; the boundary does not force that on producers. Hosts stamping +extra keys is likewise their business — a TypedDict cannot express "open" on +Python 3.12 (PEP 728 lands later). `UsageMetadataExtras` declares the aliases +hosts are known to add, for anyone who wants to describe such a mapping as a +`UsageMetadata`. + +`tests/test_loop_types.py` asserts these fields stay an open mapping, so a +later well-meant tightening fails loudly. diff --git a/tests/test_llm_runtime_extract_usage.py b/tests/test_llm_runtime_extract_usage.py index 4f81620..e26cc95 100644 --- a/tests/test_llm_runtime_extract_usage.py +++ b/tests/test_llm_runtime_extract_usage.py @@ -31,7 +31,10 @@ from types import SimpleNamespace +import pytest + from agent_core.llm import LLMResponse +from agent_core.loop_types import UsageMetadata from agent_core.runtime.loop._response import _pick_int from agent_core.runtime.loop.llm_client import extract_usage @@ -603,3 +606,71 @@ def __getitem__(self, k): assert u["prompt_tokens"] == 7 assert u["completion_tokens"] == 3 assert u["model"] == "x" + + +# --- key-set invariant --------------------------------------------------- +# +# The three response shapes are parsed by three separate branches, and a +# consumer indexing ``usage["reasoning_tokens"]`` must not work on one +# response object and raise KeyError on another. ``UsageMetadata`` states +# that invariant in the type system; a TypedDict validates nothing at +# runtime, so pin it here against the declared required keys — the type and +# all three branches move together or this fails. + + +_NATIVE = LLMResponse( + content="x", + model="m", + usage={"prompt_tokens": 1, "completion_tokens": 2}, + response_metadata={"provider_actually_used": "p"}, +) +_CANONICAL = SimpleNamespace( + usage_metadata={"input_tokens": 1, "output_tokens": 2}, + response_metadata={"model_name": "m"}, +) +_RAW = SimpleNamespace( + usage_metadata=None, + response_metadata={ + "token_usage": {"prompt_tokens": 1, "completion_tokens": 2}, + "model_name": "m", + }, +) + + +@pytest.mark.parametrize( + ("label", "response"), + [("native", _NATIVE), ("canonical", _CANONICAL), ("raw", _RAW)], +) +def test_every_branch_returns_exactly_the_required_key_set( + label: str, response: object, +) -> None: + usage = extract_usage(response) + assert usage is not None, label + assert set(usage) == UsageMetadata.__required_keys__, label + + +def test_model_is_coerced_to_str_on_the_legacy_path() -> None: + """``UsageMetadata`` declares ``model: str``; the source dict is untyped. + + A gateway echoing a non-string ``model_name`` would land that object + verbatim behind a field consumers format as text. ``provider`` was + already coerced; ``model`` was the asymmetric one. + """ + usage = extract_usage( + _resp( + usage_metadata={"input_tokens": 1, "output_tokens": 2}, + response_metadata={"model_name": {"id": "oops"}}, + ) + ) + assert usage is not None + assert isinstance(usage["model"], str) + assert usage["model"] == "{'id': 'oops'}" + + +def test_model_falls_back_to_empty_string_not_none() -> None: + usage = extract_usage( + _resp(usage_metadata={"input_tokens": 1, "output_tokens": 2}, + response_metadata={}) + ) + assert usage is not None + assert usage["model"] == "" diff --git a/tests/test_loop_types.py b/tests/test_loop_types.py index 6d3f130..6822368 100644 --- a/tests/test_loop_types.py +++ b/tests/test_loop_types.py @@ -8,7 +8,8 @@ from __future__ import annotations import asyncio -from typing import get_type_hints +from collections.abc import Mapping +from typing import Any, get_type_hints import pytest @@ -35,14 +36,28 @@ def test_usage_metadata_contract_marks_normalized_keys_required(): } ) assert lt.UsageMetadata.__optional_keys__ == frozenset( - {"input_tokens", "output_tokens", "estimated"} + {"input_tokens", "output_tokens", "total_tokens", "estimated"} ) assert "UsageMetadata" in lt.__all__ - - -def test_attempt_context_reuses_normalized_usage_contract(): - usage_type = get_type_hints(lt.LLMAttemptContext)["usage"] - assert usage_type == lt.UsageMetadata | None + assert "UsageMetadataExtras" in lt.__all__ + + +@pytest.mark.parametrize("cls_name", ["TurnContext", "LLMAttemptContext"]) +def test_observer_usage_field_stays_an_open_mapping(cls_name: str): + """``usage`` must NOT be typed as the ``UsageMetadata`` TypedDict. + + Hosts are the producers of these contexts and they hand over shapes a + TypedDict rejects outright: partial literals in test doubles + (``{"prompt_tokens": 182_000}``), alias-only budget shapes + (``{"input_tokens": .., "output_tokens": ..}``), and + ``dict(event["usage"])`` re-wraps whose value type is ``Any``. None of + those are assignable to a TypedDict, under strict OR standard mode, so + narrowing this field breaks the exact product loops the normalized + contract was introduced to unblock. ``UsageMetadata`` stays the + producer-side return type of ``extract_usage`` instead. + """ + usage_type = get_type_hints(getattr(lt, cls_name))["usage"] + assert usage_type == Mapping[str, Any] | None def _tool_result(name: str = "t", result: str = "r"):