From 3578d3e364fb287f7fe43126f557e266bac2511e Mon Sep 17 00:00:00 2001 From: Chandrakana Nandi Date: Wed, 19 Aug 2026 10:48:02 -0700 Subject: [PATCH 1/8] llm: stream Anthropic calls and give requests a real timeout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A long authoring turn (Opus thinking over a large prompt) can exceed the SDK's 600s non-streaming ceiling, and `timeout=None` explicitly DISABLES the SDK's timeouts (an explicit None is not not-given), so a socket that died silently mid-call hung the session forever — both observed on Crucible solana_vault runs. Stream every request so bytes keep flowing (no ceiling, no idle window for NAT killers to hit), and bound each httpx phase at 300s so a dead socket surfaces in minutes; for a streamed response that bounds the silence between chunks, not the whole turn. Co-Authored-By: Claude Fable 5 --- composer/llm/anthropic.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/composer/llm/anthropic.py b/composer/llm/anthropic.py index b1101fcf..7b07bbfb 100644 --- a/composer/llm/anthropic.py +++ b/composer/llm/anthropic.py @@ -293,7 +293,16 @@ def builder_for( return ChatAnthropic( model_name=self.model_name, max_tokens_to_sample=opts.tokens, - timeout=None, + # An explicit None DISABLES the SDK's timeouts (None != not-given), so a + # socket that dies silently mid-stream hangs the session forever. A float + # is a per-phase httpx timeout — for a streamed response, the max silence + # between chunks, not a cap on the whole turn. + timeout=300.0, + # Stream every request: a long authoring turn (Opus + thinking on a large + # prompt) can exceed the SDK's 600s non-streaming ceiling, and a silent + # 10-minute wait is long enough for NAT/idle killers to drop the socket + # (surfaces as APIConnectionError mid-run). Streaming keeps bytes flowing. + streaming=True, max_retries=8, stop=None, betas=betas, From 811ad47a469c006175ded048cc42a6fca421f83e Mon Sep 17 00:00:00 2001 From: Chandrakana Nandi Date: Wed, 19 Aug 2026 10:48:27 -0700 Subject: [PATCH 2/8] rag: serialize sentence-transformer encodes, add a device override The nomic model's remote code caches positional tensors per sequence length, so concurrent encodes race: every concurrent crucible_docs_search died with tensor-shape mismatches (the tool degrades to "no results", so authoring ran ungrounded and hallucinated the crucible API), and on Apple Silicon the auto-picked MPS backend segfaulted the whole process inside torch's Metal shader cache. One process-wide lock serializes encodes -- queries are short, so contention is noise -- and COMPOSER_EMBED_DEVICE=cpu lets a Mac host opt out of MPS entirely. Co-Authored-By: Claude Fable 5 --- composer/rag/db.py | 28 +++++++++++++++++++++------- composer/rag/models.py | 10 +++++++++- 2 files changed, 30 insertions(+), 8 deletions(-) diff --git a/composer/rag/db.py b/composer/rag/db.py index 0d1549fd..7d8bade9 100644 --- a/composer/rag/db.py +++ b/composer/rag/db.py @@ -1,8 +1,9 @@ from contextlib import asynccontextmanager -from typing import AsyncIterator, cast, Any, LiteralString, override, TYPE_CHECKING +from typing import AsyncIterator, ClassVar, cast, Any, LiteralString, override, TYPE_CHECKING from dataclasses import dataclass import asyncio import logging +import threading from abc import ABC, abstractmethod import os @@ -96,19 +97,32 @@ async def get_manual_section(self, headers: list[str]) -> str | None: ... + # The model's remote code caches rotary/positional tensors per sequence length, so + # concurrent encodes race and die (shape mismatches on CPU, SIGSEGV in the MPS + # shader cache). One process-wide lock serializes every encode. + _ENCODE_LOCK: ClassVar[threading.Lock] = threading.Lock() + + def _encode_query_locked(self, query: str) -> ndarray: + with self._ENCODE_LOCK: + return cast(ndarray, self.tr.encode_query(query, show_progress_bar=False)) + + def _encode_docs_locked(self, docs: list[str]) -> list[ndarray]: + with self._ENCODE_LOCK: + return cast(list[ndarray], self.tr.encode_document(docs, show_progress_bar=False)) + async def embed_query( self, query: str ) -> ndarray: - return cast(ndarray, await asyncio.to_thread( - self.tr.encode_query, f"search_query: {query}", show_progress_bar=False - )) + return await asyncio.to_thread( + self._encode_query_locked, f"search_query: {query}" + ) async def embed_docs( self, doc: list[BlockChunk] ) -> list[ndarray]: - return cast(list[ndarray], await asyncio.to_thread( - self.tr.encode_document, [f"search_document: {d.chunk}" for d in doc], show_progress_bar=False - )) + return await asyncio.to_thread( + self._encode_docs_locked, [f"search_document: {d.chunk}" for d in doc] + ) type RagConnection = str | AsyncConnectionPool[AsyncConnection[TupleRow]] diff --git a/composer/rag/models.py b/composer/rag/models.py index e5844ad4..676071b0 100644 --- a/composer/rag/models.py +++ b/composer/rag/models.py @@ -1,3 +1,4 @@ +import os from typing import TYPE_CHECKING, override from langchain_core.embeddings import Embeddings @@ -12,7 +13,14 @@ def get_model() -> SentenceTransformer: from sentence_transformers import SentenceTransformer #type: ignore def get_model() -> SentenceTransformer: - return SentenceTransformer('nomic-ai/nomic-embed-text-v1.5', trust_remote_code=True) + # COMPOSER_EMBED_DEVICE overrides the auto-picked device. torch's MPS shader + # cache is not thread-safe under concurrent encodes (SIGSEGV on Apple + # Silicon), so mac hosts should set it to "cpu". + return SentenceTransformer( + 'nomic-ai/nomic-embed-text-v1.5', + trust_remote_code=True, + device=os.environ.get("COMPOSER_EMBED_DEVICE"), + ) except ImportError: # for tests (no ST dependency) def get_model() -> "SentenceTransformer": From 7480549233ac636f8dd56452912c46048039f6b7 Mon Sep 17 00:00:00 2001 From: Chandrakana Nandi Date: Wed, 19 Aug 2026 10:48:33 -0700 Subject: [PATCH 3/8] crucible: revive checkpoint dicts at session readback State read back through the Postgres checkpointer can carry raw dicts where the schema declares models (the serializer's fallback when it cannot reconstruct the class). The readback then died on "'dict' object has no attribute 'property_title'" -- after every component session had already finished its paid authoring -- and the campaign reported empty with exit 0. Revalidate at the boundary so the readback is typed either way. Symptom fix: why the serializer falls back at all is still open, and other checkpoint readers may want the same guard. Co-Authored-By: Claude Fable 5 --- composer/rustapp/session.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/composer/rustapp/session.py b/composer/rustapp/session.py index 44418bda..5c991818 100644 --- a/composer/rustapp/session.py +++ b/composer/rustapp/session.py @@ -941,14 +941,28 @@ async def run_session[K: (RustFormalResult, RustSetupSpec)]( return SessionResult( commentary=state["result"], spec=spec, - skipped=state["skipped"], - property_checks=[(m.property_title, m.checks) for m in state["property_checks"]], - verdicts=state["verdicts"], - ran=state["ran"], + skipped=_revive(SkippedProperty, state["skipped"]), + property_checks=[ + (m.property_title, m.checks) + for m in _revive(PropertyCheckMapping, state["property_checks"]) + ], + verdicts={k: _revive_one(WireVerdict, v) for k, v in state["verdicts"].items()}, + ran=_revive(Target, state["ran"]), expected_failures=state["expected_failures"], ) +def _revive_one[M: BaseModel](ty: type[M], item: M | dict[str, Any]) -> M: + """A checkpoint round-trip may hand back a raw dict where the state declares a model + (the serializer's fallback when it cannot reconstruct the class); revalidate so the + readback is typed either way.""" + return ty.model_validate(item) if isinstance(item, dict) else item + + +def _revive[M: BaseModel](ty: type[M], items: "Sequence[M | dict[str, Any]]") -> list[M]: + return [_revive_one(ty, i) for i in items] + + def _validate_tool(deps: GateDeps, vocab: CheckVocab) -> BaseTool: return ( ValidateSpec.with_template(check=vocab.one, checks=vocab.many) From 3e17f8dbff0439432cb2ad03d4710b39661dfc3f Mon Sep 17 00:00:00 2001 From: Chandrakana Nandi Date: Fri, 21 Aug 2026 15:05:23 -0700 Subject: [PATCH 4/8] diagnostics: read streamed responses' usage from usage_metadata A streamed response carries no raw response_metadata["usage"] dict (the non-streaming path's shape), so with streaming on, UsageCallback recorded zero tokens for every call while costs kept accumulating correctly (CostAccumulator already reads the normalized usage_metadata). Fall back to usage_metadata when the raw dict is absent, translating normalized input (total, cache included) back to the raw shape (cache excluded). Verified against live streamed and non-streamed responses. Co-Authored-By: Claude Fable 5 --- composer/diagnostics/usage_callback.py | 35 ++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/composer/diagnostics/usage_callback.py b/composer/diagnostics/usage_callback.py index c03a1514..02eda78a 100644 --- a/composer/diagnostics/usage_callback.py +++ b/composer/diagnostics/usage_callback.py @@ -20,10 +20,41 @@ from langchain_core.messages import AIMessage from langchain_core.outputs import ChatGeneration, LLMResult -from graphcore.utils import get_token_usage +from graphcore.utils import TokenUsageDict, get_token_usage from composer.diagnostics.timing import get_run_summary +def _usage_of(msg: AIMessage) -> TokenUsageDict: + """Token usage of a response, whichever transport produced it. + + ``get_token_usage`` reads the raw Anthropic ``response_metadata["usage"]`` dict, + which only a non-streamed response carries; a streamed response reports usage + solely through the provider-normalized ``usage_metadata``. Fall back to that, + translating back to the raw shape: normalized ``input_tokens`` is the total + including both cache buckets, where the raw count excludes them.""" + usage = get_token_usage(msg) + if any( + usage[k] + for k in ( + "input_tokens", "output_tokens", + "cache_read_input_tokens", "cache_creation_input_tokens", + ) + ): + return usage + if (um := msg.usage_metadata) is None: + return usage + details = um.get("input_token_details") or {} + cache_read = details.get("cache_read", 0) + cache_creation = details.get("cache_creation", 0) + return { + "model_name": usage["model_name"], + "input_tokens": max(0, um["input_tokens"] - cache_read - cache_creation), + "output_tokens": um["output_tokens"], + "cache_read_input_tokens": cache_read, + "cache_creation_input_tokens": cache_creation, + } + + class UsageCallback(BaseCallbackHandler): """Records each LLM response's token usage into the active ``RunSummary``.""" @@ -45,4 +76,4 @@ def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None: if isinstance(msg, AIMessage): # get_run_summary() returns an inert throwaway outside a run, so this is # a no-op when no autoprove run is active (e.g. ad-hoc model use). - get_run_summary().record_token_usage(get_token_usage(msg)) + get_run_summary().record_token_usage(_usage_of(msg)) From 2de7120cacb9ea8afb94dd27d4e61502bfa5c0ee Mon Sep 17 00:00:00 2001 From: Chandrakana Nandi Date: Fri, 21 Aug 2026 15:35:23 -0700 Subject: [PATCH 5/8] rag: share the encode lock with the sync DefaultEmbedder Review feedback: the indexed research agents embed through DefaultEmbedder's sync Embeddings API, which called the shared sentence-transformer with no serialization -- the same race the async ComposerRAGDB wrappers just got a lock for. The lock now lives in composer.rag.models and both paths take it, so sync and async encodes cannot race each other on the same model instance either. The sync API takes the threading.Lock directly (no to_thread needed -- callers are already off the event loop or tolerate the short block; encodes are sub-second). Verified with interleaved sync-thread + async encode hammer. Co-Authored-By: Claude Fable 5 --- composer/rag/db.py | 14 ++++++-------- composer/rag/models.py | 22 ++++++++++++++++------ 2 files changed, 22 insertions(+), 14 deletions(-) diff --git a/composer/rag/db.py b/composer/rag/db.py index 7d8bade9..f691f2ec 100644 --- a/composer/rag/db.py +++ b/composer/rag/db.py @@ -1,5 +1,5 @@ from contextlib import asynccontextmanager -from typing import AsyncIterator, ClassVar, cast, Any, LiteralString, override, TYPE_CHECKING +from typing import AsyncIterator, cast, Any, LiteralString, override, TYPE_CHECKING from dataclasses import dataclass import asyncio import logging @@ -21,6 +21,7 @@ from composer.rag.types import ManualRef, BlockChunk, ManualSectionHit from composer.rag.text import code_ref_tag +from composer.rag.models import ENCODE_LOCK import sqlite3 @@ -97,17 +98,14 @@ async def get_manual_section(self, headers: list[str]) -> str | None: ... - # The model's remote code caches rotary/positional tensors per sequence length, so - # concurrent encodes race and die (shape mismatches on CPU, SIGSEGV in the MPS - # shader cache). One process-wide lock serializes every encode. - _ENCODE_LOCK: ClassVar[threading.Lock] = threading.Lock() - + # Encodes race when concurrent (see ENCODE_LOCK in composer.rag.models — the lock + # is shared with the sync DefaultEmbedder so neither path can race the other). def _encode_query_locked(self, query: str) -> ndarray: - with self._ENCODE_LOCK: + with ENCODE_LOCK: return cast(ndarray, self.tr.encode_query(query, show_progress_bar=False)) def _encode_docs_locked(self, docs: list[str]) -> list[ndarray]: - with self._ENCODE_LOCK: + with ENCODE_LOCK: return cast(list[ndarray], self.tr.encode_document(docs, show_progress_bar=False)) async def embed_query( diff --git a/composer/rag/models.py b/composer/rag/models.py index 676071b0..e4a974ce 100644 --- a/composer/rag/models.py +++ b/composer/rag/models.py @@ -1,8 +1,16 @@ import os +import threading from typing import TYPE_CHECKING, override from langchain_core.embeddings import Embeddings +# One process-wide gate for every sentence-transformer encode: the model's remote +# code caches positional tensors per sequence length, so concurrent encodes race +# (shape mismatches on CPU, SIGSEGV in torch's MPS shader cache). Shared between +# the sync Embeddings API here and the async ComposerRAGDB wrappers so the two +# paths cannot race each other on the same model instance either. +ENCODE_LOCK = threading.Lock() + # claim we always import ST if TYPE_CHECKING: from sentence_transformers import SentenceTransformer @@ -33,12 +41,14 @@ def __init__(self, model: "SentenceTransformer | None" = None): @override def embed_documents(self, texts: list[str]) -> list[list[float]]: - return self.model.encode_document( - texts - ).tolist() #type: ignore + with ENCODE_LOCK: + return self.model.encode_document( + texts + ).tolist() #type: ignore @override def embed_query(self, text: str) -> list[float]: - return self.model.encode_query( - [text] - ).tolist()[0] #type: ignore + with ENCODE_LOCK: + return self.model.encode_query( + [text] + ).tolist()[0] #type: ignore From cc397217a0ef70d5274316435bf25dca0d43afab Mon Sep 17 00:00:00 2001 From: Chandrakana Nandi Date: Mon, 24 Aug 2026 08:43:32 -0700 Subject: [PATCH 6/8] undo session readback change --- composer/rustapp/session.py | 22 ++++------------------ 1 file changed, 4 insertions(+), 18 deletions(-) diff --git a/composer/rustapp/session.py b/composer/rustapp/session.py index 5c991818..44418bda 100644 --- a/composer/rustapp/session.py +++ b/composer/rustapp/session.py @@ -941,28 +941,14 @@ async def run_session[K: (RustFormalResult, RustSetupSpec)]( return SessionResult( commentary=state["result"], spec=spec, - skipped=_revive(SkippedProperty, state["skipped"]), - property_checks=[ - (m.property_title, m.checks) - for m in _revive(PropertyCheckMapping, state["property_checks"]) - ], - verdicts={k: _revive_one(WireVerdict, v) for k, v in state["verdicts"].items()}, - ran=_revive(Target, state["ran"]), + skipped=state["skipped"], + property_checks=[(m.property_title, m.checks) for m in state["property_checks"]], + verdicts=state["verdicts"], + ran=state["ran"], expected_failures=state["expected_failures"], ) -def _revive_one[M: BaseModel](ty: type[M], item: M | dict[str, Any]) -> M: - """A checkpoint round-trip may hand back a raw dict where the state declares a model - (the serializer's fallback when it cannot reconstruct the class); revalidate so the - readback is typed either way.""" - return ty.model_validate(item) if isinstance(item, dict) else item - - -def _revive[M: BaseModel](ty: type[M], items: "Sequence[M | dict[str, Any]]") -> list[M]: - return [_revive_one(ty, i) for i in items] - - def _validate_tool(deps: GateDeps, vocab: CheckVocab) -> BaseTool: return ( ValidateSpec.with_template(check=vocab.one, checks=vocab.many) From 64b96427fbf0fe94ad94c40f7cfa85b7821c9353 Mon Sep 17 00:00:00 2001 From: Chandrakana Nandi Date: Wed, 26 Aug 2026 16:58:17 -0700 Subject: [PATCH 7/8] rely on get_normalized_token_usage and also make the usage check simpler --- composer/diagnostics/usage_callback.py | 33 ++++++++++---------------- 1 file changed, 13 insertions(+), 20 deletions(-) diff --git a/composer/diagnostics/usage_callback.py b/composer/diagnostics/usage_callback.py index 02eda78a..815a5f94 100644 --- a/composer/diagnostics/usage_callback.py +++ b/composer/diagnostics/usage_callback.py @@ -20,7 +20,7 @@ from langchain_core.messages import AIMessage from langchain_core.outputs import ChatGeneration, LLMResult -from graphcore.utils import TokenUsageDict, get_token_usage +from graphcore.utils import TokenUsageDict, get_normalized_token_usage, get_token_usage from composer.diagnostics.timing import get_run_summary @@ -29,29 +29,22 @@ def _usage_of(msg: AIMessage) -> TokenUsageDict: ``get_token_usage`` reads the raw Anthropic ``response_metadata["usage"]`` dict, which only a non-streamed response carries; a streamed response reports usage - solely through the provider-normalized ``usage_metadata``. Fall back to that, - translating back to the raw shape: normalized ``input_tokens`` is the total - including both cache buckets, where the raw count excludes them.""" + solely through the provider-normalized ``usage_metadata``. Fall back to + ``get_normalized_token_usage`` over that, translating back to the raw shape: + normalized input is the total including both cache buckets, where the raw + count excludes them.""" usage = get_token_usage(msg) - if any( - usage[k] - for k in ( - "input_tokens", "output_tokens", - "cache_read_input_tokens", "cache_creation_input_tokens", - ) - ): + if "usage" in msg.response_metadata: return usage - if (um := msg.usage_metadata) is None: - return usage - details = um.get("input_token_details") or {} - cache_read = details.get("cache_read", 0) - cache_creation = details.get("cache_creation", 0) + norm = get_normalized_token_usage(msg) + cache_read = norm["cache_read_tokens"] + cache_write = norm["cache_write_tokens"] return { - "model_name": usage["model_name"], - "input_tokens": max(0, um["input_tokens"] - cache_read - cache_creation), - "output_tokens": um["output_tokens"], + "model_name": usage["model_name"] or norm["model_name"], + "input_tokens": max(0, norm["total_input_tokens"] - cache_read - cache_write), + "output_tokens": norm["total_output_tokens"], "cache_read_input_tokens": cache_read, - "cache_creation_input_tokens": cache_creation, + "cache_creation_input_tokens": cache_write, } From df5165a19a5917737e670901287e149acb4bd9f4 Mon Sep 17 00:00:00 2001 From: Chandrakana Nandi Date: Thu, 27 Aug 2026 14:44:33 -0700 Subject: [PATCH 8/8] more general token accounting --- composer/diagnostics/timing.py | 28 ++++++++------- composer/diagnostics/usage_callback.py | 35 ++++--------------- composer/spec/source/autosetup.py | 27 +++++++++------ tests/test_token_usage.py | 48 +++++++++++++++----------- 4 files changed, 66 insertions(+), 72 deletions(-) diff --git a/composer/diagnostics/timing.py b/composer/diagnostics/timing.py index dd3c7704..a523c9cc 100644 --- a/composer/diagnostics/timing.py +++ b/composer/diagnostics/timing.py @@ -16,16 +16,18 @@ from typing import AsyncIterator, Iterable, Protocol import uuid -from graphcore.utils import TokenUsageDict +from graphcore.utils import NormalizedTokenUsage @dataclass(frozen=True) class TokenTotals: - """Raw LLM token counts accumulated across one or more calls. + """LLM token counts accumulated across one or more calls. ``input`` is fresh + (uncached) input only; the two cache buckets are counted separately. - ``from_dict`` builds one from a ``graphcore.utils.TokenUsageDict`` (its - ``input_tokens`` / ``output_tokens`` / ``cache_read_input_tokens`` / - ``cache_creation_input_tokens`` keys). + ``from_normalized`` builds one from a ``graphcore.utils.NormalizedTokenUsage``, + whose ``total_input_tokens`` *includes* both cache buckets — they are subtracted + back out so the stored fields (and everything serialized from them) keep the + same meaning they always had. """ input: int = 0 output: int = 0 @@ -45,12 +47,14 @@ def __bool__(self) -> bool: return self.input > 0 or self.output > 0 or self.cache_read > 0 or self.cache_write > 0 @classmethod - def from_dict(cls, u: TokenUsageDict) -> "TokenTotals": + def from_normalized(cls, u: NormalizedTokenUsage) -> "TokenTotals": + cache_read = u["cache_read_tokens"] + cache_write = u["cache_write_tokens"] return TokenTotals( - input=u["input_tokens"], - output=u["output_tokens"], - cache_read=u["cache_read_input_tokens"], - cache_write=u["cache_creation_input_tokens"], + input=max(0, u["total_input_tokens"] - cache_read - cache_write), + output=u["total_output_tokens"], + cache_read=cache_read, + cache_write=cache_write, ) def as_dict(self) -> dict[str, int]: @@ -145,12 +149,12 @@ def record_prover_runtime(self, ms: int, *, task_id: str | None = None) -> None: self._active_prover_reported_by_task.get(task_id, 0) + ms ) - def record_token_usage(self, usage: TokenUsageDict, *, task_id: str | None = None) -> None: + def record_token_usage(self, usage: NormalizedTokenUsage, *, task_id: str | None = None) -> None: """Accumulate one LLM call's token counts into the run-wide per-model totals and (if a task is active) into that task's in-flight bucket, later folded into its ``PhaseRecord`` by ``record_phase``. Defaults attribution to the active task.""" model = usage.get("model_name") or "unknown" - update = TokenTotals.from_dict(usage) + update = TokenTotals.from_normalized(usage) self.token_usage_by_model[model] = self.token_usage_by_model.get(model, TokenTotals()) + update if (task_id := task_id or get_current_task_id()) is not None: bucket = self._active_tokens_by_task.setdefault(task_id, {}) diff --git a/composer/diagnostics/usage_callback.py b/composer/diagnostics/usage_callback.py index 815a5f94..deb5c1c2 100644 --- a/composer/diagnostics/usage_callback.py +++ b/composer/diagnostics/usage_callback.py @@ -20,34 +20,10 @@ from langchain_core.messages import AIMessage from langchain_core.outputs import ChatGeneration, LLMResult -from graphcore.utils import TokenUsageDict, get_normalized_token_usage, get_token_usage +from graphcore.utils import get_normalized_token_usage from composer.diagnostics.timing import get_run_summary -def _usage_of(msg: AIMessage) -> TokenUsageDict: - """Token usage of a response, whichever transport produced it. - - ``get_token_usage`` reads the raw Anthropic ``response_metadata["usage"]`` dict, - which only a non-streamed response carries; a streamed response reports usage - solely through the provider-normalized ``usage_metadata``. Fall back to - ``get_normalized_token_usage`` over that, translating back to the raw shape: - normalized input is the total including both cache buckets, where the raw - count excludes them.""" - usage = get_token_usage(msg) - if "usage" in msg.response_metadata: - return usage - norm = get_normalized_token_usage(msg) - cache_read = norm["cache_read_tokens"] - cache_write = norm["cache_write_tokens"] - return { - "model_name": usage["model_name"] or norm["model_name"], - "input_tokens": max(0, norm["total_input_tokens"] - cache_read - cache_write), - "output_tokens": norm["total_output_tokens"], - "cache_read_input_tokens": cache_read, - "cache_creation_input_tokens": cache_write, - } - - class UsageCallback(BaseCallbackHandler): """Records each LLM response's token usage into the active ``RunSummary``.""" @@ -67,6 +43,9 @@ def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None: return msg = generation.message if isinstance(msg, AIMessage): - # get_run_summary() returns an inert throwaway outside a run, so this is - # a no-op when no autoprove run is active (e.g. ad-hoc model use). - get_run_summary().record_token_usage(_usage_of(msg)) + # The normalized usage_metadata is the one source every transport and + # provider fills in — the raw response_metadata["usage"] dict exists only + # on non-streamed Anthropic responses. get_run_summary() returns an inert + # throwaway outside a run, so this is a no-op when no autoprove run is + # active (e.g. ad-hoc model use). + get_run_summary().record_token_usage(get_normalized_token_usage(msg)) diff --git a/composer/spec/source/autosetup.py b/composer/spec/source/autosetup.py index 338c064b..e3f4e2b7 100644 --- a/composer/spec/source/autosetup.py +++ b/composer/spec/source/autosetup.py @@ -18,7 +18,7 @@ import asyncio from composer.prover.core import ProverOptions -from graphcore.utils import TokenUsageDict +from graphcore.utils import NormalizedTokenUsage from composer.io.context import emit_custom_event # Locators for autosetup's on-disk usage files (certora_autosetup owns that layout). from certora_autosetup.utils.paths import ( @@ -218,22 +218,27 @@ def log_complete(self, returncode: int): # owns that on-disk layout and exposes resolve_autosetup_{llm,prover}_usage_file() to locate them. -def _to_token_usage(model: str, bucket: dict) -> TokenUsageDict: - """Build a graphcore ``TokenUsageDict`` from one AutoSetup rollup bucket, - keeping only the four token fields composer tracks (AutoSetup's ``calls`` - count has no slot in ``TokenTotals`` and is dropped).""" +def _to_token_usage(model: str, bucket: dict) -> NormalizedTokenUsage: + """Build a graphcore ``NormalizedTokenUsage`` from one AutoSetup rollup bucket. + The bucket keeps the raw Anthropic convention (input excludes the cache + buckets) while the normalized form totals them, hence the sum. AutoSetup's + ``calls`` count has no slot in ``TokenTotals`` and is dropped; it reports no + thinking-token split, so ``thinking_tokens`` is 0.""" + cache_read = int(bucket.get("cache_read_input_tokens", 0)) + cache_write = int(bucket.get("cache_creation_input_tokens", 0)) return { "model_name": model, - "input_tokens": int(bucket.get("input_tokens", 0)), - "output_tokens": int(bucket.get("output_tokens", 0)), - "cache_read_input_tokens": int(bucket.get("cache_read_input_tokens", 0)), - "cache_creation_input_tokens": int(bucket.get("cache_creation_input_tokens", 0)), + "total_input_tokens": int(bucket.get("input_tokens", 0)) + cache_read + cache_write, + "total_output_tokens": int(bucket.get("output_tokens", 0)), + "cache_read_tokens": cache_read, + "cache_write_tokens": cache_write, + "thinking_tokens": 0, } -def read_autosetup_usage(project_root: Path) -> list[TokenUsageDict]: +def read_autosetup_usage(project_root: Path) -> list[NormalizedTokenUsage]: """Return AutoSetup's per-model token usage for the most recent run, one - ``TokenUsageDict`` per model — ready to feed straight into + ``NormalizedTokenUsage`` per model — ready to feed straight into ``RunSummary.record_token_usage``. Returns ``[]`` on any failure (file absent — autosetup skipped, cache hit, diff --git a/tests/test_token_usage.py b/tests/test_token_usage.py index 04596223..ea438072 100644 --- a/tests/test_token_usage.py +++ b/tests/test_token_usage.py @@ -28,15 +28,17 @@ ) from composer.diagnostics.usage_callback import UsageCallback from composer.spec.source.autosetup import read_autosetup_usage -from graphcore.utils import TokenUsageDict +from graphcore.utils import NormalizedTokenUsage -def _usage(model: str, i: int, o: int, cr: int, cw: int) -> TokenUsageDict: +def _usage(model: str, i: int, o: int, cr: int, cw: int) -> NormalizedTokenUsage: + """Normalized usage for FRESH input ``i`` — total input includes the caches.""" return { - "input_tokens": i, - "output_tokens": o, - "cache_read_input_tokens": cr, - "cache_creation_input_tokens": cw, + "total_input_tokens": i + cr + cw, + "total_output_tokens": o, + "cache_read_tokens": cr, + "cache_write_tokens": cw, + "thinking_tokens": 0, "model_name": model, } @@ -130,16 +132,17 @@ async def test_token_usage_persisted_to_run_meta_tags(): # --------------------------------------------------------------------------- # def _fake_model(callbacks): + # usage_metadata is the normalized field every transport fills in (a streamed + # response carries no raw response_metadata["usage"] dict at all). Normalized + # input_tokens is the total INCLUDING the cache buckets: 100 fresh + 5 + 2. resp = AIMessage( content="ok", - response_metadata={ - "model_name": "claude-test", - "usage": { - "input_tokens": 100, - "output_tokens": 10, - "cache_read_input_tokens": 5, - "cache_creation_input_tokens": 2, - }, + response_metadata={"model_name": "claude-test"}, + usage_metadata={ + "input_tokens": 107, + "output_tokens": 10, + "total_tokens": 117, + "input_token_details": {"cache_read": 5, "cache_creation": 2}, }, ) return FakeMessagesListChatModel(responses=[resp, resp], callbacks=callbacks) @@ -222,21 +225,24 @@ def _write_autosetup_usage( ) -def test_read_autosetup_usage_returns_token_usage_dicts(tmp_path): +def test_read_autosetup_usage_returns_normalized_usage(tmp_path): _write_autosetup_usage(tmp_path, { "claude-sonnet-4-6": _autosetup_bucket(100, 10, 5, 2), "claude-opus-4": _autosetup_bucket(50, 5, 0, 1), }) by_model = {u["model_name"]: u for u in read_autosetup_usage(tmp_path)} + # The disk bucket keeps the raw convention (input excludes caches); the + # normalized form totals them: 100 fresh + 5 read + 2 write. assert by_model["claude-sonnet-4-6"] == { "model_name": "claude-sonnet-4-6", - "input_tokens": 100, - "output_tokens": 10, - "cache_read_input_tokens": 5, - "cache_creation_input_tokens": 2, + "total_input_tokens": 107, + "total_output_tokens": 10, + "cache_read_tokens": 5, + "cache_write_tokens": 2, + "thinking_tokens": 0, } - assert by_model["claude-opus-4"]["input_tokens"] == 50 + assert by_model["claude-opus-4"]["total_input_tokens"] == 51 assert "calls" not in by_model["claude-sonnet-4-6"] # AutoSetup-only field dropped @@ -285,4 +291,4 @@ def test_autosetup_usage_fallback_newest_dir(tmp_path): timestamp="20260102_000000", write_result=False) usage = read_autosetup_usage(tmp_path) assert [u["model_name"] for u in usage] == ["new"] - assert usage[0]["input_tokens"] == 9 + assert usage[0]["total_input_tokens"] == 9