fix(types): allow model identity in usage metadata

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,60 @@
 WALL_DEADLINE_MONOTONIC_KEY = "wall_deadline_monotonic"
 
 
+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(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
+    model: str
+    prompt_tokens: int
+    completion_tokens: int
+    cache_read_tokens: int
+    cache_write_tokens: int
+    cached_tokens: int
+    cache_creation_tokens: int
+    reasoning_tokens: int
+
+
 def deadline_remaining_s(metadata: Mapping[str, Any] | None) -> float | None:
     """Return seconds to a structural lease or absolute soft deadline.
 
@@ -135,7 +189,11 @@ class TurnContext:
     thinking: str
     tool_calls: list[dict[str, Any]]
     messages: list[Message]
-    usage: dict[str, int] | 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 = ""
@@ -209,7 +267,8 @@ class LLMAttemptContext:
     recovery_action: str = ""
     duration_ms: int = 0
     ttft_ms: int | None = None
-    usage: dict[str, int] | 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
@@ -631,6 +690,8 @@ async def notify_tool_result(
     "ToolCallIntervention",
     "ToolResult",
     "TurnContext",
+    "UsageMetadata",
+    "UsageMetadataExtras",
     "deadline_remaining_s",
     "drain_background_observers",
     "merge_interventions",

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,
@@ -195,15 +196,13 @@ def extract_usage(response: Any) -> dict[str, int | str] | 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.
-        out_dict["reasoning_tokens"] = int(usage.get("reasoning_tokens", 0) or 0)
         return out_dict
 
     rmd = getattr(response, "response_metadata", None) or {}
@@ -214,7 +213,12 @@ def extract_usage(response: Any) -> dict[str, int | str] | 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")
@@ -228,7 +232,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

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.

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"] == ""

tests/test_loop_types.py

@@ -8,6 +8,8 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Mapping
+from typing import Any, get_type_hints
 
 import pytest
 
@@ -19,6 +21,45 @@
 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", "total_tokens", "estimated"}
+    )
+    assert "UsageMetadata" in lt.__all__
+    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"):
     return ToolResult(
         name=name,