Merge branch 'main' of https://github.com/andrewyng/openworker into feature/memory

This commit is contained in:
Devika Verma
2026-07-29 12:04:52 +05:30
20 changed files with 1155 additions and 58 deletions
+14 -6
View File
@@ -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.
+57 -1
View File
@@ -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,
)
)
+32
View File
@@ -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:
+18
View File
@@ -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,
)
)
+23
View File
@@ -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,
)
)
+81 -43
View File
@@ -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 `<lab>/<model>` (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<n>:<m>` 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
+31
View File
@@ -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,
)
)
+4 -1
View File
@@ -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
+1 -1
View File
@@ -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
+17 -1
View File
@@ -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") {
+66
View File
@@ -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);
});
+36 -3
View File
@@ -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<string[]>([]);
const [modelLabels, setModelLabels] = useState<Record<string, string>>({});
// {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<Record<string, number>>({});
// 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<SessionUsage>(emptyUsage());
const [surfaces, setSurfaces] = useState<SurfaceVisibility>({ 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)"
+6
View File
@@ -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<string, string>;
// {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<string, number>;
// 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";
+149 -1
View File
@@ -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) {
<span className="ml-auto" />
{/* 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 && (
<UsageChip
usage={props.usage}
contextWindow={props.contextWindow}
model={props.model}
modelLabels={props.modelLabels}
/>
)}
{/* 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<string, string>;
}) {
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) => (
<div className="flex items-baseline justify-between text-[11.5px] leading-snug">
<span className="text-faint">{label}</span>
<span className="text-ink tabular-nums">{formatTokens(value)}</span>
</div>
);
return (
<div className="relative">
<button
className="inline-flex items-center gap-1.5 px-2 py-1 rounded-lg text-[11.5px] text-muted hover:text-ink hover:bg-paper shrink-0"
onClick={() => setOpen((v) => !v)}
aria-haspopup="menu"
aria-expanded={open}
aria-label="Token usage"
title={
pct !== null
? `Token usage — ${pct}% of the context window used`
: "Token usage this session"
}
data-testid="usage-chip"
>
{pct !== null && (
<span className="w-7 h-1 rounded-full bg-line overflow-hidden" aria-hidden="true">
<span
className="block h-full bg-accent transition-all"
style={{ width: `${Math.max(pct, 4)}%` }}
/>
</span>
)}
<span className="tabular-nums">{formatTokens(total)}</span>
</button>
{open && (
<>
<div className="fixed inset-0 z-30" onClick={() => setOpen(false)} />
<div
className="absolute z-40 bottom-full mb-1 right-0 w-[280px] rounded-xl border border-line bg-panel shadow-2xl p-3"
role="menu"
data-testid="usage-popover"
>
{contextWindow ? (
<div className="mb-2.5">
<div className="text-[10.5px] uppercase tracking-[0.06em] text-faint font-semibold mb-1">
Context window
</div>
<div className="h-1.5 rounded-full bg-line overflow-hidden">
<div
className="h-full bg-accent transition-all"
style={{ width: `${pct}%` }}
/>
</div>
<div className="mt-1 text-[11.5px] text-muted tabular-nums">
{formatTokens(usage.context)} of {formatTokens(contextWindow)} · {pct}%
</div>
</div>
) : usage.context > 0 ? (
<div className="mb-2.5 text-[11.5px] text-muted tabular-nums">
In context now: {formatTokens(usage.context)} tokens
</div>
) : null}
<div className="text-[10.5px] uppercase tracking-[0.06em] text-faint font-semibold mb-1">
Session totals
</div>
<div className="flex flex-col gap-1.5">
{Object.entries(usage.byModel).map(([id, t]) => (
<div key={id}>
<div className="text-[12px] text-ink font-medium truncate" title={id}>
{labelFor(id)}
</div>
{/* 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. */}
<div className="mt-0.5 flex flex-col gap-0.5">
{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)}
</div>
</div>
))}
</div>
<div className="mt-2 pt-2 border-t border-line flex items-baseline justify-between text-[11.5px]">
<span className="text-faint">Total</span>
<span className="text-ink tabular-nums">{formatTokens(total)} tokens</span>
</div>
{model && !modelLabels?.[model] && contextWindow === undefined && (
<div className="mt-1 text-[10.5px] text-faint leading-snug">
Context meter unavailable for custom models.
</div>
)}
</div>
</>
)}
</div>
);
}
// 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).
+19
View File
@@ -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<string, TurnUsage>;
context: number;
}
export interface SessionInfo {
session_id: string;
title?: string;
+78
View File
@@ -0,0 +1,78 @@
import { describe, expect, it } from "vitest";
import { addTurnUsage, emptyUsage, formatTokens, totalTokens, usageFromMessages } from "./usage";
const turn = (over: Record<string, unknown> = {}) => ({
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");
});
});
+75
View File
@@ -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";
}
+87
View File
@@ -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
+5 -1
View File
@@ -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
+356
View File
@@ -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())