diff --git a/coworker/engine.py b/coworker/engine.py index 78a88316..3ce77d1e 100644 --- a/coworker/engine.py +++ b/coworker/engine.py @@ -353,13 +353,15 @@ class TurnEngine: if turn is None: turn = AssistantTurn() - self.messages.append(_assistant_message(turn)) + self.messages.append(_assistant_message(turn, model=self.model)) payload: dict[str, Any] = { "text": turn.text, "tool_calls": [tc.name for tc in turn.tool_calls], } if turn.reasoning: payload["reasoning"] = turn.reasoning + if turn.usage is not None: + payload["usage"] = {"model": self.model, **turn.usage.as_dict()} yield Event(EventType.ASSISTANT_MESSAGE, payload) if not turn.tool_calls: @@ -886,10 +888,11 @@ class TurnEngine: persisted/replayed. """ # Strip the display-only sidecars — `source` (connector cards), `_display` - # (e.g. filter-hidden counts), `ts` (append-time timestamps), and `reasoning` - # (thinking text) — copying only messages that carry one. Whole `notice` messages - # (error/interrupted/model-switch markers) are display-only too: dropped entirely. - _SIDECARS = ("source", "_display", "ts", "reasoning") + # (e.g. filter-hidden counts), `ts` (append-time timestamps), `reasoning` + # (thinking text), and `usage` (token counts) — copying only messages that carry + # one. Whole `notice` messages (error/interrupted/model-switch markers) are + # display-only too: dropped entirely. + _SIDECARS = ("source", "_display", "ts", "reasoning", "usage") out = [ ( {k: v for k, v in msg.items() if k not in _SIDECARS} @@ -982,12 +985,17 @@ class TurnEngine: return out -def _assistant_message(turn: AssistantTurn) -> dict[str, Any]: +def _assistant_message(turn: AssistantTurn, model: Optional[str] = None) -> dict[str, Any]: message: dict[str, Any] = { "role": "assistant", "content": turn.text or "", "ts": time.time(), } + if turn.usage is not None: + # Display/aggregation sidecar (like `reasoning`): persisted with the message, + # stripped before provider calls. Tagged with the model that produced it so + # per-model rollups survive mid-session model switches. + message["usage"] = {"model": model, **turn.usage.as_dict()} if turn.reasoning: # Display-only thinking text — rendered by the GUI, stripped for every provider # (`_outbound_messages`); provider-private replay blocks go via `extras` instead. diff --git a/coworker/providers/anthropic_provider.py b/coworker/providers/anthropic_provider.py index 579be5bb..05bc8f3a 100644 --- a/coworker/providers/anthropic_provider.py +++ b/coworker/providers/anthropic_provider.py @@ -28,10 +28,23 @@ from .base import ( ModelCapabilities, ProviderClient, StreamChunk, + TokenUsage, ToolCall, ) from .capabilities import capabilities_for + +def _usage_from(usage: Any) -> Optional[TokenUsage]: + """Messages-API usage object → normalized counts (input_tokens excludes cache).""" + if usage is None: + return None + return TokenUsage( + input=int(getattr(usage, "input_tokens", 0) or 0), + output=int(getattr(usage, "output_tokens", 0) or 0), + cache_read=int(getattr(usage, "cache_read_input_tokens", 0) or 0), + cache_write=int(getattr(usage, "cache_creation_input_tokens", 0) or 0), + ) + # Required by the Messages API; a ceiling, not a spend target. DEFAULT_MAX_TOKENS = 16000 @@ -325,6 +338,32 @@ def convert_tools(tools: Optional[list[dict[str, Any]]]) -> list[dict[str, Any]] return converted +def _add_cache_breakpoints(kwargs: dict[str, Any]) -> None: + """Opt the request into prompt caching (5-minute ephemeral, prefix-matched). + + Two breakpoints, the standard agent-loop shape: + - last system block — caches tools + system together (tools render first); + - last content block of the final message — caches the whole conversation + prefix, so each request re-reads the previous turns' cache and writes only + the new tail (append-only history keeps the prefix byte-identical). + + Outbound-only: the canonical history never carries `cache_control` (the final + message's blocks are freshly built by convert_messages — thinking replays sit + in earlier assistant turns — and the marked block is copied, not mutated). + Prefixes under the model's cacheable minimum silently don't cache; reads bill + ~0.1x and show up as `cache_read_input_tokens` (the metering's cache_read). + """ + marker = {"type": "ephemeral"} + system = kwargs.get("system") + if isinstance(system, str) and system: + kwargs["system"] = [{"type": "text", "text": system, "cache_control": marker}] + messages = kwargs.get("messages") or [] + if messages: + content = messages[-1].get("content") + if isinstance(content, list) and content: + content[-1] = {**content[-1], "cache_control": marker} + + def _reasoning_text(thinking_blocks: list[dict[str, Any]]) -> Optional[str]: """Display text for the GUI's disclosure — thinking text only (redacted stays opaque).""" text = "".join( @@ -412,6 +451,7 @@ class AnthropicProvider(ProviderClient): kwargs["system"] = system if tools: kwargs["tools"] = convert_tools(tools) + _add_cache_breakpoints(kwargs) return kwargs def complete( @@ -474,6 +514,7 @@ class AnthropicProvider(ProviderClient): raw=response, reasoning=_reasoning_text(thinking_blocks), extras=_thinking_extras(thinking_blocks), + usage=_usage_from(getattr(response, "usage", None)), ) def capabilities(self, model: str) -> ModelCapabilities: @@ -507,11 +548,18 @@ class AnthropicProvider(ProviderClient): # so both the text and the signature_delta tail are collected (in block order). thinking_accum: dict[int, dict[str, Any]] = {} stop_reason = None + usage: Optional[TokenUsage] = None last_message_delta: Any = None for event in events: kind = getattr(event, "type", None) - if kind == "content_block_start": + if kind == "message_start": + # Prompt-side counts (input + cache split) ride the opening event. + usage = ( + _usage_from(getattr(getattr(event, "message", None), "usage", None)) + or usage + ) + elif kind == "content_block_start": block = getattr(event, "content_block", None) block_kind = getattr(block, "type", None) if block_kind == "tool_use": @@ -561,6 +609,13 @@ class AnthropicProvider(ProviderClient): reason = getattr(last_message_delta, "stop_reason", None) if reason: stop_reason = reason + # Final (cumulative) output-token count rides message_delta.usage. + out = int( + getattr(getattr(event, "usage", None), "output_tokens", 0) or 0 + ) + if out: + usage = usage or TokenUsage() + usage.output = out _raise_on_refusal(stop_reason, last_message_delta) tool_calls = [] @@ -580,5 +635,6 @@ class AnthropicProvider(ProviderClient): finish_reason=_STOP_REASON_MAP.get(stop_reason, stop_reason), reasoning=_reasoning_text(thinking_blocks), extras=_thinking_extras(thinking_blocks), + usage=usage, ) ) diff --git a/coworker/providers/base.py b/coworker/providers/base.py index a5f263f6..412ab313 100644 --- a/coworker/providers/base.py +++ b/coworker/providers/base.py @@ -21,6 +21,35 @@ class ToolCall: arguments: dict[str, Any] = field(default_factory=dict) +@dataclass +class TokenUsage: + """Normalized token counts for one model round-trip. + + `input` counts only fresh (uncached) prompt tokens; cached prompt tokens are + split into `cache_read`/`cache_write`. Providers that don't report a cache + split (Ollama, most compat vendors) leave the cache fields at 0. `output` + includes thinking tokens where the vendor bills them as output (Gemini). + """ + + input: int = 0 + output: int = 0 + cache_read: int = 0 + cache_write: int = 0 + + @property + def context_tokens(self) -> int: + """Prompt-side total — what actually occupied the context window.""" + return self.input + self.cache_read + self.cache_write + + def as_dict(self) -> dict[str, int]: + return { + "input": self.input, + "output": self.output, + "cache_read": self.cache_read, + "cache_write": self.cache_write, + } + + @dataclass class AssistantTurn: """One assistant response: free text and/or a set of tool calls.""" @@ -38,6 +67,9 @@ class AssistantTurn: # owning provider consumes its own key when converting history; every other # provider must strip or ignore foreign underscore keys before its wire call. extras: dict[str, Any] = field(default_factory=dict) + # Token counts for this round-trip, normalized across providers. None when the + # backend didn't report usage (some compat servers) — never guessed. + usage: Optional[TokenUsage] = None @property def has_tool_calls(self) -> bool: diff --git a/coworker/providers/bedrock_provider.py b/coworker/providers/bedrock_provider.py index a1c76010..b24fb6f9 100644 --- a/coworker/providers/bedrock_provider.py +++ b/coworker/providers/bedrock_provider.py @@ -43,10 +43,23 @@ from .base import ( ModelCapabilities, ProviderClient, StreamChunk, + TokenUsage, ToolCall, ) from .capabilities import capabilities_for + +def _usage_from(usage: Any) -> Optional[TokenUsage]: + """Converse `usage` dict → normalized counts (`inputTokens` excludes cache).""" + if not isinstance(usage, dict): + return None + return TokenUsage( + input=int(usage.get("inputTokens") or 0), + output=int(usage.get("outputTokens") or 0), + cache_read=int(usage.get("cacheReadInputTokens") or 0), + cache_write=int(usage.get("cacheWriteInputTokens") or 0), + ) + # Converse has no required max token param but per-model defaults vary wildly (Meta's is # 512 — an agent turn gets truncated mid-tool-call); 4096 fits every family's ceiling. DEFAULT_MAX_TOKENS = 4096 @@ -381,6 +394,7 @@ class _BedrockConverseClient(ProviderClient): finish_reason=_STOP_REASON_MAP.get(stop_reason, stop_reason), raw=response, reasoning="".join(reasoning_parts) or None, + usage=_usage_from(response.get("usage")), ) def stream( @@ -400,6 +414,7 @@ class _BedrockConverseClient(ProviderClient): reasoning_parts: list[str] = [] tool_accum: dict[int, dict[str, str]] = {} stop_reason = None + usage: Optional[TokenUsage] = None for event in response.get("stream") or []: if "contentBlockStart" in event: @@ -427,6 +442,8 @@ class _BedrockConverseClient(ProviderClient): yield StreamChunk(reasoning_delta=thought) elif "messageStop" in event: stop_reason = event["messageStop"].get("stopReason") or stop_reason + elif "metadata" in event: + usage = _usage_from(event["metadata"].get("usage")) or usage tool_calls = [ ToolCall( @@ -442,6 +459,7 @@ class _BedrockConverseClient(ProviderClient): tool_calls=tool_calls, finish_reason=_STOP_REASON_MAP.get(stop_reason, stop_reason), reasoning="".join(reasoning_parts) or None, + usage=usage, ) ) diff --git a/coworker/providers/gemini_provider.py b/coworker/providers/gemini_provider.py index 93ef234a..9e58bbe6 100644 --- a/coworker/providers/gemini_provider.py +++ b/coworker/providers/gemini_provider.py @@ -30,10 +30,26 @@ from .base import ( ModelCapabilities, ProviderClient, StreamChunk, + TokenUsage, ToolCall, ) from .capabilities import capabilities_for + +def _usage_from(meta: Any) -> Optional[TokenUsage]: + """`usage_metadata` → normalized counts. `prompt_token_count` INCLUDES the cached + share; thinking tokens are billed as output, so they fold into `output`.""" + if meta is None: + return None + prompt = int(getattr(meta, "prompt_token_count", 0) or 0) + cached = int(getattr(meta, "cached_content_token_count", 0) or 0) + return TokenUsage( + input=max(prompt - cached, 0), + output=int(getattr(meta, "candidates_token_count", 0) or 0) + + int(getattr(meta, "thoughts_token_count", 0) or 0), + cache_read=cached, + ) + # Gemini finishReason → the engine's OpenAI-shaped finish_reason vocabulary. STOP maps to # "tool_calls" instead when the turn contains function calls (Gemini has no distinct reason). _FINISH_REASON_MAP = { @@ -467,6 +483,7 @@ class GeminiProvider(ProviderClient): raw=response, reasoning="".join(parsed.thoughts) or None, extras=_signature_extras(parsed.text_sig, parsed.call_sigs), + usage=_usage_from(getattr(response, "usage_metadata", None)), ) def capabilities(self, model: str) -> ModelCapabilities: @@ -491,10 +508,15 @@ class GeminiProvider(ProviderClient): finish = None text_sig: Optional[str] = None call_sigs: list[Optional[str]] = [] + usage: Optional[TokenUsage] = None # Unlike Anthropic, function_call parts arrive whole (args are a complete dict per # part), so there is no JSON accumulation — just collect parts across chunks. for chunk in client.models.generate_content_stream(**kwargs): + # Counts are cumulative per chunk; the last one seen is the final total. + chunk_usage = _usage_from(getattr(chunk, "usage_metadata", None)) + if chunk_usage is not None: + usage = chunk_usage parsed = _parse_candidate(chunk) for thought in parsed.thoughts: thought_parts.append(thought) @@ -520,5 +542,6 @@ class GeminiProvider(ProviderClient): finish_reason=_map_finish(finish, bool(tool_calls)), reasoning="".join(thought_parts) or None, extras=_signature_extras(text_sig, call_sigs), + usage=usage, ) ) diff --git a/coworker/providers/matrix.py b/coworker/providers/matrix.py index f4e7d004..9c3e0d79 100644 --- a/coworker/providers/matrix.py +++ b/coworker/providers/matrix.py @@ -11,6 +11,11 @@ falls back to the conservative heuristics in ``capabilities.py`` at their own ri degraded results. Ids verified against vendor/reseller catalogs on 2026-07-04; refresh the reseller rows when catalogs rotate (they rename on every model generation). +Context windows (``context_window``, tokens) feed the GUI's context-fill meter. Entries +where the vendor spec wasn't re-checked stay ``None`` — the meter simply hides rather than +showing a made-up denominator. Values entered 2026-07-28 from vendor docs; verify alongside +the id refresh. + Resellers: Together + Fireworks + OpenRouter. TODO: add Groq entries here AND its descriptor in ``registry.py`` once the current provider surface is tested — deliberately deferred to bound how much needs verifying at once. @@ -19,6 +24,7 @@ deferred to bound how much needs verifying at once. from __future__ import annotations from dataclasses import dataclass +from typing import Optional from .base import ModelCapabilities @@ -37,6 +43,9 @@ _AGENTIC_VISION = ModelCapabilities( class ModelEntry: label: str # UI display name, e.g. "GLM-5.2 · via Together" caps: ModelCapabilities = _AGENTIC + # Max context length in tokens (prompt side), for the GUI's context-fill meter. + # None = not verified against the vendor spec yet; the meter hides. + context_window: Optional[int] = None MATRIX: dict[str, ModelEntry] = { @@ -44,32 +53,38 @@ MATRIX: dict[str, ModelEntry] = { # GPT-5.6 (2026-07-09): number = generation, Sol/Terra/Luna = capability tiers. # Bare "gpt-5.6" aliases to Sol server-side; we list the explicit tier ids only. # Rolling out — accounts without access get a friendly error (providers/errors.py). - "gpt-5.6-sol": ModelEntry("GPT-5.6 Sol · OpenAI", _AGENTIC_VISION), - "gpt-5.6-terra": ModelEntry("GPT-5.6 Terra · OpenAI", _AGENTIC_VISION), - "gpt-5.6-luna": ModelEntry("GPT-5.6 Luna · OpenAI", _AGENTIC_VISION), - "gpt-5.5": ModelEntry("GPT-5.5 · OpenAI", _AGENTIC_VISION), + "gpt-5.6-sol": ModelEntry("GPT-5.6 Sol · OpenAI", _AGENTIC_VISION, 400_000), + "gpt-5.6-terra": ModelEntry("GPT-5.6 Terra · OpenAI", _AGENTIC_VISION, 400_000), + "gpt-5.6-luna": ModelEntry("GPT-5.6 Luna · OpenAI", _AGENTIC_VISION, 400_000), + "gpt-5.5": ModelEntry("GPT-5.5 · OpenAI", _AGENTIC_VISION, 400_000), # Fable 5 (2026-06-09) is GA; its Mythos 5 sibling is approved-orgs-only, so it # stays out of a picker meant for the public. "anthropic:claude-fable-5": ModelEntry( - "Claude Fable 5 · Anthropic", _AGENTIC_VISION + "Claude Fable 5 · Anthropic", _AGENTIC_VISION, 1_000_000 ), "anthropic:claude-opus-4-8": ModelEntry( - "Claude Opus 4.8 · Anthropic", _AGENTIC_VISION + "Claude Opus 4.8 · Anthropic", _AGENTIC_VISION, 200_000 ), "anthropic:claude-sonnet-4-6": ModelEntry( - "Claude Sonnet 4.6 · Anthropic", _AGENTIC_VISION + "Claude Sonnet 4.6 · Anthropic", _AGENTIC_VISION, 200_000 ), "anthropic:claude-haiku-4-5": ModelEntry( - "Claude Haiku 4.5 · Anthropic", _AGENTIC_VISION + "Claude Haiku 4.5 · Anthropic", _AGENTIC_VISION, 200_000 ), # Gemini 3 (thought signatures required in tool loops — carried via the `_gemini` # message sidecar, see gemini_provider.py; ids from the vendor catalog 2026-07-22). "gemini:gemini-3.1-pro-preview": ModelEntry( - "Gemini 3.1 Pro · Google", _AGENTIC_VISION + "Gemini 3.1 Pro · Google", _AGENTIC_VISION, 1_048_576 + ), + "gemini:gemini-3.6-flash": ModelEntry( + "Gemini 3.6 Flash · Google", _AGENTIC_VISION, 1_048_576 + ), + "gemini:gemini-2.5-pro": ModelEntry( + "Gemini 2.5 Pro · Google", _AGENTIC_VISION, 1_048_576 + ), + "gemini:gemini-2.5-flash": ModelEntry( + "Gemini 2.5 Flash · Google", _AGENTIC_VISION, 1_048_576 ), - "gemini:gemini-3.6-flash": ModelEntry("Gemini 3.6 Flash · Google", _AGENTIC_VISION), - "gemini:gemini-2.5-pro": ModelEntry("Gemini 2.5 Pro · Google", _AGENTIC_VISION), - "gemini:gemini-2.5-flash": ModelEntry("Gemini 2.5 Flash · Google", _AGENTIC_VISION), # -- direct OpenAI-compatible vendors ---------------------------------------- # Muse Spark (Meta Model API, public preview 2026-07-09): multimodal + tools via # their OpenAI-compat surface. Vision yes; PDFs unverified over compat — falls @@ -80,62 +95,78 @@ MATRIX: dict[str, ModelEntry] = { tools=True, vision=True, parallel_tool_calls=True, streaming=True ), ), - "zai:glm-5.2": ModelEntry("GLM-5.2 · Z AI"), - "deepseek:deepseek-v4-flash": ModelEntry("DeepSeek V4 Flash · DeepSeek"), - "deepseek:deepseek-v4-pro": ModelEntry("DeepSeek V4 Pro · DeepSeek"), - "kimi:kimi-k2.6": ModelEntry("Kimi K2.6 · Moonshot"), + "zai:glm-5.2": ModelEntry("GLM-5.2 · Z AI", _AGENTIC, 128_000), + "deepseek:deepseek-v4-flash": ModelEntry( + "DeepSeek V4 Flash · DeepSeek", _AGENTIC, 128_000 + ), + "deepseek:deepseek-v4-pro": ModelEntry( + "DeepSeek V4 Pro · DeepSeek", _AGENTIC, 128_000 + ), + "kimi:kimi-k2.6": ModelEntry("Kimi K2.6 · Moonshot", _AGENTIC, 256_000), "minimax:MiniMax-M2.5": ModelEntry("MiniMax M2.5 · MiniMax"), - "qwen:qwen3-max": ModelEntry("Qwen3 Max · Alibaba"), - "xai:grok-4.3": ModelEntry("Grok 4.3 · xAI"), - "mistral:mistral-large-latest": ModelEntry("Mistral Large · Mistral"), + "qwen:qwen3-max": ModelEntry("Qwen3 Max · Alibaba", _AGENTIC, 256_000), + "xai:grok-4.3": ModelEntry("Grok 4.3 · xAI", _AGENTIC, 256_000), + "mistral:mistral-large-latest": ModelEntry( + "Mistral Large · Mistral", _AGENTIC, 128_000 + ), # -- resellers (their model namespaces, verbatim) ----------------------------- "together:thinkingmachines/Inkling": ModelEntry("Inkling · via Together"), - "together:zai-org/GLM-5.2": ModelEntry("GLM-5.2 · via Together"), + "together:zai-org/GLM-5.2": ModelEntry("GLM-5.2 · via Together", _AGENTIC, 128_000), # Kimi K3 (2026-07-16) is not on Together yet — weights land ~07-27; revisit then. - "together:moonshotai/Kimi-K2.7-Code": ModelEntry("Kimi K2.7 Code · via Together"), - "together:moonshotai/Kimi-K2.6": ModelEntry("Kimi K2.6 · via Together"), + "together:moonshotai/Kimi-K2.7-Code": ModelEntry( + "Kimi K2.7 Code · via Together", _AGENTIC, 256_000 + ), + "together:moonshotai/Kimi-K2.6": ModelEntry( + "Kimi K2.6 · via Together", _AGENTIC, 256_000 + ), "together:deepseek-ai/DeepSeek-V4-Pro": ModelEntry( - "DeepSeek V4 Pro · via Together" + "DeepSeek V4 Pro · via Together", _AGENTIC, 128_000 ), "together:meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": ModelEntry( - "Llama 4 Maverick · via Together" + "Llama 4 Maverick · via Together", _AGENTIC, 1_000_000 ), "fireworks:accounts/fireworks/models/glm-5p2": ModelEntry( - "GLM-5.2 · via Fireworks" + "GLM-5.2 · via Fireworks", _AGENTIC, 128_000 ), "fireworks:accounts/fireworks/models/kimi-k2p6": ModelEntry( - "Kimi K2.6 · via Fireworks" + "Kimi K2.6 · via Fireworks", _AGENTIC, 256_000 ), "fireworks:accounts/fireworks/models/deepseek-v4-pro": ModelEntry( - "DeepSeek V4 Pro · via Fireworks" + "DeepSeek V4 Pro · via Fireworks", _AGENTIC, 128_000 ), "fireworks:accounts/fireworks/models/llama4-maverick-instruct-basic": ModelEntry( - "Llama 4 Maverick · via Fireworks" + "Llama 4 Maverick · via Fireworks", _AGENTIC, 1_000_000 ), # OpenRouter slugs are lowercase `/` (checked against their catalog # 2026-07-25); same labs as above, one key for all of them. - "openrouter:z-ai/glm-5.2": ModelEntry("GLM-5.2 · via OpenRouter"), - "openrouter:moonshotai/kimi-k2.6": ModelEntry("Kimi K2.6 · via OpenRouter"), - "openrouter:deepseek/deepseek-v4-pro": ModelEntry("DeepSeek V4 Pro · via OpenRouter"), + "openrouter:z-ai/glm-5.2": ModelEntry("GLM-5.2 · via OpenRouter", _AGENTIC, 128_000), + "openrouter:moonshotai/kimi-k2.6": ModelEntry( + "Kimi K2.6 · via OpenRouter", _AGENTIC, 256_000 + ), + "openrouter:deepseek/deepseek-v4-pro": ModelEntry( + "DeepSeek V4 Pro · via OpenRouter", _AGENTIC, 128_000 + ), "openrouter:meta-llama/llama-4-maverick": ModelEntry( - "Llama 4 Maverick · via OpenRouter" + "Llama 4 Maverick · via OpenRouter", _AGENTIC, 1_000_000 ), # -- cloud accounts (models running in the user's own AWS/GCP) ---------------- # Bedrock ids carry a family segment (claude/ → native Anthropic path, other/ → # Converse) plus AWS's own `-v:` version suffix. Some regions require the # `us.`/`eu.` cross-region inference-profile prefix — custom add-model accepts those. "bedrock:claude/anthropic.claude-sonnet-4-6-v1:0": ModelEntry( - "Claude Sonnet 4.6 · AWS Bedrock", _AGENTIC_VISION + "Claude Sonnet 4.6 · AWS Bedrock", _AGENTIC_VISION, 200_000 ), "bedrock:claude/anthropic.claude-haiku-4-5-v1:0": ModelEntry( - "Claude Haiku 4.5 · AWS Bedrock", _AGENTIC_VISION + "Claude Haiku 4.5 · AWS Bedrock", _AGENTIC_VISION, 200_000 + ), + "bedrock:other/amazon.nova-2-pro-v1:0": ModelEntry( + "Nova 2 Pro · AWS Bedrock", _AGENTIC, 300_000 ), - "bedrock:other/amazon.nova-2-pro-v1:0": ModelEntry("Nova 2 Pro · AWS Bedrock"), "bedrock:other/meta.llama4-maverick-17b-instruct-v1:0": ModelEntry( - "Llama 4 Maverick · AWS Bedrock" + "Llama 4 Maverick · AWS Bedrock", _AGENTIC, 1_000_000 ), "bedrock:other/mistral.mistral-large-3-v1:0": ModelEntry( - "Mistral Large 3 · AWS Bedrock" + "Mistral Large 3 · AWS Bedrock", _AGENTIC, 128_000 ), # Live-verified on Converse 2026-07-26 (complete/stream/tool round trip); asked for # two tool calls it emits them one at a time, so parallel stays off. @@ -148,22 +179,22 @@ MATRIX: dict[str, ModelEntry] = { # Vertex ids carry a family segment too (gemini/ and claude/ → native paths, # openweight/ → the MaaS OpenAI-compat endpoint, keeping the publisher segment). "vertex:gemini/gemini-3.1-pro-preview": ModelEntry( - "Gemini 3.1 Pro · Vertex AI", _AGENTIC_VISION + "Gemini 3.1 Pro · Vertex AI", _AGENTIC_VISION, 1_048_576 ), "vertex:gemini/gemini-3.6-flash": ModelEntry( - "Gemini 3.6 Flash · Vertex AI", _AGENTIC_VISION + "Gemini 3.6 Flash · Vertex AI", _AGENTIC_VISION, 1_048_576 ), "vertex:claude/claude-sonnet-4-6": ModelEntry( - "Claude Sonnet 4.6 · Vertex AI", _AGENTIC_VISION + "Claude Sonnet 4.6 · Vertex AI", _AGENTIC_VISION, 200_000 ), "vertex:claude/claude-haiku-4-5": ModelEntry( - "Claude Haiku 4.5 · Vertex AI", _AGENTIC_VISION + "Claude Haiku 4.5 · Vertex AI", _AGENTIC_VISION, 200_000 ), "vertex:openweight/meta/llama-4-maverick-17b-128e-instruct-maas": ModelEntry( - "Llama 4 Maverick · Vertex AI" + "Llama 4 Maverick · Vertex AI", _AGENTIC, 1_000_000 ), "vertex:openweight/qwen/qwen3-coder-480b-a35b-instruct-maas": ModelEntry( - "Qwen3 Coder · Vertex AI" + "Qwen3 Coder · Vertex AI", _AGENTIC, 256_000 ), } @@ -177,6 +208,13 @@ def model_labels() -> dict[str, str]: return {mid: e.label for mid, e in MATRIX.items()} +def model_context_windows() -> dict[str, int]: + """Full-id → context-window map (verified entries only), for the GUI's fill meter.""" + return { + mid: e.context_window for mid, e in MATRIX.items() if e.context_window + } + + def models_for_provider(provider: str) -> list[str]: """BARE model ids (prefix stripped) the matrix curates for a provider — feeds the Settings pane's suggestions and the composer picker so both stay in lockstep with the diff --git a/coworker/providers/openai_provider.py b/coworker/providers/openai_provider.py index 42edec85..5e480e7c 100644 --- a/coworker/providers/openai_provider.py +++ b/coworker/providers/openai_provider.py @@ -15,6 +15,7 @@ from .base import ( ModelCapabilities, ProviderClient, StreamChunk, + TokenUsage, ToolCall, ) from .capabilities import capabilities_for @@ -91,9 +92,30 @@ def _param_fix_retry(kwargs: dict[str, Any], exc: Exception) -> dict[str, Any]: fixed = dict(kwargs) fixed["max_completion_tokens"] = fixed.pop("max_tokens") return fixed + if "stream_options" in msg and "stream_options" in kwargs: + # Older compat servers don't know the usage opt-in; drop it, lose only metering. + fixed = dict(kwargs) + fixed.pop("stream_options") + return fixed raise exc +def _usage_from(usage: Any) -> Optional[TokenUsage]: + """chat.completions usage → normalized counts. `prompt_tokens` INCLUDES cached + tokens, so the cached share is subtracted into `cache_read`; no write-side split + exists on this API shape.""" + if usage is None: + return None + prompt = int(getattr(usage, "prompt_tokens", 0) or 0) + details = getattr(usage, "prompt_tokens_details", None) + cached = int(getattr(details, "cached_tokens", 0) or 0) + return TokenUsage( + input=max(prompt - cached, 0), + output=int(getattr(usage, "completion_tokens", 0) or 0), + cache_read=cached, + ) + + class OpenAIProvider(ProviderClient): def __init__( self, @@ -174,6 +196,7 @@ class OpenAIProvider(ProviderClient): finish_reason=getattr(choice, "finish_reason", None), raw=response, reasoning=_delta_reasoning(message), + usage=_usage_from(getattr(response, "usage", None)), ) def capabilities(self, model: str) -> ModelCapabilities: @@ -191,6 +214,9 @@ class OpenAIProvider(ProviderClient): "model": model, "messages": _strip_foreign_sidecars(messages), "stream": True, + # Usage on the final chunk (empty `choices`). Compat servers that reject + # the option get a one-shot retry without it (_param_fix_retry). + "stream_options": {"include_usage": True}, **settings, } if tools: @@ -202,6 +228,7 @@ class OpenAIProvider(ProviderClient): reasoning_parts: list[str] = [] tool_accum: dict[int, dict[str, str]] = {} finish_reason = None + usage: Optional[TokenUsage] = None # Up to two param-fix retries: effort and max_tokens can BOTH need fixing. for _ in range(2): @@ -213,6 +240,9 @@ class OpenAIProvider(ProviderClient): else: chunks = client.chat.completions.create(**kwargs) for chunk in chunks: + chunk_usage = _usage_from(getattr(chunk, "usage", None)) + if chunk_usage is not None: + usage = chunk_usage choices = getattr(chunk, "choices", None) if not choices: continue @@ -262,6 +292,7 @@ class OpenAIProvider(ProviderClient): tool_calls=tool_calls, finish_reason=finish_reason, reasoning="".join(reasoning_parts) or None, + usage=usage, ) ) diff --git a/coworker/server/manager.py b/coworker/server/manager.py index 1d042ca0..c838ea39 100644 --- a/coworker/server/manager.py +++ b/coworker/server/manager.py @@ -1769,7 +1769,7 @@ class SessionManager: selectable = [m for m in self._curated_models() if _selectable(m)] if self.model not in selectable: selectable.insert(0, self.model) - from ..providers.matrix import model_labels + from ..providers.matrix import model_context_windows, model_labels return { "provider": "openai", @@ -1778,6 +1778,9 @@ class SessionManager: # Curated-matrix display names ({full id → "GLM-5.2 · via Together"}) so every # picker shows human labels; custom models absent here render their raw id. "model_labels": model_labels(), + # {full id → context window in tokens}, verified matrix entries only — + # drives the composer's context-fill meter (absent id → meter hides). + "model_context_windows": model_context_windows(), "has_key": env_key or stored, # Provider-agnostic "can this default model actually run?" — true when the default # model's provider is configured (any provider, not just OpenAI). Drives the GUI's diff --git a/pyproject.toml b/pyproject.toml index 8eaef6ed..687dad70 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ dependencies = [ "docstring_parser", "pyyaml>=6", # persona manifest frontmatter (YAML) "pydantic>=2", - "mcp>=1.1", # MCP client (stdio + streamable-http); we use our own async layer on it + "mcp>=1.1,<2", # MCP client (stdio + streamable-http); 2.0 removed streamablehttp_client "httpx>=0.27", # sync outbound senders for messaging connectors (send_message tool) "websockets>=13", # managed Slack relay client transport (relay_client.py) "ddgs>=9", # keyless default web-search provider (DuckDuckGo); Tavily/Brave use httpx diff --git a/surfaces/gui/e2e/fixtures.ts b/surfaces/gui/e2e/fixtures.ts index 7df802f5..4224bff7 100644 --- a/surfaces/gui/e2e/fixtures.ts +++ b/surfaces/gui/e2e/fixtures.ts @@ -46,6 +46,11 @@ const SETTINGS = { "anthropic:claude-opus-4-8": "Claude Opus 4.8 · Anthropic", "zai:glm-5.2": "GLM-5.2 · Z AI", }, + // Context windows (subset — mirrors /v1/settings.model_context_windows); drives the + // composer usage chip's context-fill meter. + model_context_windows: { + "anthropic:claude-opus-4-8": 200_000, + }, }; const PERSONAS = { @@ -704,7 +709,18 @@ export async function mockApi(page: import("@playwright/test").Page) { send("assistant_delta", { text: msg.text }); // Echo the model the message carried — pins the model-per-message contract (the // composer's visible model must ride on every user_message; 2026-07-04 fix). - send("assistant_message", { text: `Echo: ${msg.text} [model=${msg.model || "none"}]` }); + // `usage` mirrors the real engine's assistant_message sidecar (OPE-42): fixed + // counts per turn so the usage-chip specs can assert exact accumulation. + send("assistant_message", { + text: `Echo: ${msg.text} [model=${msg.model || "none"}]`, + usage: { + model: msg.model || "anthropic:claude-opus-4-8", + input: 1_000, + output: 200, + cache_read: 8_000, + cache_write: 800, + }, + }); send("turn_done"); } else if (msg.type === "approval") { if (pendingTool === "run_shell") { diff --git a/surfaces/gui/e2e/usage-chip.spec.ts b/surfaces/gui/e2e/usage-chip.spec.ts new file mode 100644 index 00000000..bfb98750 --- /dev/null +++ b/surfaces/gui/e2e/usage-chip.spec.ts @@ -0,0 +1,66 @@ +// Token-usage chip (OPE-42): after a turn reports usage, a quiet meter+count chip appears +// in the composer's bottom row; clicking it opens the per-model breakdown popover with the +// context-window fill. The fake agent attaches fixed usage to every echo turn +// (input 1k / output 200 / cache_read 8k / cache_write 800 — 10k per turn), and the +// settings fixture maps the default model to a 200k context window. +import { expect } from "@playwright/test"; +import { test } from "./fixtures"; + +test("usage chip appears after a turn and opens the breakdown popover", async ({ page }) => { + await page.goto("/"); + await page.getByText("Draft the launch note").first().click(); + + // Fresh session: no usage yet — the chip is hidden entirely. + await expect(page.getByTestId("usage-chip")).toHaveCount(0); + + const box = page.getByPlaceholder(/Ask the coworker/); + await box.fill("hello"); + await box.press("Enter"); + await expect(page.getByText("Echo: hello", { exact: false }).first()).toBeVisible({ + timeout: 10_000, + }); + + // Chip shows the session total (1k + 200 + 8k + 800 = 10k). + const chip = page.getByTestId("usage-chip"); + await expect(chip).toContainText("10k"); + + // Popover: context fill (9.8k prompt-side of 200k = 5%) + per-model breakdown. + await chip.click(); + const pop = page.getByTestId("usage-popover"); + await expect(pop).toBeVisible(); + await expect(pop).toContainText("Context window"); + await expect(pop).toContainText("9.8k of 200k · 5%"); + await expect(pop).toContainText("Session totals"); + await expect(pop).toContainText("Claude Opus 4.8 · Anthropic"); + // Cache split present → the input rows read as components of Total input. + await expect(pop).toContainText("Uncached input"); + await expect(pop).toContainText("Cache reads"); + await expect(pop).toContainText("Cache writes"); + // Total input = fresh 1k + cache_read 8k + cache_write 800 (cumulative billed input). + await expect(pop).toContainText("Total input"); + await expect(pop).toContainText("9.8k"); + await expect(pop).toContainText("10k tokens"); + + // Second turn accumulates (totals double), and the scrim click closes the popover. + await page.mouse.click(10, 10); + await expect(pop).toHaveCount(0); + await box.fill("again"); + await box.press("Enter"); + await expect(page.getByText("Echo: again", { exact: false }).first()).toBeVisible({ + timeout: 10_000, + }); + await expect(chip).toContainText("20k"); +}); + +test("usage resets on a new session", async ({ page }) => { + await page.goto("/"); + await page.getByText("Draft the launch note").first().click(); + const box = page.getByPlaceholder(/Ask the coworker/); + await box.fill("hello"); + await box.press("Enter"); + await expect(page.getByTestId("usage-chip")).toContainText("10k", { timeout: 10_000 }); + + // "＋ New session" wipes the transcript — and the usage accumulation with it. + await page.getByRole("button", { name: /New session/ }).first().click(); + await expect(page.getByTestId("usage-chip")).toHaveCount(0); +}); diff --git a/surfaces/gui/src/App.tsx b/surfaces/gui/src/App.tsx index 1f50216b..ecca5bf7 100644 --- a/surfaces/gui/src/App.tsx +++ b/surfaces/gui/src/App.tsx @@ -31,10 +31,19 @@ import { type SurfaceVisibility, type WorkspaceCommandTrust, } from "./api"; -import type { ApprovalDecision, Attachment, Item, SessionInfo, TodoItem, WsEvent } from "./types"; +import type { + ApprovalDecision, + Attachment, + Item, + SessionInfo, + SessionUsage, + TodoItem, + WsEvent, +} from "./types"; import { isProjectScoped } from "./personaScope"; import { baseName } from "./paths"; import { itemsFromMessages } from "./itemsFromMessages"; +import { addTurnUsage, emptyUsage, usageFromMessages } from "./usage"; import { streamMode } from "./streamGate"; import { InboxItemCard } from "./components/InboxItemCard"; import { isTauri, platformOS, startWindowDrag } from "./tauri"; @@ -156,6 +165,12 @@ export function App() { const [model, setModel] = useState("gpt-5.6-sol"); const [models, setModels] = useState([]); const [modelLabels, setModelLabels] = useState>({}); + // {full model id → context window in tokens} from the curated matrix (verified only); + // drives the composer usage chip's context-fill meter. + const [modelContextWindows, setModelContextWindows] = useState>({}); + // Per-session token usage (OPE-42): rebuilt from the transcript on session load, + // accumulated live from assistant_message events, reset with the transcript. + const [usage, setUsage] = useState(emptyUsage()); const [surfaces, setSurfaces] = useState({ cowork: true, chat: false, code: false }); const [mode, setMode] = useState("interactive"); const [connected, setConnected] = useState(false); @@ -395,9 +410,12 @@ export function App() { setBranch(null); } try { - setItems(itemsFromMessages(await getSessionMessages(last.session_id))); + const messages = await getSessionMessages(last.session_id); + setItems(itemsFromMessages(messages)); + setUsage(usageFromMessages(messages)); } catch { setItems([]); + setUsage(emptyUsage()); } setSessionId(last.session_id); setShowGate(false); @@ -486,6 +504,7 @@ export function App() { .then((s) => { setModels(s.models || []); setModelLabels(s.model_labels || {}); + setModelContextWindows(s.model_context_windows || {}); setModelReady(s.model_ready); if (s.surfaces) setSurfaces(s.surfaces); }) @@ -601,6 +620,7 @@ export function App() { setReasoningStream(reasoningRef.current + (d.text || "")); break; case "assistant_message": { + if (d.usage) setUsage((u) => addTurnUsage(u, d.usage)); // The event's reasoning is authoritative (covers background-delivered turns); // the local buffer is the fallback for older servers. const reasoning = d.reasoning || reasoningRef.current; @@ -899,6 +919,7 @@ export function App() { const target = forAgent || agent; setSurface("session"); // return to the conversation view if we were on a sub-view setItems([]); + setUsage(emptyUsage()); setStreaming(""); setTodo([]); setRunning(false); @@ -975,8 +996,10 @@ export function App() { try { const messages = await getSessionMessages(id); setItems(itemsFromMessages(messages)); + setUsage(usageFromMessages(messages)); } catch { setItems([]); + setUsage(emptyUsage()); } }; const switchAgent = async (name: string) => { @@ -989,6 +1012,7 @@ export function App() { setAgent(name); setItems([]); + setUsage(emptyUsage()); setStreaming(""); setTodo([]); setRunning(false); @@ -1017,9 +1041,12 @@ export function App() { else setShowGate(true); setSessionId(target.sessionId); try { - setItems(itemsFromMessages(await getSessionMessages(target.sessionId))); + const messages = await getSessionMessages(target.sessionId); + setItems(itemsFromMessages(messages)); + setUsage(usageFromMessages(messages)); } catch { setItems([]); + setUsage(emptyUsage()); } return; } @@ -1043,6 +1070,7 @@ export function App() { setShowGate(false); setGateCreate(false); setItems([]); + setUsage(emptyUsage()); setStreaming(""); setTodo([]); setSessionId(newId()); @@ -1055,6 +1083,7 @@ export function App() { const target = forAgent || agent; setSurface("session"); setItems([]); + setUsage(emptyUsage()); setStreaming(""); setTodo([]); setRunning(false); @@ -1079,6 +1108,7 @@ export function App() { // Archiving the open chat: leave it and start fresh (it moves to the Archived section). if (archived && id === sessionId) { setItems([]); + setUsage(emptyUsage()); setStreaming(""); setTodo([]); setRunning(false); @@ -1091,6 +1121,7 @@ export function App() { refreshSessions(); if (id === sessionId) { setItems([]); + setUsage(emptyUsage()); setStreaming(""); setTodo([]); setRunning(false); @@ -1569,6 +1600,8 @@ export function App() { onUnattendedChange={agent !== "chat" ? toggleUnattended : undefined} prefill={composerPrefill} resetKey={sessionId} + usage={usage} + contextWindow={modelContextWindows[model]} placeholder={ agent === "code" ? "Ask the coder to build, fix, or explain… (drop or paste files)" diff --git a/surfaces/gui/src/api.ts b/surfaces/gui/src/api.ts index 93021221..9b40d3ee 100644 --- a/surfaces/gui/src/api.ts +++ b/surfaces/gui/src/api.ts @@ -141,6 +141,9 @@ export interface ConversationMessage { tool_calls?: any[]; tool_call_id?: string; source?: MessageSource; + // Token counts for the round-trip that produced an assistant message + // ({model, input, output, cache_read, cache_write}); absent on older servers. + usage?: import("./types").TurnUsage; [key: string]: any; } @@ -691,6 +694,9 @@ export interface ModelSettings { sessions_peek?: number; // Curated-matrix display names ({full id → "GLM-5.2 · via Together"}); custom models absent. model_labels?: Record; + // {full id → context window in tokens}, verified matrix entries only — drives the + // composer's context-fill meter (absent id → the meter hides). Optional for older backends. + model_context_windows?: Record; // Token savings (PDF attachments): fallback for models without native PDF support, // and attach-time thresholds. Optional so the GUI is robust to an older backend. pdf_fallback?: "text" | "images"; diff --git a/surfaces/gui/src/components/Composer.tsx b/surfaces/gui/src/components/Composer.tsx index 852e4d6f..5d638ee1 100644 --- a/surfaces/gui/src/components/Composer.tsx +++ b/surfaces/gui/src/components/Composer.tsx @@ -1,7 +1,8 @@ import { useEffect, useLayoutEffect, useRef, useState, type ReactNode } from "react"; -import type { Attachment } from "../types"; +import type { Attachment, SessionUsage } from "../types"; import { isPdfFile, readFile } from "../attach"; import { getSettings, inspectPdf } from "../api"; +import { formatTokens, totalTokens } from "../usage"; import { Dropdown, type Option } from "./Dropdown"; import { Icon } from "./Icon"; import { Toggle } from "./Toggle"; @@ -77,6 +78,12 @@ interface Props { resetKey?: string; // Surface-specific hint shown in the empty textarea. placeholder?: string; + // Per-session token usage (OPE-42) — absent/empty hides the usage chip entirely + // (older servers, backends that don't report usage, fresh sessions). + usage?: SessionUsage; + // Context-window size (tokens) of the ACTIVE model, from the curated matrix; + // undefined hides the fill meter (unverified/custom models) but keeps the counts. + contextWindow?: number; } export function Composer(props: Props) { @@ -462,6 +469,18 @@ export function Composer(props: Props) { + {/* token usage (OPE-42) — a quiet meter+count chip; hidden until the server + reports usage. Fill = context-window occupancy (bounded), count = session + consumption (unbounded, so never a fill). */} + {!dictation?.recording && props.usage && totalTokens(props.usage) > 0 && ( + + )} + {/* model — a quiet chip, now for the session's whole life (§17 rev 2026-07-22: mid-session switching shipped, so the picker stays actionable; the topbar subtitle still states the current model). */} @@ -546,6 +565,135 @@ export function Composer(props: Props) { ); } +// Token-usage chip + popover (OPE-42). Trigger: a tiny context-fill meter (only when the +// active model's window is known) + the session's total token count. Click → per-model +// breakdown. Tokens only, never dollars (true cost is unknowable client-side — discounted +// pricing, per-provider cache billing). +function UsageChip({ + usage, + contextWindow, + model, + modelLabels, +}: { + usage: SessionUsage; + contextWindow?: number; + model: string; + modelLabels?: Record; +}) { + const [open, setOpen] = useState(false); + const total = totalTokens(usage); + const pct = contextWindow + ? Math.min(100, Math.round((usage.context / contextWindow) * 100)) + : null; + const labelFor = (id: string) => + id === "unknown" ? "Unknown model" : modelLabels?.[id] || shortModel(id); + // One field per line, session-summed (owner ask 2026-07-28). Values are cumulative + // across the whole session, never just the last turn; "Input" is the fresh + // (uncached) share — the cached share sits in the cache rows at its own price. + const stat = (label: string, value: number) => ( +
+ {label} + {formatTokens(value)} +
+ ); + return ( +
+ + {open && ( + <> +
setOpen(false)} /> +
+ {contextWindow ? ( +
+
+ Context window +
+
+
+
+
+ {formatTokens(usage.context)} of {formatTokens(contextWindow)} · {pct}% +
+
+ ) : usage.context > 0 ? ( +
+ In context now: {formatTokens(usage.context)} tokens +
+ ) : null} +
+ Session totals +
+
+ {Object.entries(usage.byModel).map(([id, t]) => ( +
+
+ {labelFor(id)} +
+ {/* Every row is a session sum. With a cache split, the input rows are + the three BILLING CLASSES of input (each priced differently) and + read as components: uncached + cache reads + cache writes = total. + Without one (Ollama, compat vendors), plain "Input" says it all. */} +
+ {t.cache_read + t.cache_write > 0 ? ( + <> + {stat("Uncached input", t.input)} + {stat("Cache reads", t.cache_read)} + {stat("Cache writes", t.cache_write)} + {stat("Total input", t.input + t.cache_read + t.cache_write)} + + ) : ( + stat("Input", t.input) + )} + {stat("Output", t.output)} +
+
+ ))} +
+
+ Total + {formatTokens(total)} tokens +
+ {model && !modelLabels?.[model] && contextWindow === undefined && ( +
+ Context meter unavailable for custom models. +
+ )} +
+ + )} +
+ ); +} + // The composer's Mode menu (§22): a quiet "Mode ⌄" chip opening the five permission options with // the current one marked, plus — when the session supports it — the "Send approvals to Inbox" // toggle at the bottom (the old standalone InboxControl, folded in). diff --git a/surfaces/gui/src/types.ts b/surfaces/gui/src/types.ts index 04881495..cb8c23fd 100644 --- a/surfaces/gui/src/types.ts +++ b/surfaces/gui/src/types.ts @@ -39,6 +39,25 @@ export interface TodoItem { status: "pending" | "in_progress" | "done"; } +// Per-round-trip token counts, as attached by the server to assistant messages and +// the assistant_message event (`{model, input, output, cache_read, cache_write}`). +// Absent on older servers and on backends that don't report usage. +export interface TurnUsage { + model?: string | null; + input: number; + output: number; + cache_read: number; + cache_write: number; +} + +// Per-session accumulation, keyed by model id (multiple models when the user +// switched mid-session). `context` = the latest round-trip's prompt-side total — +// what currently occupies the active model's context window. +export interface SessionUsage { + byModel: Record; + context: number; +} + export interface SessionInfo { session_id: string; title?: string; diff --git a/surfaces/gui/src/usage.test.ts b/surfaces/gui/src/usage.test.ts new file mode 100644 index 00000000..82e85322 --- /dev/null +++ b/surfaces/gui/src/usage.test.ts @@ -0,0 +1,78 @@ +import { describe, expect, it } from "vitest"; +import { addTurnUsage, emptyUsage, formatTokens, totalTokens, usageFromMessages } from "./usage"; + +const turn = (over: Record = {}) => ({ + model: "anthropic:claude-fable-5", + input: 100, + output: 20, + cache_read: 50, + cache_write: 10, + ...over, +}); + +describe("addTurnUsage", () => { + it("accumulates per model and tracks the latest prompt-side total as context", () => { + let u = addTurnUsage(emptyUsage(), turn()); + u = addTurnUsage(u, turn({ input: 40, output: 5, cache_read: 200, cache_write: 0 })); + const m = u.byModel["anthropic:claude-fable-5"]; + expect(m).toEqual({ + model: "anthropic:claude-fable-5", + input: 140, + output: 25, + cache_read: 250, + cache_write: 10, + }); + // context = LAST turn's input + cache_read + cache_write, not a sum. + expect(u.context).toBe(240); + }); + + it("keys separate models separately (mid-session switch)", () => { + let u = addTurnUsage(emptyUsage(), turn()); + u = addTurnUsage(u, turn({ model: "gpt-5.5", input: 7 })); + expect(Object.keys(u.byModel).sort()).toEqual(["anthropic:claude-fable-5", "gpt-5.5"]); + }); + + it("ignores malformed payloads and clamps negatives", () => { + expect(addTurnUsage(emptyUsage(), null)).toEqual(emptyUsage()); + expect(addTurnUsage(emptyUsage(), "x")).toEqual(emptyUsage()); + const u = addTurnUsage(emptyUsage(), turn({ input: -5, output: "9" })); + expect(u.byModel["anthropic:claude-fable-5"].input).toBe(0); + expect(u.byModel["anthropic:claude-fable-5"].output).toBe(9); + }); + + it("buckets a missing model id under 'unknown'", () => { + const u = addTurnUsage(emptyUsage(), turn({ model: undefined })); + expect(u.byModel["unknown"].input).toBe(100); + }); +}); + +describe("usageFromMessages", () => { + it("folds assistant usage sidecars and ignores everything else", () => { + const u = usageFromMessages([ + { role: "user", content: "hi" }, + { role: "assistant", content: "a", usage: turn() as any }, + { role: "tool", content: "r", usage: turn() as any }, // wrong role — ignored + { role: "assistant", content: "b" }, // no sidecar (older server) — ignored + { role: "assistant", content: "c", usage: turn({ input: 300 }) as any }, + ]); + expect(u.byModel["anthropic:claude-fable-5"].input).toBe(400); + expect(u.context).toBe(360); + }); +}); + +describe("totalTokens / formatTokens", () => { + it("totals across models and directions", () => { + let u = addTurnUsage(emptyUsage(), turn()); + u = addTurnUsage(u, turn({ model: "gpt-5.5" })); + expect(totalTokens(u)).toBe(360); + }); + + it("formats humane numbers", () => { + expect(formatTokens(0)).toBe("0"); + expect(formatTokens(980)).toBe("980"); + expect(formatTokens(12_400)).toBe("12.4k"); + expect(formatTokens(982_000)).toBe("982k"); + expect(formatTokens(1_240_000)).toBe("1.24M"); + expect(formatTokens(NaN)).toBe("0"); + }); +}); diff --git a/surfaces/gui/src/usage.ts b/surfaces/gui/src/usage.ts new file mode 100644 index 00000000..135a3fa1 --- /dev/null +++ b/surfaces/gui/src/usage.ts @@ -0,0 +1,75 @@ +// Per-session token-usage accumulation (OPE-42). Pure functions so the reducer is +// unit-testable without the app. The server attaches a `usage` sidecar +// ({model, input, output, cache_read, cache_write}) to assistant messages and to the +// assistant_message event; older servers and non-reporting backends send none, and +// everything here no-ops gracefully in that case. + +import type { ConversationMessage } from "./api"; +import type { SessionUsage, TurnUsage } from "./types"; + +export function emptyUsage(): SessionUsage { + return { byModel: {}, context: 0 }; +} + +const num = (v: any): number => { + const n = Number(v); + return Number.isFinite(n) && n > 0 ? n : 0; +}; + +/** Fold one turn's usage sidecar into the session accumulation. */ +export function addTurnUsage(prev: SessionUsage, raw: any): SessionUsage { + if (!raw || typeof raw !== "object") return prev; + const turn: TurnUsage = { + model: typeof raw.model === "string" && raw.model ? raw.model : null, + input: num(raw.input), + output: num(raw.output), + cache_read: num(raw.cache_read), + cache_write: num(raw.cache_write), + }; + const key = turn.model || "unknown"; + const cur = prev.byModel[key]; + return { + byModel: { + ...prev.byModel, + [key]: { + model: turn.model, + input: (cur?.input || 0) + turn.input, + output: (cur?.output || 0) + turn.output, + cache_read: (cur?.cache_read || 0) + turn.cache_read, + cache_write: (cur?.cache_write || 0) + turn.cache_write, + }, + }, + // Prompt-side total of the LATEST round-trip = what currently sits in the + // context window (not a sum — each request resends the whole history). + context: turn.input + turn.cache_read + turn.cache_write, + }; +} + +/** Rebuild the accumulation from a replayed transcript (session load/switch). */ +export function usageFromMessages(messages: ConversationMessage[]): SessionUsage { + let acc = emptyUsage(); + for (const m of messages || []) { + if (m.role === "assistant" && m.usage) acc = addTurnUsage(acc, m.usage); + } + return acc; +} + +/** All tokens consumed this session, across models and directions (chip headline). */ +export function totalTokens(u: SessionUsage): number { + return Object.values(u.byModel).reduce( + (sum, t) => sum + t.input + t.output + t.cache_read + t.cache_write, + 0, + ); +} + +/** 980 → "980", 12_400 → "12.4k", 982_000 → "982k", 1_240_000 → "1.24M". */ +export function formatTokens(n: number): string { + if (!Number.isFinite(n) || n <= 0) return "0"; + if (n < 1000) return String(Math.round(n)); + if (n < 1_000_000) { + const k = n / 1000; + return (k < 100 ? k.toFixed(1).replace(/\.0$/, "") : String(Math.round(k))) + "k"; + } + const m = n / 1_000_000; + return (m < 100 ? m.toFixed(2).replace(/\.?0+$/, "") : String(Math.round(m))) + "M"; +} diff --git a/tests/test_anthropic_caching.py b/tests/test_anthropic_caching.py new file mode 100644 index 00000000..f4aea26f --- /dev/null +++ b/tests/test_anthropic_caching.py @@ -0,0 +1,87 @@ +"""Prompt-caching breakpoints on the Anthropic provider (OPE-42 follow-up). + +The provider opts every request into 5-minute ephemeral caching: one breakpoint on +the last system block (tools + system) and one on the final message's last content +block (conversation prefix). Outbound-only — persisted history stays clean. +""" + +from __future__ import annotations + +from coworker.providers.anthropic_provider import AnthropicProvider + +MARKER = {"type": "ephemeral"} + + +def _kwargs(messages, tools=None): + return AnthropicProvider(client=object())._request_kwargs( + model="claude-haiku-4-5", messages=messages, tools=tools, settings={} + ) + + +def test_system_becomes_cached_block_list(): + kwargs = _kwargs( + [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"}, + ] + ) + assert kwargs["system"] == [ + {"type": "text", "text": "be terse", "cache_control": MARKER} + ] + + +def test_last_message_last_block_carries_breakpoint(): + kwargs = _kwargs( + [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "reply"}, + {"role": "user", "content": "second"}, + ] + ) + messages = kwargs["messages"] + assert messages[-1]["content"][-1]["cache_control"] == MARKER + # Only the FINAL message is marked — earlier turns stay unmarked so the + # prefix bytes match the previous request's cache. + for message in messages[:-1]: + assert all("cache_control" not in b for b in message["content"]) + + +def test_tool_result_last_block_carries_breakpoint(): + kwargs = _kwargs( + [ + {"role": "user", "content": "run it"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "t1", + "type": "function", + "function": {"name": "ls", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "t1", "content": "README.md"}, + ] + ) + last = kwargs["messages"][-1]["content"][-1] + assert last["type"] == "tool_result" + assert last["cache_control"] == MARKER + + +def test_persisted_history_is_never_mutated(): + history = [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"}, + ] + _kwargs(history) + assert history == [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"}, + ] + + +def test_no_system_prompt_is_fine(): + kwargs = _kwargs([{"role": "user", "content": "hi"}]) + assert "system" not in kwargs + assert kwargs["messages"][-1]["content"][-1]["cache_control"] == MARKER diff --git a/tests/test_anthropic_provider.py b/tests/test_anthropic_provider.py index 4f234e36..8f5b6065 100644 --- a/tests/test_anthropic_provider.py +++ b/tests/test_anthropic_provider.py @@ -263,7 +263,11 @@ def test_complete_text_turn_with_defaults(): and not turn.has_tool_calls ) assert fake.kwargs["model"] == "claude-sonnet-4-6" - assert fake.kwargs["system"] == "sys" + # System rides as a block list so the last block can carry the prompt-cache + # breakpoint (see _add_cache_breakpoints). + assert fake.kwargs["system"] == [ + {"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}} + ] assert fake.kwargs["max_tokens"] == DEFAULT_MAX_TOKENS # required param, injected assert "tools" not in fake.kwargs diff --git a/tests/test_token_usage.py b/tests/test_token_usage.py new file mode 100644 index 00000000..8848a330 --- /dev/null +++ b/tests/test_token_usage.py @@ -0,0 +1,356 @@ +"""Token-usage metering — provider capture, normalization, engine plumbing. + +Fakes follow the provider test convention: SimpleNamespace objects mimicking each +SDK's response surface, dict events for Bedrock's Converse stream. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import aisuite as ai +from coworker.engine import TurnEngine +from coworker.events import EventType +from coworker.permissions import PermissionEngine +from coworker.providers import ( + AssistantTurn, + ModelCapabilities, + ProviderClient, +) +from coworker.providers.anthropic_provider import AnthropicProvider +from coworker.providers.base import TokenUsage +from coworker.providers.bedrock_provider import _BedrockConverseClient +from coworker.providers.gemini_provider import GeminiProvider +from coworker.providers.matrix import model_context_windows +from coworker.providers.openai_provider import OpenAIProvider +from coworker.tools import ToolRegistry + + +def _final_turn(chunks): + return chunks[-1].turn + + +# -- TokenUsage --------------------------------------------------------------------- + + +def test_context_tokens_is_prompt_side_total(): + usage = TokenUsage(input=100, output=50, cache_read=300, cache_write=20) + assert usage.context_tokens == 420 + assert usage.as_dict() == { + "input": 100, + "output": 50, + "cache_read": 300, + "cache_write": 20, + } + + +# -- Anthropic ---------------------------------------------------------------------- + + +class _FakeAnthropicClient: + def __init__(self, events): + def create(**kwargs): + self.kwargs = kwargs + return events + + self.messages = SimpleNamespace(create=create) + self.beta = SimpleNamespace(messages=SimpleNamespace(create=create)) + + +def test_anthropic_stream_captures_usage(): + events = [ + SimpleNamespace( + type="message_start", + message=SimpleNamespace( + usage=SimpleNamespace( + input_tokens=7, + output_tokens=1, + cache_read_input_tokens=100, + cache_creation_input_tokens=25, + ) + ), + ), + SimpleNamespace( + type="content_block_start", + index=0, + content_block=SimpleNamespace(type="text"), + ), + SimpleNamespace( + type="content_block_delta", + index=0, + delta=SimpleNamespace(type="text_delta", text="hi"), + ), + SimpleNamespace( + type="message_delta", + delta=SimpleNamespace(stop_reason="end_turn"), + usage=SimpleNamespace(output_tokens=42), + ), + SimpleNamespace(type="message_stop"), + ] + provider = AnthropicProvider(client=_FakeAnthropicClient(events)) + turn = _final_turn( + list(provider.stream(model="m", messages=[{"role": "user", "content": "x"}])) + ) + assert turn.usage == TokenUsage(input=7, output=42, cache_read=100, cache_write=25) + + +def test_anthropic_complete_captures_usage(): + response = SimpleNamespace( + content=[SimpleNamespace(type="text", text="hi")], + stop_reason="end_turn", + usage=SimpleNamespace( + input_tokens=10, + output_tokens=5, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + ), + ) + provider = AnthropicProvider(client=_FakeAnthropicClient(response)) + turn = provider.complete(model="m", messages=[{"role": "user", "content": "x"}]) + assert turn.usage == TokenUsage(input=10, output=5) + + +def test_anthropic_stream_without_usage_leaves_none(): + events = [ + SimpleNamespace(type="message_start"), # no message/usage attrs + SimpleNamespace( + type="message_delta", delta=SimpleNamespace(stop_reason="end_turn") + ), + ] + provider = AnthropicProvider(client=_FakeAnthropicClient(events)) + turn = _final_turn( + list(provider.stream(model="m", messages=[{"role": "user", "content": "x"}])) + ) + assert turn.usage is None + + +# -- OpenAI-compat ------------------------------------------------------------------ + + +class _FakeOpenAIClient: + def __init__(self, chunks, *, reject_stream_options=False): + self.calls = [] + + def create(**kwargs): + self.calls.append(kwargs) + if reject_stream_options and "stream_options" in kwargs: + raise RuntimeError("unknown parameter: 'stream_options'") + return chunks + + self.chat = SimpleNamespace(completions=SimpleNamespace(create=create)) + + +def _openai_chunks(): + return [ + SimpleNamespace( + choices=[ + SimpleNamespace( + delta=SimpleNamespace(content="hi", tool_calls=None), + finish_reason="stop", + ) + ], + usage=None, + ), + # Usage arrives on a final empty-choices chunk (include_usage contract). + SimpleNamespace( + choices=[], + usage=SimpleNamespace( + prompt_tokens=140, + completion_tokens=9, + prompt_tokens_details=SimpleNamespace(cached_tokens=40), + ), + ), + ] + + +def test_openai_stream_requests_and_captures_usage(): + fake = _FakeOpenAIClient(_openai_chunks()) + provider = OpenAIProvider(client=fake) + turn = _final_turn( + list(provider.stream(model="m", messages=[{"role": "user", "content": "x"}])) + ) + assert fake.calls[0]["stream_options"] == {"include_usage": True} + # Cached share is carved out of prompt_tokens into cache_read. + assert turn.usage == TokenUsage(input=100, output=9, cache_read=40) + + +def test_openai_stream_retries_without_stream_options_when_rejected(): + fake = _FakeOpenAIClient(_openai_chunks(), reject_stream_options=True) + provider = OpenAIProvider(client=fake) + turn = _final_turn( + list(provider.stream(model="m", messages=[{"role": "user", "content": "x"}])) + ) + assert "stream_options" not in fake.calls[-1] + assert turn.text == "hi" # the turn still completes; only metering is lost + + +def test_openai_complete_captures_usage_without_cache_details(): + response = SimpleNamespace( + choices=[ + SimpleNamespace( + message=SimpleNamespace(content="hi", tool_calls=None), + finish_reason="stop", + ) + ], + usage=SimpleNamespace( + prompt_tokens=30, completion_tokens=4, prompt_tokens_details=None + ), + ) + fake = _FakeOpenAIClient(response) + provider = OpenAIProvider(client=fake) + turn = provider.complete(model="m", messages=[{"role": "user", "content": "x"}]) + assert turn.usage == TokenUsage(input=30, output=4) + + +# -- Gemini ------------------------------------------------------------------------- + + +class _FakeGeminiClient: + def __init__(self, responses): + def generate_content_stream(**kwargs): + self.kwargs = kwargs + return iter(responses) + + def generate_content(**kwargs): + self.kwargs = kwargs + return responses[0] + + self.models = SimpleNamespace( + generate_content=generate_content, + generate_content_stream=generate_content_stream, + ) + + +def _gemini_response(text, usage_metadata=None): + return SimpleNamespace( + candidates=[ + SimpleNamespace( + content=SimpleNamespace( + parts=[SimpleNamespace(text=text, function_call=None)] + ), + finish_reason=SimpleNamespace(name="STOP"), + ) + ], + usage_metadata=usage_metadata, + ) + + +def test_gemini_stream_keeps_last_usage_metadata(): + responses = [ + _gemini_response( + "he", + SimpleNamespace( + prompt_token_count=90, + candidates_token_count=1, + cached_content_token_count=50, + thoughts_token_count=0, + ), + ), + _gemini_response( + "y", + # Cumulative — the last chunk carries the final totals. + SimpleNamespace( + prompt_token_count=90, + candidates_token_count=12, + cached_content_token_count=50, + thoughts_token_count=6, + ), + ), + ] + provider = GeminiProvider(client=_FakeGeminiClient(responses)) + turn = _final_turn( + list(provider.stream(model="m", messages=[{"role": "user", "content": "x"}])) + ) + # input = prompt minus cached; thinking tokens fold into output. + assert turn.usage == TokenUsage(input=40, output=18, cache_read=50) + + +# -- Bedrock (Converse) ------------------------------------------------------------- + + +def test_bedrock_converse_stream_captures_metadata_usage(): + fake = SimpleNamespace( + converse_stream=lambda **kwargs: { + "stream": [ + {"contentBlockDelta": {"delta": {"text": "hi"}, "contentBlockIndex": 0}}, + {"messageStop": {"stopReason": "end_turn"}}, + { + "metadata": { + "usage": { + "inputTokens": 11, + "outputTokens": 3, + "cacheReadInputTokens": 8, + "cacheWriteInputTokens": 2, + } + } + }, + ] + } + ) + client = _BedrockConverseClient(client=fake) + turn = _final_turn( + list(client.stream(model="m", messages=[{"role": "user", "content": "x"}])) + ) + assert turn.usage == TokenUsage(input=11, output=3, cache_read=8, cache_write=2) + + +# -- engine plumbing ---------------------------------------------------------------- + + +class _UsageProvider(ProviderClient): + def complete(self, *, model, messages, tools=None, **settings): + return AssistantTurn( + text="done", + finish_reason="stop", + usage=TokenUsage(input=100, output=20, cache_read=5), + ) + + def capabilities(self, model): + return ModelCapabilities() + + +def _run_engine(tmp_path): + registry = ToolRegistry() + registry.register_all(ai.toolkits.files(root=str(tmp_path), allow_write=True)) + engine = TurnEngine( + provider=_UsageProvider(), + registry=registry, + permissions=PermissionEngine(workspace_root=tmp_path), + model="gpt-5.5", + ) + + async def _collect(): + return [ev async for ev in engine.run("hello")] + + return engine, asyncio.run(_collect()) + + +def test_engine_attaches_usage_to_event_and_message(tmp_path): + engine, events = _run_engine(tmp_path) + assistant = next(ev for ev in events if ev.type == EventType.ASSISTANT_MESSAGE) + expected = { + "model": "gpt-5.5", + "input": 100, + "output": 20, + "cache_read": 5, + "cache_write": 0, + } + assert assistant.data["usage"] == expected + persisted = next(m for m in engine.messages if m.get("role") == "assistant") + assert persisted["usage"] == expected + + +def test_outbound_messages_strip_usage_sidecar(tmp_path): + engine, _ = _run_engine(tmp_path) + assert all("usage" not in m for m in engine._outbound_messages()) + + +# -- matrix ------------------------------------------------------------------------- + + +def test_model_context_windows_covers_verified_entries_only(): + windows = model_context_windows() + assert windows["anthropic:claude-fable-5"] == 1_000_000 + assert "together:thinkingmachines/Inkling" not in windows # unverified stays absent + assert all(isinstance(v, int) and v > 0 for v in windows.values())