mirror of
https://github.com/andrewyng/openworker.git
synced 2026-09-03 04:49:26 +00:00
Merge branch 'main' of https://github.com/andrewyng/openworker into feature/memory
This commit is contained in:
+14
-6
@@ -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.
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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") {
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
@@ -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)"
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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");
|
||||
});
|
||||
});
|
||||
@@ -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";
|
||||
}
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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())
|
||||
Reference in New Issue
Block a user