Skip to content

fix(types): allow model identity in usage metadata - #9

Merged
zhanghanduo merged 3 commits into
mainfrom
fix/usage-metadata-type
Sep 1, 2026
Merged

zhanghanduo merged 3 commits into
mainfrom
fix/usage-metadata-type

Conversation

@zhanghanduo

@zhanghanduo zhanghanduo commented Aug 31, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Align the loop observer usage contract with extract_usage.

Normalized usage contains integer token counters plus string provider and
model identity. TurnContext and LLMAttemptContext still declared every
value as int, causing product agent loops to fail Pyright once they consumed
the shared LLM runtime.

Contract shape

The fix is layered rather than uniform — the producer side is precise, the
observer fields stay open:

Site Type Why
extract_usage return UsageMetadata (9 required keys) sole producer; key set is verifiable
TurnContext.usage, LLMAttemptContext.usage Mapping[str, Any] | None products are the producers and build shapes a TypedDict rejects
a consumer wanting precision annotates its own parameter UsageMetadata opt-in, not imposed

An earlier revision narrowed the two dataclass fields to the UsageMetadata
TypedDict. That broke the producers the contract exists to serve: neither
dict[str, int] nor dict[str, Any] is assignable to a TypedDict, in strict
or standard mode. Real downstream shapes that fail under it:

  • partial literals in test doubles — usage={"prompt_tokens": 182_000}
    (FrontierAgent), usage={"prompt_tokens": 210_000} (ApodexHarness);
  • alias-only budget shapes — usage={"input_tokens": …, "output_tokens": …},
    which miroharness/components/observers/budget_observer.py reads;
  • production: dict(event["usage"]) re-wrapped out of a dict[str, Any]
    attempt event, stamped with provider / model and passed to
    LLMAttemptContext — miroharness/core/runtime/loop/agent_loop.py:537-543.

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 (TypedDict and dict are mutually
non-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

  • Key-set invariant is now enforced at runtime. A TypedDict validates
    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 raises
    KeyError on another". Each branch is pinned against
    UsageMetadata.__required_keys__.
  • model is coerced to str on the legacy path, symmetric with
    provider. It is read off an untyped gateway metadata dict, so a gateway
    echoing a non-string model_name would land that object behind a field
    UsageMetadata declares str and consumers format as text.
  • UsageMetadataExtras is public and gained total_tokens, so a host can
    describe the aliases it stamps (ApodexHarness native clients carry
    total_tokens on LLMResponse.usage). It was private, leaving the optional
    half unreachable.
  • The field-annotation test asserts the mapping stays open — reversed from
    asserting the TypedDict — and covers TurnContext, not just
    LLMAttemptContext, so a later well-meant tightening fails loudly.
  • docs/llm-runtime-boundary.md documents the contract and the reasoning.

Validation

  • Ruff passed
  • Pyright: 0 errors, 0 warnings
  • pytest: 285 passed
  • uv lock --check
  • wheel and sdist build passed
  • Nine real downstream call sites from FrontierAgent
    (refactor/agent-core-runtime-contracts) and ApodexHarness
    (refactor/agent-core-llm-runtime) reproduced as a Pyright probe against
    this branch: 0 errors under standard mode, down from 6. The 4 remaining
    under strict are reportUnknown* on ApodexHarness' own
    event.get("usage") returning Any — independent of this field's
    declaration and present before this PR.

@zhanghanduo
zhanghanduo force-pushed the fix/usage-metadata-type branch from 3566b3d to bd0b757 Compare August 31, 2026 23:42
zhanghanduo and others added 2 commits September 1, 2026 08:38
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>
@zhanghanduo
zhanghanduo merged commit c417534 into main Sep 1, 2026
1 check passed
@zhanghanduo
zhanghanduo deleted the fix/usage-metadata-type branch September 1, 2026 01:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant