fix(types): allow model identity in usage metadata - #9
Merged
Merged
Conversation
zhanghanduo
force-pushed
the
fix/usage-metadata-type
branch
from
August 31, 2026 23:42
3566b3d to
bd0b757
Compare
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) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Align the loop observer usage contract with
extract_usage.Normalized usage contains integer token counters plus string
providerandmodelidentity.TurnContextandLLMAttemptContextstill declared everyvalue as
int, causing product agent loops to fail Pyright once they consumedthe shared LLM runtime.
Contract shape
The fix is layered rather than uniform — the producer side is precise, the
observer fields stay open:
extract_usagereturnUsageMetadata(9 required keys)TurnContext.usage,LLMAttemptContext.usageMapping[str, Any] | NoneUsageMetadataAn earlier revision narrowed the two dataclass fields to the
UsageMetadataTypedDict. That broke the producers the contract exists to serve: neither
dict[str, int]nordict[str, Any]is assignable to a TypedDict, in strictor standard mode. Real downstream shapes that fail under it:
usage={"prompt_tokens": 182_000}(FrontierAgent),
usage={"prompt_tokens": 210_000}(ApodexHarness);usage={"input_tokens": …, "output_tokens": …},which
miroharness/components/observers/budget_observer.pyreads;dict(event["usage"])re-wrapped out of adict[str, Any]attempt event, stamped with
provider/modeland passed toLLMAttemptContext—miroharness/core/runtime/loop/agent_loop.py:537-543.Mappingrather thandict[str, Any]because it is covariant in its valuetype: it accepts all of the above and a
UsageMetadata, which adict[str, Any]field would reject (TypedDict anddictare mutuallynon-assignable). Hosts stamping extra keys is likewise their business — a
TypedDict cannot express "open" on Python 3.12 (PEP 728 lands later).
Also in this PR
nothing at import time, so declaring the shape left the invariant it replaced
unguarded — "all three branches return the same key set, or a consumer
indexing
usage["reasoning_tokens"]works on one response object and raisesKeyErroron another". Each branch is pinned againstUsageMetadata.__required_keys__.modelis coerced tostron the legacy path, symmetric withprovider. It is read off an untyped gateway metadata dict, so a gatewayechoing a non-string
model_namewould land that object behind a fieldUsageMetadatadeclaresstrand consumers format as text.UsageMetadataExtrasis public and gainedtotal_tokens, so a host candescribe the aliases it stamps (ApodexHarness native clients carry
total_tokensonLLMResponse.usage). It was private, leaving the optionalhalf unreachable.
asserting the TypedDict — and covers
TurnContext, not justLLMAttemptContext, so a later well-meant tightening fails loudly.docs/llm-runtime-boundary.mddocuments the contract and the reasoning.Validation
uv lock --check(
refactor/agent-core-runtime-contracts) and ApodexHarness(
refactor/agent-core-llm-runtime) reproduced as a Pyright probe againstthis branch: 0 errors under standard mode, down from 6. The 4 remaining
under strict are
reportUnknown*on ApodexHarness' ownevent.get("usage")returningAny— independent of this field'sdeclaration and present before this PR.