From da5c0b9252c11e7edef1d28e5ad55ee9da3da416 Mon Sep 17 00:00:00 2001 From: Josh Joseph Date: Mon, 12 Oct 2026 19:06:46 +0000 Subject: [PATCH 1/7] Show the smart-routing savings estimate in the Claude mod The launch-scoped ug-smart-router mod now draws a dim row above the prompt with the estimated savings from subagent smart routing and, when Smart Router Orchestrator is active, the skill's version from SKILL.claude.md. The mod sums each turn request's tokens by served and baseline model and runs a stdlib-only pricer against cached endpoint rates, which each routed launch refreshes in the background. The pricer, usage file and version reach the mod as plugin options. ENABLE_SMART_ROUTING_SAVINGS=0 opts a launch out. Co-authored-by: Isaac --- docs/claude-mod.md | 21 ++ src/ucode/agents/claude_mods/register.ts | 7 +- src/ucode/agents/claude_mods/savings.ts | 274 +++++++++++++++++ src/ucode/databricks.py | 32 ++ src/ucode/smart_routing/orchestrator.py | 22 ++ src/ucode/smart_routing/pricing.py | 291 ++++++++++++++++++ src/ucode/smart_routing/savings.py | 168 +++++++++++ src/ucode/smart_routing/v2.py | 143 ++++++++- tests/README.md | 7 + tests/conftest.py | 2 + tests/integration/README.md | 5 +- tests/integration/utils/sql.py | 36 ++- tests/native_claude/session.py | 6 +- tests/native_claude/test_mod.py | 61 +++- tests/test_claude_mod.py | 72 ++++- tests/test_claude_smart_routing_v2.py | 213 +++++++++++++- tests/test_claude_windows_smart_routing.py | 1 + tests/test_databricks.py | 55 ++++ tests/test_skills.py | 42 +++ tests/test_smart_routing_pricing.py | 300 +++++++++++++++++++ tests/test_smart_routing_savings.py | 324 +++++++++++++++++++++ 21 files changed, 2056 insertions(+), 26 deletions(-) create mode 100644 src/ucode/agents/claude_mods/savings.ts create mode 100644 src/ucode/smart_routing/pricing.py create mode 100644 src/ucode/smart_routing/savings.py create mode 100644 tests/test_smart_routing_pricing.py create mode 100644 tests/test_smart_routing_savings.py diff --git a/docs/claude-mod.md b/docs/claude-mod.md index 736260a5a..df81b321e 100644 --- a/docs/claude-mod.md +++ b/docs/claude-mod.md @@ -27,3 +27,24 @@ UG passes the resolved routing flags and recipe in the existing launch settings `pluginConfigs["ug-smart-router"].options`. Claude supplies these to `register(on, options)`. `unset` lists originally absent flags because native options cannot carry null values. The source remains static; no additional launch snapshot file is created. + +## Savings estimate + +`hooks/savings.ts` draws a dim row above the prompt: the estimated savings from smart +routing and, when the launch activates Smart Router Orchestrator, the skill's +`metadata.version` from `SKILL.claude.md`. For each turn request, main agent and +subagents, the mod sums token usage by the model that served it and its baseline. Main +is its own baseline, and a subagent's baseline is the main model in effect when it ran. +Under first-prompt routing, requests before the user changes the main model are priced +against the model the session started on. + +After each turn the mod writes those sums to `savings_usage`, a file in the launch's +temporary directory. It then runs the `savings_pricer` argv +(`python -P -m ucode.smart_routing.savings`), which prices the sums from the +workspace's cached endpoint rates and prints the row's text. Each routed launch +refreshes that cache on a background thread. An estimate is hidden, not undercounted, +when any served or baseline model has no rate. The orchestrator version is passed +as `orchestrator_version`. + +The estimate is on by default; `ENABLE_SMART_ROUTING_SAVINGS=0` opts a launch out +and leaves the three options at their declared defaults. diff --git a/src/ucode/agents/claude_mods/register.ts b/src/ucode/agents/claude_mods/register.ts index 04a8adee7..44a95a086 100644 --- a/src/ucode/agents/claude_mods/register.ts +++ b/src/ucode/agents/claude_mods/register.ts @@ -1,6 +1,9 @@ import type { On, PluginOptions } from 'claude-code'; +import { registerSavings, savingsSessionStart } from './savings'; + /** Entry point for UG's in-process Claude hooks. */ -export function register(on: On, _options: PluginOptions): void { - on('session.start', ($, e, next) => next(e)); +export function register(on: On, options: PluginOptions): void { + on('session.start', savingsSessionStart); + registerSavings(on, options); } diff --git a/src/ucode/agents/claude_mods/savings.ts b/src/ucode/agents/claude_mods/savings.ts new file mode 100644 index 000000000..2aaebe84f --- /dev/null +++ b/src/ucode/agents/claude_mods/savings.ts @@ -0,0 +1,274 @@ +import type { On, PluginOptions } from 'claude-code'; + +// Smart-routing savings estimate, drawn above the prompt with the Smart Router +// Orchestrator skill version. Per turn request (main agent and subagents) it sums +// token usage by (agent, baseline, served, before_user_switch), where served is +// the model the request named (the model service billed; the id the API reports +// back is a gateway deployment's, which drifts) and baseline is what the request +// would have used without routing: the main model in effect for a subagent, served +// itself for main (routing never changes the main model mid-session). First-prompt +// routing switches the main model once, before the first request, so a later +// change is the user's; requests before it are flagged, and priced from +// start_model under that routing. +// Claude Code's side queries (title, compaction, helpers) are not turn steps, so +// they never count. A hooks module has no Node APIs and the price table lives in +// Python, so after each turn it writes the sums to the launch's savings_usage +// file and runs the savings_pricer command (ucode.smart_routing.savings) on it. +// Everything is best effort and swallows errors. + +const USAGE_VERSION = 1; +const PRICER_TIMEOUT_MS = 20_000; +const TOKEN_KEYS = [ + 'input_tokens', + 'output_tokens', + 'cache_creation_input_tokens', + 'cache_read_input_tokens', +] as const; +const CACHE_SPLIT_KEYS = ['ephemeral_5m_input_tokens', 'ephemeral_1h_input_tokens'] as const; + +type Totals = Record<(typeof TOKEN_KEYS)[number], number> & { + cache_creation?: Partial>; +}; +type Entry = { + agent: string; + baseline: string; + served: string; + before_user_switch: boolean; + usage: Totals; +}; + +const num = (value: unknown): number => + typeof value === 'number' && Number.isFinite(value) ? value : 0; + +const str = (value: unknown): string | null => (typeof value === 'string' && value ? value : null); + +// Set from the launch's plugin options by registerSavings. +let pricer: string[] = []; +let usagePath: string | null = null; +let enabled = false; +let version: string | null = null; + +const entries = new Map(); +// The model the session started on, before any first-prompt routing. Set once. +let startModel: string | null = null; +// The model the main loop last named, so a subagent's baseline is a real model +// id ($.session.model() may return an alias). +let lastMainModel: string | null = null; +// The model the main loop's first request named (first-prompt routing's pick). +let firstMainModel: string | null = null; +// Whether a main request has since named another model, which only the user does. +let userSwitched = false; +let seeding: Promise | undefined; +let estimate: string | null = null; +let dirty = false; +let running = false; +let rerun = false; + +async function mainModel($: any): Promise { + try { + return str(await $.session.model()); + } catch { + return null; + } +} + +function accumulate( + agent: string, + baseline: string, + served: string, + beforeUserSwitch: boolean, + usage: any, +): void { + const key = JSON.stringify([agent, baseline, served, beforeUserSwitch]); + let entry = entries.get(key); + if (!entry) { + entry = { + agent, + baseline, + served, + before_user_switch: beforeUserSwitch, + usage: { + input_tokens: 0, + output_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + }, + }; + entries.set(key, entry); + } + for (const k of TOKEN_KEYS) entry.usage[k] += num(usage[k]); + // The 5m/1h cache-write split is only summed where the API reported it. + const split = usage.cache_creation; + if (split && typeof split === 'object') { + const sums = entry.usage.cache_creation ?? {}; + for (const k of CACHE_SPLIT_KEYS) if (k in split) sums[k] = (sums[k] ?? 0) + num(split[k]); + entry.usage.cache_creation = sums; + } +} + +// A mod reload resets module state while the session carries on, so the first +// hook of a fresh module instance picks up where the last one left off from the +// file it wrote. A missing or malformed file just means starting empty. +async function loadUsage($: any): Promise { + try { + const saved = JSON.parse(await $.fs.read(usagePath)); + if (!saved || saved.version !== USAGE_VERSION || !Array.isArray(saved.entries)) return; + for (const item of saved.entries) { + const agent = str(item?.agent); + const baseline = str(item?.baseline); + const served = str(item?.served); + if (agent && baseline && served && item.usage && typeof item.usage === 'object') { + accumulate(agent, baseline, served, item.before_user_switch === true, item.usage); + dirty = true; + } + } + startModel ??= str(saved.start_model); + lastMainModel ??= str(saved.main_model); + firstMainModel ??= str(saved.first_main_model); + userSwitched ||= saved.user_switched === true; + } catch {} +} + +// State is seeded once per module instance, by whichever hook runs first; +// concurrent hooks wait on the same load. +async function ready($: any): Promise { + if (!enabled) return false; + seeding ??= loadUsage($); + await seeding; + return true; +} + +// The pricer's last stdout line; a failed run drops the estimate, which would be stale. +async function runPricer($: any): Promise { + try { + const run = await $.process.run([...pricer, '--usage', usagePath], { + timeoutMs: PRICER_TIMEOUT_MS, + }); + if (run?.exitCode === 0) { + const out = JSON.parse(String(run.stdout ?? '').trim().split('\n').pop() ?? ''); + if (out && typeof out === 'object') return str(out.savings); + } + } catch {} + return null; +} + +async function writeAndPrice($: any): Promise { + dirty = false; + try { + await $.fs.write( + usagePath, + JSON.stringify({ + version: USAGE_VERSION, + start_model: startModel, + main_model: lastMainModel, + first_main_model: firstMainModel, + user_switched: userSwitched, + entries: Array.from(entries.values()), + }), + ); + } catch (error) { + // Retry at the next turn.complete rather than pricing a file that is not there. + dirty = true; + throw error; + } + const next = await runPricer($); + if (next !== estimate) { + estimate = next; + // Claude Code only redraws the band when its props change, not when module state does. + $.ui.invalidate('ui.render'); + } +} + +// One pricer run at a time; a request that arrives mid-run just asks for another. +async function flush($: any): Promise { + if (running) { + rerun = true; + return; + } + running = true; + try { + do { + rerun = false; + await writeAndPrice($); + } while (rerun); + } catch { + } finally { + running = false; + } +} + +const line = (): string | null => + [estimate, version && `Smart Router Orchestrator v${version}`].filter(Boolean).join(' ยท ') || + null; + +// The module's session.start hook, which register.ts installs (a second unmatched +// hook on the same event fails the module load). It runs before the first prompt, +// so the start model is captured before any routing happens. session.start can +// fire again (/clear), by when first-prompt routing may have switched the model, +// so the start model is only ever set once. The flush is not awaited: it spawns a +// process and Claude Code waits on this hook before the first prompt. +export async function savingsSessionStart($: any, e: any, next: any): Promise { + if (await ready($)) { + startModel ??= await mainModel($); + void flush($); + } + return next(e); +} + +export function registerSavings(on: On, options: PluginOptions): void { + const raw = options as Record; + pricer = + Array.isArray(raw.savings_pricer) && raw.savings_pricer.every(arg => typeof arg === 'string') + ? (raw.savings_pricer as string[]) + : []; + usagePath = str(raw.savings_usage); + enabled = pricer.length > 0 && usagePath !== null; + version = str(raw.orchestrator_version); + + // Fires for every request to a model, subagents included (e.agentId is set). + // yield* forwards the streamed response untouched; only the finished result's + // usage is read. + on('turn.step', async function* ($, e, next) { + const request = e as any; + const agent = str(request.agentId); + const active = await ready($); + let baseline: string | null = null; + if (active) { + if (agent === null) { + lastMainModel = str(request.model) ?? lastMainModel; + firstMainModel ??= lastMainModel; + userSwitched ||= lastMainModel !== firstMainModel; + } else baseline = lastMainModel ?? (await mainModel($)); + } + // Read before the request runs: a concurrent main request may switch it. + const beforeUserSwitch = !userSwitched; + const result = yield* next(e); + try { + const usage = (result as any)?.usage; + const served = str(request.model) ?? str(usage?.model); + // Main is its own baseline: routing leaves the main model alone. + const base = agent === null ? served : baseline; + if (active && usage && served && base) { + accumulate(agent ?? 'main', base, served, beforeUserSwitch, usage); + dirty = true; + } + } catch {} + return result; + }); + + // Price at the end of each turn, off the turn's critical path. + on('turn.complete', async ($, e, next) => { + const result = await next(e); + if ((await ready($)) && dirty) void flush($); + return result; + }); + + on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => { + const theirs = await next(e); + const text = line(); + if (text === null) return theirs; + const { Box, Text } = $.ui.resolve(e); + const row = Text({ dimColor: true, children: [text] }); + return theirs ? Box({ flexDirection: 'column', children: [row, theirs] }) : row; + }); +} diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 7ba376617..8945b82de 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2411,6 +2411,38 @@ def fetch_external_model_prices(workspace: str, token: str) -> tuple[list[dict], return models, None +# Per-token rates for Databricks-hosted (`system.ai`) model services from the gateway's +# pay-per-token endpoint-rates API, which the smart-routing savings estimate prices usage with. +# The model names travel as repeated `databricks_hosted_model_services` query parameters. +_ENDPOINT_RATES_API_PATH = "/api/ai-gateway/v2/endpoint-rates:batchGet" +_SYSTEM_AI_MODEL_PREFIX = "system.ai." + + +def fetch_endpoint_rates( + workspace: str, token: str, model_services: list[str] +) -> tuple[list[dict], str | None]: + """Batch-fetch per-token rates for ``system.ai`` model services. + + Returns ``(rates, reason)`` with each rate the raw ``EndpointRate`` entry; models the caller + can't see are absent rather than errors. Names outside ``system.ai`` are dropped because the + API rejects them. ``reason`` is non-None on failure so callers omit the estimate, not fail. + """ + names = sorted({name for name in model_services if name.startswith(_SYSTEM_AI_MODEL_PREFIX)}) + if not names: + return [], "no system.ai model services to price" + query = urlencode([("databricks_hosted_model_services", name) for name in names]) + url = f"https://{workspace_hostname(workspace)}{_ENDPOINT_RATES_API_PATH}?{query}" + payload, reason = _http_get_json(url, token, timeout=30) + if reason is not None: + return [], reason + if not isinstance(payload, dict): + return [], "endpoint-rates returned an unexpected response shape" + rates = payload.get("model_service_rates") + if not isinstance(rates, list): + return [], None + return [rate for rate in rates if isinstance(rate, dict)], None + + # The `update_mask` paths a config PATCH sends. The server rejects paths outside its mutable set, # so this omits `spec_version` (an estore-internal format marker, still sent in the body; naming it # in the mask is the 400 this fixes) and the reserved/deprecated fields (`display_name`, `tracing`, diff --git a/src/ucode/smart_routing/orchestrator.py b/src/ucode/smart_routing/orchestrator.py index 8bc6afe16..85d24ebde 100644 --- a/src/ucode/smart_routing/orchestrator.py +++ b/src/ucode/smart_routing/orchestrator.py @@ -59,6 +59,28 @@ def skill_directory() -> Path: return skills._skills_source() / skills.SMART_ROUTER_ORCHESTRATOR_SKILL +def skill_version(*, agent: str) -> str | None: + """The ``metadata.version`` in the agent's skill frontmatter, or None when it can't be read.""" + try: + text = skills.skill_entrypoint(skill_directory(), agent).read_text(encoding="utf-8") + except (OSError, UnicodeError): + return None + lines = text.splitlines() + if not lines or lines[0].strip() != "---": + return None + in_metadata = False + for line in lines[1:]: + if line.strip() == "---": + break + if not line.startswith((" ", "\t")): + in_metadata = line.rstrip() == "metadata:" + continue + key, separator, value = line.strip().partition(":") + if in_metadata and separator and key == "version": + return value.strip().strip("\"'") or None + return None + + def add_claude_agents(plugin_dir: Path) -> None: """Load roles alongside the router's exact-model agents, only for this launch.""" if feature_enabled(agent=AGENT_CLAUDE): diff --git a/src/ucode/smart_routing/pricing.py b/src/ucode/smart_routing/pricing.py new file mode 100644 index 000000000..cee68cfda --- /dev/null +++ b/src/ucode/smart_routing/pricing.py @@ -0,0 +1,291 @@ +"""Per-token model prices and the cost of coding-agent token usage. + +The launcher fetches rates from the AI Gateway's endpoint-rates API (``prices_from_endpoint_rates``) +and caches them per workspace (``price_cache_path``) for the savings pricer to read; this module +defines that cache and the arithmetic. A token class with no rate makes a response unpriceable, and +the pricer drops the estimate rather than undercount it. The API returns per-million-token +dollar rates for input, output, and cache tokens (5-minute and 1-hour writes, and reads), which the +estimate needs because fixed multipliers don't hold across models (Opus 5.5 cache reads bill at +0.05x input, Opus 4.8's at 0.1x). A model whose org has no DBU-to-dollar conversion stays unpriced. + +Stdlib-only on purpose: the savings pricer imports this in a fresh process after every turn, and the +CLI's usual imports (Rich, Typer, the Databricks SDK, even ``urllib.request``) take about a second. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import re +import tempfile +import time +from collections.abc import Iterable, Mapping +from dataclasses import dataclass, fields +from decimal import Decimal, InvalidOperation +from pathlib import Path +from typing import Any, NamedTuple + +_PRICE_CACHE_VERSION = 1 +_MILLION = Decimal(1_000_000) +_CONTEXT_SUFFIX_RE = re.compile(r"\[[^\]]*\]$") +_ANTHROPIC_AIGW_PREFIX_RE = re.compile(r"^anthropic-aigw-[0-9a-f]{8}-") +# The gateway reports some served models by their provider id, e.g. Bedrock's +# `anthropic.claude-haiku-4-5-20251001-v1:0` for a request to `system.ai.claude-haiku-4-5`. +_PROVIDER_PREFIX_RE = re.compile(r"^(?:(?:us|eu|apac|au|jp|global)\.)?anthropic\.") +_PROVIDER_SUFFIX_RE = re.compile(r"(?:-20\d{6})?(?:-v\d+(?::\d+)?)?$") +# Other served models are reported by a deployment id that drifts from the model service name: +# `glm-5.3-flash` for `system.ai.glm-5-3-flash`, `glm-5-3-colo-on-sp-v1` for `system.ai.glm-5-3`. +_DEPLOYMENT_SUFFIX_RE = re.compile(r"-colo-on-[a-z0-9]+$") + + +def model_key(model: str) -> str: + """Collapse the spellings of one model id that Claude Code and the gateway use. + + Mirrors ``routing.unwrap_anthropic_gateway_model`` + ``routing.normalize_model`` without + importing ``routing``, whose ``urllib.request`` import alone costs ~0.2s of pricer startup. + Also drops a ``[1m]`` context-window selector (a window, not a price), and the provider prefix, + date/version suffix and deployment spelling of a served id, so it keys the same as the + requested model. + """ + name = _CONTEXT_SUFFIX_RE.sub("", (model or "").strip().lower()) + name = _ANTHROPIC_AIGW_PREFIX_RE.sub("", name).rsplit("/", 1)[-1] + for prefix in ("databricks-", "system.ai."): + if name.startswith(prefix): + name = name[len(prefix) :] + break + name = _PROVIDER_PREFIX_RE.sub("", name).split("@", 1)[0] + name = _DEPLOYMENT_SUFFIX_RE.sub("", _PROVIDER_SUFFIX_RE.sub("", name)) + # Model service names spell versions with dashes. + return name.replace(".", "-") + + +@dataclass(frozen=True) +class ModelPrice: + """USD per million tokens by token class; ``None`` where the gateway reports no rate.""" + + input: Decimal | None = None + output: Decimal | None = None + cache_read: Decimal | None = None + cache_write_5m: Decimal | None = None + cache_write_1h: Decimal | None = None + # Rates for requests whose prompt exceeds the threshold (e.g. Sonnet 4.x above 200k tokens). + long_context_threshold: int | None = None + long_context: ModelPrice | None = None + + +_RATE_FIELDS = tuple( + field.name for field in fields(ModelPrice) if not field.name.startswith("long_context") +) + + +class TokenUsage(NamedTuple): + """One model response's billed tokens, split by the classes that are priced separately.""" + + input: int = 0 + cache_write_5m: int = 0 + cache_write_1h: int = 0 + cache_read: int = 0 + output: int = 0 + + @classmethod + def from_message_usage( + cls, usage: Mapping[str, Any], *, uncovered_writes_1h: bool = False + ) -> TokenUsage: + """Read an Anthropic Messages ``usage`` object as Claude Code records it in transcripts. + + Cache writes are split by TTL because a 1-hour write bills at twice the input rate and a + 5-minute write at 1.25x; ug enables 1-hour caching. Writes the breakdown doesn't cover are + billed as 5-minute writes, the API's default TTL. A caller whose usage never carries the + breakdown, but whose session ran with 1-hour caching, sets ``uncovered_writes_1h`` so those + writes aren't undercounted at the 5-minute rate. + """ + creation = usage.get("cache_creation") + creation = creation if isinstance(creation, Mapping) else {} + write_1h = _count(creation.get("ephemeral_1h_input_tokens")) + write_5m = _count(creation.get("ephemeral_5m_input_tokens")) + uncovered = max(_count(usage.get("cache_creation_input_tokens")) - write_1h - write_5m, 0) + if uncovered_writes_1h: + write_1h += uncovered + else: + write_5m += uncovered + return cls( + input=_count(usage.get("input_tokens")), + cache_write_5m=write_5m, + cache_write_1h=write_1h, + cache_read=_count(usage.get("cache_read_input_tokens")), + output=_count(usage.get("output_tokens")), + ) + + @property + def prompt_tokens(self) -> int: + """Every input-side token, which is what long-context price thresholds measure.""" + return self.input + self.cache_write_5m + self.cache_write_1h + self.cache_read + + +def _count(value: object) -> int: + return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else 0 + + +def token_cost(price: ModelPrice, tokens: TokenUsage) -> Decimal | None: + """USD cost of one response, or None when a token class it used has no rate. + + Returning None rather than pricing a missing rate at zero keeps callers from showing a figure + that silently undercounts. + """ + if ( + price.long_context is not None + and price.long_context_threshold is not None + and tokens.prompt_tokens > price.long_context_threshold + ): + price = price.long_context + total = Decimal(0) + for count, rate in ( + (tokens.input, price.input), + (tokens.cache_write_5m, price.cache_write_5m), + (tokens.cache_write_1h, price.cache_write_1h), + (tokens.cache_read, price.cache_read), + (tokens.output, price.output), + ): + if count <= 0: + continue + if rate is None: + return None + total += Decimal(count) * rate + return total / _MILLION + + +def _rate(raw: object) -> Decimal | None: + if isinstance(raw, bool) or not isinstance(raw, int | float | str) or not str(raw).strip(): + return None + try: + rate = Decimal(str(raw).strip()) + except InvalidOperation: + return None + return rate if rate.is_finite() and rate >= 0 else None + + +def _dump_price(price: ModelPrice) -> dict[str, Any]: + dumped: dict[str, Any] = { + name: str(value) for name in _RATE_FIELDS if (value := getattr(price, name)) is not None + } + if price.long_context is not None and price.long_context_threshold is not None: + dumped["long_context_threshold"] = price.long_context_threshold + dumped["long_context"] = _dump_price(price.long_context) + return dumped + + +def _load_price(raw: object) -> ModelPrice | None: + if not isinstance(raw, Mapping): + return None + rates = {name: _rate(raw.get(name)) for name in _RATE_FIELDS} + threshold = raw.get("long_context_threshold") + long_context = _load_price(raw.get("long_context")) + if isinstance(threshold, bool) or not isinstance(threshold, int) or long_context is None: + threshold, long_context = None, None + return ModelPrice(**rates, long_context_threshold=threshold, long_context=long_context) + + +# The endpoint-rates API groups each model's costs by unit; the estimate is in dollars. +_DOLLARS_UNIT = "USD" +# `TokenType` enum names โ†’ the ``ModelPrice`` field each bills. CACHE_CREATION is the API's default +# 5-minute cache write; CACHE_CREATION_1H is the 1-hour write ug enables. +_TOKEN_TYPE_FIELDS = { + "TOKEN_TYPE_INPUT": "input", + "TOKEN_TYPE_OUTPUT": "output", + "TOKEN_TYPE_CACHE_CREATION": "cache_write_5m", + "TOKEN_TYPE_CACHE_CREATION_1H": "cache_write_1h", + "TOKEN_TYPE_CACHE_READ": "cache_read", +} + + +def prices_from_endpoint_rates(rates: Iterable[object]) -> dict[str, ModelPrice]: + """Map endpoint-rates ``EndpointRate`` entries to prices keyed by their model service. + + Reads the ``USD`` group of each entry's ``costs`` (a per-million-token ``cost`` per + ``token_type``). The API omits the dollar group when the org has no DBU-to-dollar conversion + configured; those models are left unpriced rather than shown in DBUs. + """ + prices: dict[str, ModelPrice] = {} + for rate in rates: + if not isinstance(rate, Mapping): + continue + service = rate.get("model_service") + costs = rate.get("costs") + if not isinstance(service, str) or not service or not isinstance(costs, list): + continue + dollars = next( + (c for c in costs if isinstance(c, Mapping) and c.get("unit") == _DOLLARS_UNIT), None + ) + token_costs = dollars.get("token_costs") if isinstance(dollars, Mapping) else None + if not isinstance(token_costs, list): + continue + rate_fields: dict[str, Decimal] = {} + for entry in token_costs: + if not isinstance(entry, Mapping): + continue + field = _TOKEN_TYPE_FIELDS.get(entry.get("token_type")) + value = _rate(entry.get("cost")) + if field is not None and value is not None: + rate_fields[field] = value + price = ModelPrice( + input=rate_fields.get("input"), + output=rate_fields.get("output"), + cache_read=rate_fields.get("cache_read"), + cache_write_5m=rate_fields.get("cache_write_5m"), + cache_write_1h=rate_fields.get("cache_write_1h"), + ) + if price.input is not None or price.output is not None: + prices[service] = price + return prices + + +def price_cache_path(app_dir: Path, workspace: str) -> Path: + """The workspace's price cache; dollar rates depend on its org's DBU conversion.""" + host = workspace.strip().lower().removeprefix("https://").rstrip("/") + digest = hashlib.sha256(host.encode("utf-8")).hexdigest()[:16] + return app_dir / f"model-prices-{digest}.json" + + +def write_price_cache( + path: Path, prices: Mapping[str, ModelPrice], *, now: float | None = None +) -> None: + """Atomically replace the price cache so a concurrent pricer never reads a partial file. + + ``prices`` is keyed by any spelling of the model id; entries are stored under ``model_key``. + """ + payload = { + "version": _PRICE_CACHE_VERSION, + "fetched_at": now if now is not None else time.time(), + "models": {model_key(model): _dump_price(price) for model, price in prices.items()}, + } + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.", suffix=".tmp") + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(payload, handle, separators=(",", ":")) + os.replace(temporary, path) + except BaseException: + Path(temporary).unlink(missing_ok=True) + raise + + +def read_price_cache(path: Path) -> tuple[dict[str, ModelPrice], str] | None: + """Load prices keyed by ``model_key``, plus a fingerprint that changes on every refresh.""" + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + if not isinstance(payload, dict) or payload.get("version") != _PRICE_CACHE_VERSION: + return None + models = payload.get("models") + if not isinstance(models, dict): + return None + prices = { + key: price + for key, raw in models.items() + if isinstance(key, str) and (price := _load_price(raw)) is not None + } + if not prices: + return None + return prices, str(payload.get("fetched_at")) diff --git a/src/ucode/smart_routing/savings.py b/src/ucode/smart_routing/savings.py new file mode 100644 index 000000000..a25cbba25 --- /dev/null +++ b/src/ucode/smart_routing/savings.py @@ -0,0 +1,168 @@ +"""Price the Claude mod's token sums into a smart-routing savings estimate. + +The savings estimate assumes every token the session used, main agent and subagents alike, would +otherwise have run on the baseline main-agent model:: + + saved = sum(tokens * price(baseline)) - sum(tokens * price(model that served them)) + +per token class. The mod (``agents/claude_mods/savings.ts``) sums each turn request's usage by +(agent, baseline, served) and runs this module on that file after every turn; it prints one JSON +line with the segment the mod draws. A main-agent request is its own baseline, and a subagent's is +the main model in effect when it ran, so switching models with ``/model`` never reprices earlier +work. Under first-prompt routing, where the router also picks the main model, requests from before +the user changed the main model are priced against the model the session started on. A negative +result (routing chose pricier models) is shown as a cost increase. + +Stdlib-only besides ``pricing`` on purpose: the mod runs this in a fresh process after every turn, +and the CLI's usual imports take about a second. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from decimal import ROUND_HALF_UP, Decimal +from pathlib import Path + +from ucode.smart_routing.pricing import ( + ModelPrice, + TokenUsage, + model_key, + read_price_cache, + token_cost, +) + +MODULE = "ucode.smart_routing.savings" +# The file the mod writes its token sums to, in the launch directory; the mod reads its path from +# the plugin options, so only the version is shared with it. +USAGE_FILENAME = "savings-usage.json" +_USAGE_VERSION = 1 +_CENT = Decimal("0.01") + + +def pricer_argv(*, python: str, price_cache: Path, baseline_session_start: bool) -> list[str]: + """The command the Claude mod runs (plus ``--usage PATH``) to price its token sums.""" + # -P keeps the launch directory off sys.path, so a project file can't shadow a stdlib module. + argv = [python, "-P", "-m", MODULE, "--price-cache", str(price_cache)] + if baseline_session_start: + argv.append("--baseline-session-start") + return argv + + +def _response_costs( + prices: dict[str, ModelPrice], served_key: str, baseline_key: str, tokens: TokenUsage +) -> tuple[Decimal, Decimal] | None: + """``(actual, baseline)`` dollars for one entry, or None when either model can't be priced.""" + served_price = prices.get(served_key) + baseline_price = prices.get(baseline_key) + if served_price is None or baseline_price is None: + return None + actual = token_cost(served_price, tokens) + baseline = token_cost(baseline_price, tokens) + if actual is None or baseline is None: + return None + return actual, baseline + + +def usage_savings( + usage_path: Path, *, price_cache: Path, baseline_session_start: bool +) -> tuple[Decimal, Decimal] | None: + """The ``(saved, baseline)`` dollars for the mod's token sums, or None with nothing to claim. + + None until routing changed some model, and whenever an entry is malformed, baseline-less or + unpriced: an undercounted figure would be worse than none. A file that can't be read as the + mod's format at all is an error, not an empty estimate. + """ + document = json.loads(usage_path.read_text(encoding="utf-8")) + if not isinstance(document, dict) or document.get("version") != _USAGE_VERSION: + raise ValueError(f"unsupported savings usage file: {usage_path}") + entries = document.get("entries") + if not isinstance(entries, list): + raise ValueError(f"unsupported savings usage file: {usage_path}") + fixed_baseline = None + if baseline_session_start: + start = document.get("start_model") + if not isinstance(start, str) or not start: + return None + fixed_baseline = start + cached = read_price_cache(price_cache) + if cached is None: + return None + prices = cached[0] + + actual_total = baseline_total = Decimal(0) + rerouted = False + for entry in entries: + if not isinstance(entry, dict): + return None + usage, served = entry.get("usage"), entry.get("served") + baseline = entry.get("baseline") + if fixed_baseline is not None and entry.get("before_user_switch") is True: + baseline = fixed_baseline + if not isinstance(usage, dict) or not isinstance(served, str) or not served: + return None + if not isinstance(baseline, str) or not baseline: + return None + # The mod's usage has no cache-write TTL split, and ug launches Claude with 1-hour caching. + tokens = TokenUsage.from_message_usage(usage, uncovered_writes_1h=True) + if not any(tokens): + continue + served_key, baseline_key = model_key(served), model_key(baseline) + costs = _response_costs(prices, served_key, baseline_key, tokens) + if costs is None: + return None + actual_total += costs[0] + baseline_total += costs[1] + rerouted = rerouted or served_key != baseline_key + if not rerouted or baseline_total <= 0: + return None + return baseline_total - actual_total, baseline_total + + +def savings_text(saved: Decimal, baseline: Decimal) -> str: + """A money-bag estimate when positive, an honest cost increase when not.""" + percent = (abs(saved) / baseline * 100).quantize(Decimal(1), rounding=ROUND_HALF_UP) + magnitude = abs(saved) + amount = ( + "<$0.01" if magnitude < _CENT else f"${magnitude.quantize(_CENT, rounding=ROUND_HALF_UP):,}" + ) + if saved >= 0: + return f"๐Ÿ’ฐ Est. saved with smart routing: {amount} ({percent}%)" + return f"Smart routing cost {amount} more ({percent}%)" + + +def usage_output(usage_path: Path, *, price_cache: Path, baseline_session_start: bool) -> str: + """The one JSON line the mod reads: its savings segment, or null.""" + # The mod must get a parseable line whatever goes wrong. + savings = None + try: + computed = usage_savings( + usage_path, price_cache=price_cache, baseline_session_start=baseline_session_start + ) + savings = None if computed is None else savings_text(*computed) + except Exception: # noqa: BLE001 + pass + return json.dumps({"savings": savings}, ensure_ascii=False) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + prog=MODULE, description="Price the Claude mod's token sums as a JSON line." + ) + parser.add_argument("--usage", type=Path, required=True) + parser.add_argument("--price-cache", type=Path, required=True) + parser.add_argument("--baseline-session-start", action="store_true") + args = parser.parse_args(argv) + line = usage_output( + args.usage, + price_cache=args.price_cache, + baseline_session_start=args.baseline_session_start, + ) + # Write bytes so the money bag survives a non-UTF-8 stdout locale. + sys.stdout.buffer.write((line + "\n").encode("utf-8")) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 987a4c1d6..d1f64b41b 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -7,6 +7,7 @@ import socket import subprocess import sys +import threading import time import urllib.request from collections.abc import Callable, Mapping, MutableMapping @@ -40,6 +41,7 @@ from ucode.databricks import ( AnthropicModelCatalog, build_auth_token_argv, + fetch_endpoint_rates, get_databricks_token, list_anthropic_model_catalog, list_anthropic_models, @@ -51,7 +53,14 @@ release_file_lock, ) from ucode.skills import SMART_ROUTER_ORCHESTRATOR_SKILL, SMART_ROUTER_SKILL, install_skill -from ucode.smart_routing import claude_routing, codex_interposer, orchestrator, routing +from ucode.smart_routing import ( + claude_routing, + codex_interposer, + orchestrator, + pricing, + routing, + savings, +) from ucode.smart_routing.claude_hooks import ( FIRST_PROMPT_SOCKET_ENV, sync_first_prompt_hook, @@ -77,6 +86,15 @@ CLAUDE_ROUTE_SELECTION_TIMEOUT_S = 20.0 CLAUDE_ROUTED_AGENT_PREFIX = "ucode-route-" CLAUDE_ROUTING_PLUGIN_NAME = "ug-smart-router" +# Default on; "0" opts a launch out of the smart-routing savings estimate. +ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR = "ENABLE_SMART_ROUTING_SAVINGS" +# Native plugin options for the savings mod, at the values Claude fills in for a launch that omits +# them: no pricer (the estimate is off) and no orchestrator skill version to show. +CLAUDE_SAVINGS_OPTION_DEFAULTS: dict[str, str | list[str]] = { + "savings_pricer": [], + "savings_usage": "", + "orchestrator_version": "", +} CLAUDE_ROUTED_AGENT_PROMPT = ( "Complete the delegated task exactly as requested. Follow the parent agent's instructions and " "return a concise report of your findings or changes." @@ -331,6 +349,88 @@ def _claude_mod_launch_options(env: Mapping[str, str] | None = None) -> dict: } +def savings_enabled(env: Mapping[str, str] | None = None) -> bool: + """Whether a routed Claude launch shows the savings estimate (default on; ``0`` opts out).""" + source = os.environ if env is None else env + return source.get(ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, "1") != "0" + + +def _claude_savings_launch_options( + launch_dir: Path, *, price_cache: Path, baseline_session_start: bool +) -> dict[str, str | list[str]]: + """Options for the savings mod: the pricer it runs after each turn, the file it sums usage + into, and the orchestrator skill version it shows when the launch activates that skill.""" + options = dict(CLAUDE_SAVINGS_OPTION_DEFAULTS) + if savings_enabled(): + options["savings_pricer"] = savings.pricer_argv( + python=sys.executable, + price_cache=price_cache, + baseline_session_start=baseline_session_start, + ) + options["savings_usage"] = str(launch_dir / savings.USAGE_FILENAME) + if orchestrator.feature_enabled(agent=AGENT_CLAUDE): + options["orchestrator_version"] = orchestrator.skill_version(agent=AGENT_CLAUDE) or "" + return options + + +# Claude settings env keys that name the main model a session may start on. +_MAIN_MODEL_ENV_KEYS = ( + "ANTHROPIC_MODEL", + "ANTHROPIC_DEFAULT_FABLE_MODEL", + "ANTHROPIC_DEFAULT_OPUS_MODEL", + "ANTHROPIC_DEFAULT_SONNET_MODEL", + "ANTHROPIC_DEFAULT_HAIKU_MODEL", +) + + +def _savings_model_services(models: list[str | None], settings: dict) -> list[str]: + """The ``system.ai`` model services a session's responses can be priced against. + + Covers the routable models plus the main model, which the baseline prices every token at and + which can come from a pinned ``--model`` or Claude's default-model env rather than the catalog. + """ + env = settings.get("env") + env = env if isinstance(env, dict) else {} + candidates = [*models, *(env.get(key) for key in _MAIN_MODEL_ENV_KEYS)] + services: set[str] = set() + for model in candidates: + if not isinstance(model, str) or not model: + continue + name = _canonical_claude_model_id(_unwrapped_claude_model_id(model.strip())) + name = name.removesuffix("[1m]") + if name.startswith("system.ai."): + services.add(name) + return sorted(services) + + +def _refresh_savings_prices( + workspace: str, token: str, model_services: list[str], price_cache: Path +) -> None: + """Fetch this launch's per-token rates and cache them for the savings pricer. + + On failure the previous cache stays; without one the estimate stays hidden. + """ + try: + rates, _reason = fetch_endpoint_rates(workspace, token, model_services) + prices = pricing.prices_from_endpoint_rates(rates) + if prices: + pricing.write_price_cache(price_cache, prices) + except Exception: # noqa: BLE001 - a background refresh must never surface in the agent's TUI + return + + +def _start_savings_price_refresh( + workspace: str, token: str, model_services: list[str], price_cache: Path +) -> None: + """Refresh prices off the launch path so a slow rates API never delays Claude's startup.""" + threading.Thread( + target=_refresh_savings_prices, + args=(workspace, token, model_services, price_cache), + name="ug-savings-prices", + daemon=True, + ).start() + + def _write_routed_claude_plugin(plugin_dir: Path, model_ids: list[str]) -> None: """Write exact-model agents for launch-scoped loading through --plugin-dir.""" write_json_file( @@ -358,6 +458,25 @@ def _write_routed_claude_plugin(plugin_dir: Path, model_ids: list[str]) -> None: "title": "Recipe", "description": "Original routing recipe", }, + "savings_pricer": { + "type": "string", + "multiple": True, + "default": CLAUDE_SAVINGS_OPTION_DEFAULTS["savings_pricer"], + "title": "Savings pricer", + "description": "Command that prices the session's token sums", + }, + "savings_usage": { + "type": "string", + "default": CLAUDE_SAVINGS_OPTION_DEFAULTS["savings_usage"], + "title": "Savings usage file", + "description": "Launch-scoped file the token sums are written to", + }, + "orchestrator_version": { + "type": "string", + "default": CLAUDE_SAVINGS_OPTION_DEFAULTS["orchestrator_version"], + "title": "Orchestrator skill version", + "description": "Smart Router Orchestrator skill version, when it is active", + }, }, }, ) @@ -379,8 +498,9 @@ def _write_routed_claude_plugin(plugin_dir: Path, model_ids: list[str]) -> None: ), ) orchestrator.add_claude_agents(plugin_dir) - source = Path(__file__).parents[1] / "agents" / "claude_mods" / "register.ts" - write_text_file(plugin_dir / "hooks" / "register.ts", source.read_text(encoding="utf-8")) + sources = Path(__file__).parents[1] / "agents" / "claude_mods" + for name in ("register.ts", "savings.ts"): + write_text_file(plugin_dir / "hooks" / name, (sources / name).read_text(encoding="utf-8")) write_json_file(plugin_dir / "hooks" / "hooks.json", {"modules": ["./register.ts"]}) @@ -617,9 +737,24 @@ def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: session_path = _prepare_smart_router_session("claude") env[SESSION_ENV_VAR] = str(session_path) env[SESSION_PYTHON_ENV_VAR] = os.environ[SESSION_PYTHON_ENV_VAR] + # The savings pricer reads rates from this cache, since it runs after every turn and + # can't wait on the network; the refresh below fills it. + price_cache = pricing.price_cache_path(APP_DIR, workspace) settings.setdefault("pluginConfigs", {}).setdefault(CLAUDE_ROUTING_PLUGIN_NAME, {})[ "options" - ] = _claude_mod_launch_options({**resolve_environment(agent=AGENT_CLAUDE), **env}) + ] = { + **_claude_mod_launch_options({**resolve_environment(agent=AGENT_CLAUDE), **env}), + **_claude_savings_launch_options( + launch_dir, price_cache=price_cache, baseline_session_start=route_first_prompt + ), + } + if savings_enabled(): + _start_savings_price_refresh( + workspace, + token, + _savings_model_services([*model_ids, launch_model], settings), + price_cache, + ) try: write_json_file(settings_path, settings) _write_routed_claude_plugin(plugin_dir, model_ids) diff --git a/tests/README.md b/tests/README.md index 6d9c2b959..b26dc1a3a 100644 --- a/tests/README.md +++ b/tests/README.md @@ -17,6 +17,13 @@ Claude mod entry-point coverage lives in `test_claude_mod.py` (generated plugin `native_claude/test_mod.py` suite. Claude 2.1.290 validates the module and runs its TypeScript event-forwarding test in CI. Request metadata and toggle behavior are covered by a follow-up change. See [the mod guide](../docs/claude-mod.md). +Savings estimate coverage: `test_smart_routing_savings.py` prices the mod's token sums +(per-entry and first-prompt baselines, unpriced or malformed entries, the pricer's single +JSON line). `test_smart_routing_pricing.py` covers endpoint rates and the price cache. +`test_claude_mod.py` and `test_claude_smart_routing_v2.py` cover the savings options, +opt-out and price refresh. In the native suite, the options reach the mod unchanged and +a real session records its start model. Token accumulation in `turn.step` and the +above-prompt row are not covered automatically. Dedicated CUJs reuse `integration/utils` session/terminal mechanics and file-task and transcript readers, not its config stubs or pytest fixtures. Prompt/model diff --git a/tests/conftest.py b/tests/conftest.py index ab2f4f3d3..11c76e4cc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -77,6 +77,8 @@ def reject_privileged_write(path, _desired_text): monkeypatch.delenv("ENABLE_SMART_ROUTER_ORCHESTRATOR", raising=False) monkeypatch.delenv("SMART_ROUTER_CONFIG_VERSION", raising=False) monkeypatch.delenv("SMART_ROUTER_NAME", raising=False) + # A developer's savings opt-out would otherwise drop the savings options from launch tests. + monkeypatch.delenv("ENABLE_SMART_ROUTING_SAVINGS", raising=False) # On Windows, resolve_command swaps a bare program name for whatever `shutil.which` # finds on the developer's PATH (e.g. a real `codex.CMD`). Rebind only the compatibility # helper's `shutil` so argv stays host-independent; helper tests patch `which` explicitly. diff --git a/tests/integration/README.md b/tests/integration/README.md index c78ec6137..69a00ec3c 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -423,7 +423,10 @@ configuration, each journey injects only a tracing-enabled CodingAgentConfig inp `claude_tracing.json` / `codex_tracing.json` fixtures) through the suite's managed-config stub; the agents, inference, OTLP export, and table verification remain real. Each adds the prompt's UUID as a trace-safe `ug_integration_marker` attribute, resolves the destination table from the workspace tracing -configuration, waits 30 seconds, and queries that table through an existing SQL warehouse. +configuration, waits 30 seconds, and queries that table through an existing SQL warehouse +(a running one if any, else a serverless one). A query may wait up to 10 minutes for the +warehouse to start or queue it before its 180-second execution limit applies, and a timeout +reports which of the two it hit. The tests assert that a span with the marker arrived and identifies the requested model. Scoped discovery additionally requires Model Services diff --git a/tests/integration/utils/sql.py b/tests/integration/utils/sql.py index 03feb222e..40e879852 100644 --- a/tests/integration/utils/sql.py +++ b/tests/integration/utils/sql.py @@ -8,6 +8,10 @@ DEFAULT_TRACE_TABLE_PREFIX = "unity_gateway" TRACE_TABLE_SUFFIX = "_otel_spans" +# A statement stays PENDING while its warehouse starts or queues it; a stopped classic or pro +# warehouse can take minutes to start, which is warehouse latency, not a missing trace. +WAREHOUSE_WAIT_SECONDS = 600 +STATEMENT_RUN_SECONDS = 180 def resolve_trace_table(workspace: str, bearer: str) -> str: @@ -45,7 +49,7 @@ def _quote_identifier(identifier: str) -> str: def resolve_warehouse_id(workspace: str, bearer: str) -> str: - """Choose an existing warehouse, preferring one that is already running.""" + """Choose an existing warehouse: a running one, else a serverless one (quick to start).""" request = urllib.request.Request( f"{workspace.rstrip('/')}/api/2.0/sql/warehouses", headers={"Authorization": f"Bearer {bearer}"}, @@ -62,7 +66,15 @@ def resolve_warehouse_id(workspace: str, bearer: str) -> str: running = next( (warehouse for warehouse in warehouses if warehouse.get("state") == "RUNNING"), None ) - return str((running or warehouses[0])["id"]) + serverless = next( + ( + warehouse + for warehouse in warehouses + if warehouse.get("enable_serverless_compute") is True + ), + None, + ) + return str((running or serverless or warehouses[0])["id"]) def query_count( @@ -88,9 +100,23 @@ def query_count( with urllib.request.urlopen(request, timeout=60) as response: # noqa: S310 result = json.load(response) - deadline = time.monotonic() + 180 - while result.get("status", {}).get("state") in {"PENDING", "RUNNING"}: - assert time.monotonic() < deadline, "SQL statement did not finish within 180 seconds" + statement_id = result.get("statement_id") + submitted_at = time.monotonic() + running_since: float | None = None + while (state := result.get("status", {}).get("state")) in {"PENDING", "RUNNING"}: + now = time.monotonic() + if state == "RUNNING" and running_since is None: + running_since = now + if running_since is None: + assert now - submitted_at < WAREHOUSE_WAIT_SECONDS, ( + f"SQL statement {statement_id} stayed PENDING on warehouse {warehouse_id} " + f"for {WAREHOUSE_WAIT_SECONDS} seconds" + ) + else: + assert now - running_since < STATEMENT_RUN_SECONDS, ( + f"SQL statement {statement_id} on warehouse {warehouse_id} did not finish " + f"within {STATEMENT_RUN_SECONDS} seconds of running" + ) time.sleep(2) poll = urllib.request.Request( f"{workspace.rstrip('/')}/api/2.0/sql/statements/{result['statement_id']}", diff --git a/tests/native_claude/session.py b/tests/native_claude/session.py index b9b85bdca..379f44222 100644 --- a/tests/native_claude/session.py +++ b/tests/native_claude/session.py @@ -19,7 +19,7 @@ class Session: - def __init__(self, binary, root, env, api="http://127.0.0.1:1"): + def __init__(self, binary, root, env, api="http://127.0.0.1:1", options=None): self.root = root root.mkdir() self.plugin = root / "plugin" @@ -37,7 +37,9 @@ def __init__(self, binary, root, env, api="http://127.0.0.1:1"): "CLAUDE_CODE_EXTRA_BODY": '{"caller_field":"keep"}', }, "pluginConfigs": { - CLAUDE_ROUTING_PLUGIN_NAME: {"options": _claude_mod_launch_options(env)} + CLAUDE_ROUTING_PLUGIN_NAME: { + "options": {**_claude_mod_launch_options(env), **(options or {})} + } }, } process_env = { diff --git a/tests/native_claude/test_mod.py b/tests/native_claude/test_mod.py index b1bc3bb10..c1661218a 100644 --- a/tests/native_claude/test_mod.py +++ b/tests/native_claude/test_mod.py @@ -2,16 +2,36 @@ from __future__ import annotations +import json import os import shutil import subprocess +import sys +import time from pathlib import Path import pytest from session import Session -from ucode.smart_routing import config -from ucode.smart_routing.v2 import _claude_mod_launch_options, _write_routed_claude_plugin +from ucode.smart_routing import config, savings +from ucode.smart_routing.v2 import ( + CLAUDE_SAVINGS_OPTION_DEFAULTS, + _claude_mod_launch_options, + _write_routed_claude_plugin, +) + + +def _savings_options(root: Path) -> dict: + return { + # A space and a non-ASCII character check the argv survives the options round trip. + "savings_pricer": savings.pricer_argv( + python=sys.executable, + price_cache=root / "no prices ฮป.json", + baseline_session_start=False, + ), + "savings_usage": str(root / savings.USAGE_FILENAME), + "orchestrator_version": "1.2.0", + } @pytest.fixture(scope="module") @@ -48,10 +68,41 @@ def test_native_options_are_launch_local(native_binary, tmp_path): "SMART_ROUTER_CONFIG_VERSION": config.SUBAGENT_ORCH_V0_CLAUDE_ONLY, "SMART_ROUTER_NAME": 'custom "quoted" ฮป', } + savings_options = _savings_options(tmp_path / "a") with ( - Session(native_binary, tmp_path / "a", env) as a, + Session(native_binary, tmp_path / "a", env, options=savings_options) as a, Session(native_binary, tmp_path / "b", {}) as b, ): - assert a.observe()["options"] == _claude_mod_launch_options(env) - assert b.observe()["options"] == _claude_mod_launch_options({}) + assert a.observe()["options"] == {**_claude_mod_launch_options(env), **savings_options} + # Claude fills in the declared defaults for options a launch omits. + assert b.observe()["options"] == { + **_claude_mod_launch_options({}), + **CLAUDE_SAVINGS_OPTION_DEFAULTS, + } assert not (a.plugin / "hooks/launch.json").exists() + + +def test_savings_mod_records_the_start_model(native_binary, tmp_path): + root = tmp_path / "session" + options = _savings_options(root) + usage = Path(options["savings_usage"]) + with Session(native_binary, root, {}, options=options) as session: + session.observe() + # The mod writes its sums off the session.start critical path. + deadline = time.monotonic() + 10 + while True: + try: + recorded = json.loads(usage.read_text(encoding="utf-8")) + break + except (OSError, ValueError): + assert time.monotonic() < deadline, "the savings mod never wrote its usage file" + time.sleep(0.1) + + assert recorded == { + "version": 1, + "start_model": "claude-sonnet-4-6", + "main_model": None, + "first_main_model": None, + "user_switched": False, + "entries": [], + } diff --git a/tests/test_claude_mod.py b/tests/test_claude_mod.py index ea1812935..deeabf481 100644 --- a/tests/test_claude_mod.py +++ b/tests/test_claude_mod.py @@ -1,27 +1,40 @@ """Generated plugin contents for the native Claude mod entry point.""" import json +import sys from pathlib import Path import pytest from ucode.constants import AGENT_CLAUDE, SMART_ROUTING_ENV_KEYS -from ucode.smart_routing import config, v2 +from ucode.smart_routing import config, orchestrator, savings, v2 def test_plugin_contains_mod_entry_point(tmp_path): v2._write_routed_claude_plugin(tmp_path, []) assert json.loads((tmp_path / "hooks/hooks.json").read_text()) == {"modules": ["./register.ts"]} - assert (tmp_path / "hooks/register.ts").read_bytes() == ( - Path(v2.__file__).parents[1] / "agents" / "claude_mods" / "register.ts" - ).read_bytes() + sources = Path(v2.__file__).parents[1] / "agents" / "claude_mods" + for name in ("register.ts", "savings.ts"): + assert (tmp_path / "hooks" / name).read_bytes() == (sources / name).read_bytes() assert sorted(path.name for path in (tmp_path / "hooks").iterdir()) == [ "hooks.json", "register.ts", + "savings.ts", ] +def test_plugin_declares_the_savings_options_at_their_defaults(tmp_path): + v2._write_routed_claude_plugin(tmp_path, []) + + user_config = json.loads((tmp_path / ".claude-plugin/plugin.json").read_text())["userConfig"] + assert { + name: user_config[name]["default"] for name in v2.CLAUDE_SAVINGS_OPTION_DEFAULTS + } == v2.CLAUDE_SAVINGS_OPTION_DEFAULTS + # The pricer is an argv list, so paths with spaces reach the mod intact. + assert user_config["savings_pricer"]["multiple"] is True + + @pytest.mark.parametrize("version", [*config._VERSIONS, "", "unsupported_v99"]) def test_launch_options_resolve_versions_without_mutating_input(version): env = { @@ -52,3 +65,54 @@ def test_launch_options_preserve_empty_and_absent_values(): "unset": ["ENABLE_SMART_ROUTING_SUBAGENT_ONLY", "ENABLE_SMART_ROUTER_ORCHESTRATOR"], "recipe": "task_v3", } + + +def test_savings_is_on_by_default_and_opts_out_with_zero(monkeypatch): + assert v2.savings_enabled() + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, "0") + assert not v2.savings_enabled() + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, "1") + assert v2.savings_enabled() + + +@pytest.mark.parametrize("first_prompt", [False, True]) +def test_savings_options_give_the_mod_a_pricer_and_a_launch_local_usage_file( + tmp_path, first_prompt +): + price_cache = tmp_path / "model-prices.json" + + assert v2._claude_savings_launch_options( + tmp_path, price_cache=price_cache, baseline_session_start=first_prompt + ) == { + "savings_pricer": savings.pricer_argv( + python=sys.executable, price_cache=price_cache, baseline_session_start=first_prompt + ), + "savings_usage": str(tmp_path / savings.USAGE_FILENAME), + "orchestrator_version": "", + } + + +def test_opting_out_of_savings_leaves_its_options_at_their_defaults(tmp_path, monkeypatch): + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, "0") + + assert ( + v2._claude_savings_launch_options( + tmp_path, price_cache=tmp_path / "model-prices.json", baseline_session_start=False + ) + == v2.CLAUDE_SAVINGS_OPTION_DEFAULTS + ) + + +@pytest.mark.parametrize("savings_flag", ["1", "0"]) +def test_savings_options_show_the_orchestrator_skill_version_when_it_is_active( + tmp_path, monkeypatch, savings_flag +): + monkeypatch.setenv("ENABLE_SMART_ROUTER_ORCHESTRATOR", "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, savings_flag) + + options = v2._claude_savings_launch_options( + tmp_path, price_cache=tmp_path / "model-prices.json", baseline_session_start=False + ) + + assert options["orchestrator_version"] == orchestrator.skill_version(agent=AGENT_CLAUDE) + assert options["orchestrator_version"] diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 58ab1fda5..7b66a9574 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -7,6 +7,7 @@ import sys import threading import time +from decimal import Decimal from pathlib import Path from unittest.mock import Mock @@ -15,7 +16,8 @@ from ucode.agents import LaunchOptions, claude from ucode.constants import AGENT_CLAUDE from ucode.databricks import AnthropicModelCatalog -from ucode.smart_routing import claude_hooks, routing, v2 +from ucode.smart_routing import claude_hooks, pricing, routing, savings, v2 +from ucode.smart_routing.pricing import ModelPrice if sys.platform != "win32": from ucode.smart_routing import claude_pty @@ -25,6 +27,17 @@ ) +_START_SAVINGS_PRICE_REFRESH = v2._start_savings_price_refresh + + +@pytest.fixture(autouse=True) +def savings_price_refresh(monkeypatch) -> Mock: + """Routed launches refresh savings prices on a background thread; keep it off the network.""" + refresh = Mock() + monkeypatch.setattr(v2, "_start_savings_price_refresh", refresh) + return refresh + + def _plugin_agent_models(plugin_dir: Path) -> set[str]: models = set() for agent_path in (plugin_dir / "agents").glob("*.md"): @@ -430,7 +443,9 @@ def fake_run(_argv, **kwargs): "theme": "dark", } - def test_subagent_only_launch_skips_first_prompt_routing(self, tmp_path, monkeypatch): + def test_subagent_only_launch_skips_first_prompt_routing( + self, tmp_path, monkeypatch, savings_price_refresh + ): user_settings = tmp_path / "settings.json" user_settings.write_text(json.dumps({"model": "opus"})) monkeypatch.delenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, raising=False) @@ -500,9 +515,21 @@ def send_signal(self, _signal): assert exc.value.code == 4 settings = captured["settings"] + price_cache = pricing.price_cache_path(tmp_path, "https://example.com") assert settings["pluginConfigs"][v2.CLAUDE_ROUTING_PLUGIN_NAME] == { - "options": v2._claude_mod_launch_options({**os.environ, **settings["env"]}) + "options": { + **v2._claude_mod_launch_options({**os.environ, **settings["env"]}), + **v2._claude_savings_launch_options( + captured["plugin_dir"].parent, + price_cache=price_cache, + baseline_session_start=False, + ), + } } + # The `opus` alias isn't a model service, so only the catalog id is priced. + savings_price_refresh.assert_called_once_with( + "https://example.com", "token", ["system.ai.claude-opus-4-8"], price_cache + ) env = settings["env"] assert env[v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR] == "1" for key in (v2.SESSION_ENV_VAR, v2.SESSION_PYTHON_ENV_VAR): @@ -524,6 +551,186 @@ def send_signal(self, _signal): assert json.loads(user_settings.read_text()) == {"model": "opus"} +class TestSavingsLaunch: + """Routed launches hand the savings mod its pricer and refresh the rates it prices with.""" + + @staticmethod + def _launch(monkeypatch, tmp_path, *, first_prompt: bool) -> dict: + user_settings = tmp_path / "settings.json" + user_settings.write_text(json.dumps({"model": "opus"})) + routing_flag = ( + v2.ENABLE_SMART_ROUTING_ENV_VAR if first_prompt else v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR + ) + monkeypatch.setenv(routing_flag, "1") + monkeypatch.setattr(v2, "APP_DIR", tmp_path) + monkeypatch.setattr(v2, "CLAUDE_PTY_LOG", tmp_path / "v2.log") + monkeypatch.setattr(v2, "_model_picker_catalog", lambda: None) + monkeypatch.setattr(v2, "get_databricks_token", lambda *_args, **_kwargs: "token") + monkeypatch.setattr(v2, "build_auth_token_argv", lambda *_args, **_kwargs: ["ug"]) + monkeypatch.setattr( + v2, + "list_anthropic_model_catalog", + lambda *_args: AnthropicModelCatalog( + model_ids=["system.ai.claude-opus-4-8"], model_id_to_display_name={} + ), + ) + captured: dict = {} + + def capture(argv) -> None: + captured["launch_dir"] = Path(argv[argv.index("--plugin-dir") + 1]).parent + settings = json.loads(Path(argv[argv.index("--settings") + 1]).read_text()) + captured["options"] = settings["pluginConfigs"][v2.CLAUDE_ROUTING_PLUGIN_NAME][ + "options" + ] + + def run_pty(argv, **_kwargs) -> int: + capture(argv) + return 0 + + def popen(argv, **_kwargs) -> Mock: + capture(argv) + return Mock(wait=lambda: 0) + + if first_prompt: + monkeypatch.setattr(claude_pty, "run_claude_pty", run_pty) + else: + monkeypatch.setattr(v2.subprocess, "Popen", popen) + + with pytest.raises(SystemExit) as exc: + v2.launch_claude( + {"workspace": "https://example.com"}, + [], + binary="claude", + user_settings_path=user_settings, + launch_model="opus", + compose_settings=lambda _args: ({}, []), + launch_model_args=claude._launch_model_args, + model_name=claude._maybe_add_1m_suffix, + ) + assert exc.value.code == 0 + return captured + + @pytest.mark.parametrize("first_prompt", [False, pytest.param(True, marks=requires_pty)]) + def test_first_prompt_routing_prices_against_the_start_model( + self, monkeypatch, tmp_path, first_prompt + ): + captured = self._launch(monkeypatch, tmp_path, first_prompt=first_prompt) + + assert captured["options"]["savings_pricer"] == savings.pricer_argv( + python=sys.executable, + price_cache=pricing.price_cache_path(tmp_path, "https://example.com"), + baseline_session_start=first_prompt, + ) + assert captured["options"]["savings_usage"] == str( + captured["launch_dir"] / savings.USAGE_FILENAME + ) + + def test_opting_out_skips_the_pricer_and_the_price_refresh( + self, monkeypatch, tmp_path, savings_price_refresh + ): + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, "0") + + captured = self._launch(monkeypatch, tmp_path, first_prompt=False) + + assert { + key: captured["options"][key] for key in v2.CLAUDE_SAVINGS_OPTION_DEFAULTS + } == v2.CLAUDE_SAVINGS_OPTION_DEFAULTS + savings_price_refresh.assert_not_called() + + +class TestSavingsPriceRefresh: + def test_prices_routable_and_main_models_by_system_ai_name(self): + services = v2._savings_model_services( + [ + "system.ai.claude-opus-4-8[1m]", + "anthropic-aigw-73ea02b2-system.ai.glm-5-2", + "databricks-claude-sonnet-5", + "opus", + None, + ], + { + "env": { + "ANTHROPIC_DEFAULT_OPUS_MODEL": "system.ai.claude-opus-5-5", + "UNRELATED": "system.ai.not-a-main-model", + } + }, + ) + + assert services == [ + "system.ai.claude-opus-4-8", + "system.ai.claude-opus-5-5", + "system.ai.claude-sonnet-5", + "system.ai.glm-5-2", + ] + + def test_caches_the_fetched_rates(self, monkeypatch, tmp_path): + cache = tmp_path / "model-prices.json" + requests = [] + + def fetch(workspace, token, model_services): + requests.append((workspace, token, model_services)) + token_costs = [ + {"token_type": "TOKEN_TYPE_INPUT", "cost": 5}, + {"token_type": "TOKEN_TYPE_OUTPUT", "cost": 25}, + ] + rate = { + "model_service": "system.ai.claude-opus-4-8", + "costs": [{"unit": "USD", "token_costs": token_costs}], + } + return [rate], None + + monkeypatch.setattr(v2, "fetch_endpoint_rates", fetch) + + v2._refresh_savings_prices( + "https://example.com", "token", ["system.ai.claude-opus-4-8"], cache + ) + + assert requests == [("https://example.com", "token", ["system.ai.claude-opus-4-8"])] + cached = pricing.read_price_cache(cache) + assert cached is not None + assert cached[0] == { + "claude-opus-4-8": ModelPrice(input=Decimal("5"), output=Decimal("25")) + } + + @pytest.mark.parametrize( + "fetch", + [ + lambda *_args: ([], "HTTP 404 Not Found"), + lambda *_args: (_ for _ in ()).throw(RuntimeError("rates API down")), + ], + ) + def test_failed_refresh_keeps_the_previous_cache(self, monkeypatch, tmp_path, fetch): + cache = tmp_path / "model-prices.json" + pricing.write_price_cache( + cache, {"claude-opus-4-8": ModelPrice(input=Decimal("5"))}, now=1.0 + ) + monkeypatch.setattr(v2, "fetch_endpoint_rates", fetch) + + v2._refresh_savings_prices( + "https://example.com", "token", ["system.ai.claude-opus-4-8"], cache + ) + + cached = pricing.read_price_cache(cache) + assert cached is not None and cached[1] == "1.0" + + def test_refresh_runs_on_a_daemon_thread(self, monkeypatch, tmp_path): + done = threading.Event() + calls = [] + + def refresh(*args): + thread = threading.current_thread() + calls.append((thread is threading.main_thread(), thread.daemon, args)) + done.set() + + monkeypatch.setattr(v2, "_refresh_savings_prices", refresh) + cache = tmp_path / "model-prices.json" + + _START_SAVINGS_PRICE_REFRESH("https://example.com", "token", ["m"], cache) + + assert done.wait(5) + assert calls == [(False, True, ("https://example.com", "token", ["m"], cache))] + + class TestV2ModelPickerDiscovery: """modelPicker takes priority over gateway model discovery for smart routing.""" diff --git a/tests/test_claude_windows_smart_routing.py b/tests/test_claude_windows_smart_routing.py index b6db0f462..bb5266f0f 100644 --- a/tests/test_claude_windows_smart_routing.py +++ b/tests/test_claude_windows_smart_routing.py @@ -54,6 +54,7 @@ def test_windows_subagent_routing_uses_native_binary_without_unix_imports(tmp_pa monkeypatch.delenv(v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR, raising=False) monkeypatch.setattr(v2, "APP_DIR", tmp_path) monkeypatch.setattr(v2, "install_skill", lambda *_args, **_kwargs: None) + monkeypatch.setattr(v2, "_start_savings_price_refresh", lambda *_args, **_kwargs: None) monkeypatch.setattr(v2, "_model_picker_catalog", lambda: None) monkeypatch.setattr(v2, "get_databricks_token", lambda *_args, **_kwargs: "token") monkeypatch.setattr(v2, "build_auth_token_argv", lambda *_args, **_kwargs: ["ug"]) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 49c88faa1..2402badf2 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -5453,3 +5453,58 @@ def test_binary_mode_is_unchanged(self, monkeypatch): script = r"import sys; sys.stdout.buffer.write(b'\xff')" result = db_mod.run([sys.executable, "-c", script], capture_output=True) assert result.stdout == b"\xff" + + +class TestFetchEndpointRates: + def test_requests_system_ai_services_as_query_params(self, monkeypatch): + seen = {} + glm_rate = { + "model_service": "system.ai.glm-5-3", + "costs": [ + {"unit": "USD", "token_costs": [{"token_type": "TOKEN_TYPE_INPUT", "cost": 1.4}]} + ], + } + + def fake_get(url, token, **kwargs): + seen.update(url=url, token=token) + return {"model_service_rates": [glm_rate, "not-a-rate"]}, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + rates, reason = db_mod.fetch_endpoint_rates( + WS, "tok", ["system.ai.glm-5-3", "claude-opus-4-8", "system.ai.claude-sonnet-5"] + ) + + assert reason is None + assert rates == [glm_rate] + # Names outside system.ai are dropped (the API rejects them); the rest travel as repeated, + # sorted `databricks_hosted_model_services` query parameters. + assert seen["token"] == "tok" + assert seen["url"] == ( + f"{WS}/api/ai-gateway/v2/endpoint-rates:batchGet" + "?databricks_hosted_model_services=system.ai.claude-sonnet-5" + "&databricks_hosted_model_services=system.ai.glm-5-3" + ) + + def test_skips_the_request_without_system_ai_services(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_http_get_json", lambda *args, **kwargs: pytest.fail("no request expected") + ) + + assert db_mod.fetch_endpoint_rates(WS, "tok", ["claude-opus-4-8"]) == ( + [], + "no system.ai model services to price", + ) + + @pytest.mark.parametrize( + ("response", "expected"), + [ + ((None, "HTTP 404 Not Found"), ([], "HTTP 404 Not Found")), + ((["unexpected"], None), ([], "endpoint-rates returned an unexpected response shape")), + (({}, None), ([], None)), + ], + ) + def test_reports_failures_without_raising(self, monkeypatch, response, expected): + monkeypatch.setattr(db_mod, "_http_get_json", lambda *args, **kwargs: response) + + assert db_mod.fetch_endpoint_rates(WS, "tok", ["system.ai.glm-5-3"]) == expected diff --git a/tests/test_skills.py b/tests/test_skills.py index a55737bf0..38e741e40 100644 --- a/tests/test_skills.py +++ b/tests/test_skills.py @@ -1,5 +1,6 @@ """Tests for skills shipped with Unity Gateway.""" +import re import shlex from pathlib import Path @@ -108,6 +109,47 @@ def test_orchestrator_injects_the_installed_agent_instructions( assert context.endswith(installed.joinpath("SKILL.md").read_text()) +def test_orchestrator_skill_version_reads_the_shipped_skill(): + assert re.fullmatch(r"\d+\.\d+\.\d+", orchestrator.skill_version(agent="claude") or "") + + +@pytest.mark.parametrize( + ("text", "expected"), + [ + ('---\nname: x\nmetadata:\n version: "1.2.0"\n---\nbody', "1.2.0"), + ("---\nmetadata:\n owner: ug\n version: 2.0\n---\n", "2.0"), + # Only metadata's version counts, and only inside the frontmatter. + ("---\nversion: 9\nmetadata:\n owner: ug\n---\nmetadata:\n version: 9\n", None), + ('---\nmetadata:\n version: ""\n---\n', None), + ("no frontmatter", None), + ], +) +def test_orchestrator_skill_version_is_the_frontmatter_metadata_version( + bundled_skill, text, expected +): + source_skill, _home = bundled_skill + _write_skill(source_skill.parent, skills.SMART_ROUTER_ORCHESTRATOR_SKILL, text) + + assert orchestrator.skill_version(agent="claude") == expected + + +def test_orchestrator_skill_version_follows_the_agent_entrypoint(bundled_skill): + source_skill, _home = bundled_skill + source = _write_skill( + source_skill.parent, + skills.SMART_ROUTER_ORCHESTRATOR_SKILL, + "---\nmetadata:\n version: 1.0.0\n---\n", + ) + source.joinpath("SKILL.claude.md").write_text("---\nmetadata:\n version: 2.0.0\n---\n") + + assert orchestrator.skill_version(agent="claude") == "2.0.0" + assert orchestrator.skill_version(agent="codex") == "1.0.0" + + +def test_orchestrator_skill_version_is_none_without_the_skill(bundled_skill): + assert orchestrator.skill_version(agent="claude") is None + + def test_uninstalls_one_skill_from_agent_roots(bundled_skill): source_skill, home = bundled_skill _write_skill(source_skill.parent, "second-skill") diff --git a/tests/test_smart_routing_pricing.py b/tests/test_smart_routing_pricing.py new file mode 100644 index 000000000..90bdd561d --- /dev/null +++ b/tests/test_smart_routing_pricing.py @@ -0,0 +1,300 @@ +"""Tests for per-token model pricing used by the smart-routing savings estimate.""" + +from __future__ import annotations + +import json +from decimal import Decimal + +import pytest + +from ucode.smart_routing import pricing +from ucode.smart_routing.pricing import ModelPrice, TokenUsage + +# Catalog list prices (USD per million tokens) for the two Claude route arms. +OPUS = ModelPrice( + input=Decimal("5"), + output=Decimal("25"), + cache_read=Decimal("0.5"), + cache_write_5m=Decimal("6.25"), + cache_write_1h=Decimal("10"), +) +SONNET_4 = ModelPrice( + input=Decimal("3"), + output=Decimal("15"), + cache_read=Decimal("0.3"), + cache_write_5m=Decimal("3.75"), + long_context_threshold=200_000, + long_context=ModelPrice( + input=Decimal("6"), + output=Decimal("22.5"), + cache_read=Decimal("0.6"), + cache_write_5m=Decimal("7.5"), + ), +) + + +class TestModelKey: + @pytest.mark.parametrize( + "model", + [ + "claude-opus-4-8", + "system.ai.claude-opus-4-8", + "system.ai.claude-opus-4-8[1m]", + "databricks-claude-opus-4-8", + "anthropic-aigw-1a2b3c4d-claude-opus-4-8", + "SYSTEM.AI.Claude-Opus-4-8", + "anthropic.claude-opus-4-8", + "global.anthropic.claude-opus-4-8", + "us.anthropic.claude-opus-4-8-20260101-v1:0", + "claude-opus-4-8@default", + ], + ) + def test_collapses_gateway_spellings(self, model): + assert pricing.model_key(model) == "claude-opus-4-8" + + def test_matches_a_served_bedrock_id_to_the_requested_model(self): + # The gateway reports Haiku responses by their Bedrock id. + assert pricing.model_key("anthropic.claude-haiku-4-5-20251001-v1:0") == pricing.model_key( + "system.ai.claude-haiku-4-5" + ) + + def test_unwraps_non_claude_gateway_models(self): + assert pricing.model_key("anthropic-aigw-73ea02b2-system.ai.glm-5-2") == "glm-5-2" + + @pytest.mark.parametrize( + ("served", "requested"), + [ + ("glm-5.3-flash", "system.ai.glm-5-3-flash"), + ("glm-5-3-colo-on-sp-v1", "system.ai.glm-5-3"), + ("kimi-k3-colo-on-sp-v1", "system.ai.kimi-k3"), + ], + ) + def test_matches_a_served_deployment_id_to_the_requested_model(self, served, requested): + assert pricing.model_key(served) == pricing.model_key(requested) + + def test_keeps_distinct_models_distinct(self): + assert pricing.model_key("system.ai.claude-sonnet-5") != pricing.model_key( + "system.ai.claude-opus-4-8" + ) + + +class TestTokenUsage: + def test_splits_cache_writes_by_ttl(self): + tokens = TokenUsage.from_message_usage( + { + "input_tokens": 2, + "cache_creation_input_tokens": 6509, + "cache_read_input_tokens": 143654, + "output_tokens": 695, + "cache_creation": { + "ephemeral_1h_input_tokens": 6000, + "ephemeral_5m_input_tokens": 9, + }, + } + ) + + # The 500 writes the breakdown doesn't cover bill at the default 5-minute TTL. + assert tokens == TokenUsage( + input=2, cache_write_5m=509, cache_write_1h=6000, cache_read=143654, output=695 + ) + assert tokens.prompt_tokens == 2 + 6509 + 143654 + + def test_treats_missing_breakdown_as_five_minute_writes(self): + tokens = TokenUsage.from_message_usage({"cache_creation_input_tokens": 40}) + + assert tokens == TokenUsage(cache_write_5m=40) + + def test_can_bill_writes_the_breakdown_does_not_cover_as_one_hour_writes(self): + unsplit = {"cache_creation_input_tokens": 40} + partial = { + "cache_creation_input_tokens": 6509, + "cache_creation": {"ephemeral_1h_input_tokens": 6000, "ephemeral_5m_input_tokens": 9}, + } + + assert TokenUsage.from_message_usage(unsplit, uncovered_writes_1h=True) == TokenUsage( + cache_write_1h=40 + ) + # The breakdown's own 5-minute writes stay 5-minute; only the 500 uncovered move to 1-hour. + assert TokenUsage.from_message_usage(partial, uncovered_writes_1h=True) == TokenUsage( + cache_write_5m=9, cache_write_1h=6500 + ) + + @pytest.mark.parametrize("bad", [None, -5, True, "12", 1.5]) + def test_ignores_non_count_values(self, bad): + assert TokenUsage.from_message_usage({"input_tokens": bad, "output_tokens": 3}) == ( + TokenUsage(output=3) + ) + + +class TestTokenCost: + def test_prices_every_token_class_at_its_own_rate(self): + tokens = TokenUsage( + input=1_000, cache_write_5m=2_000, cache_write_1h=3_000, cache_read=100_000, output=500 + ) + + # 1k*5 + 2k*6.25 + 3k*10 + 100k*0.5 + 500*25 = 110,000 micro-dollars. + assert pricing.token_cost(OPUS, tokens) == Decimal("0.11") + + def test_uses_long_context_rates_above_the_threshold(self): + short = TokenUsage(input=100_000, output=1_000) + long = TokenUsage(input=150_000, cache_read=60_000, output=1_000) + + assert pricing.token_cost(SONNET_4, short) == Decimal("0.315") + assert pricing.token_cost(SONNET_4, long) == Decimal("0.9585") + + def test_returns_none_when_a_used_class_has_no_rate(self): + # SONNET_4 publishes no 1-hour write rate. + assert pricing.token_cost(SONNET_4, TokenUsage(cache_write_1h=10)) is None + + def test_missing_rate_for_an_unused_class_is_fine(self): + assert pricing.token_cost(SONNET_4, TokenUsage(input=100_000)) == Decimal("0.3") + + +class TestPriceCache: + def test_round_trips_prices_under_model_keys(self, tmp_path): + path = tmp_path / "cache" / "model-prices.json" + pricing.write_price_cache( + path, {"system.ai.claude-opus-4-8": OPUS, "claude-sonnet-4": SONNET_4}, now=100.0 + ) + + cached = pricing.read_price_cache(path) + + assert cached is not None + prices, fingerprint = cached + assert prices == {"claude-opus-4-8": OPUS, "claude-sonnet-4": SONNET_4} + assert fingerprint == "100.0" + assert [entry.name for entry in path.parent.iterdir()] == ["model-prices.json"] + + def test_fingerprint_changes_on_refresh(self, tmp_path): + path = tmp_path / "model-prices.json" + pricing.write_price_cache(path, {"claude-opus-4-8": OPUS}, now=100.0) + first = pricing.read_price_cache(path) + pricing.write_price_cache(path, {"claude-opus-4-8": OPUS}, now=200.0) + + assert first is not None and pricing.read_price_cache(path) is not None + assert first[1] != pricing.read_price_cache(path)[1] + + @pytest.mark.parametrize( + "content", + [ + "not json", + json.dumps({"version": 999, "models": {"claude-opus-4-8": {"input": "5"}}}), + json.dumps({"version": 1, "models": {}}), + json.dumps({"version": 1, "models": []}), + ], + ) + def test_unusable_cache_reads_as_absent(self, tmp_path, content): + path = tmp_path / "model-prices.json" + path.write_text(content) + + assert pricing.read_price_cache(path) is None + + def test_missing_cache_reads_as_absent(self, tmp_path): + assert pricing.read_price_cache(tmp_path / "missing.json") is None + + +class TestEndpointRates: + @staticmethod + def _rate(service: str, *, usd: dict | None = None, dbu: dict | None = None) -> dict: + """An ``EndpointRate`` entry: costs grouped by unit, each a token_type->cost (per 1M).""" + costs = [] + for unit, table in (("DBU", dbu), ("USD", usd)): + if table is not None: + costs.append( + { + "unit": unit, + "token_costs": [ + {"token_type": token_type, "cost": cost} + for token_type, cost in table.items() + ], + } + ) + return {"model_service": service, "costs": costs} + + def test_reads_dollar_rates_keyed_by_model_service(self): + rates = [ + self._rate( + "system.ai.claude-fable-5", + dbu={"TOKEN_TYPE_INPUT": 142.858}, # priced in USD; the DBU group is ignored + usd={"TOKEN_TYPE_INPUT": 10.00006, "TOKEN_TYPE_OUTPUT": 50.00002}, + ), + self._rate( + "system.ai.glm-5-3", + usd={"TOKEN_TYPE_INPUT": 1.4, "TOKEN_TYPE_OUTPUT": 4.3999998}, + ), + ] + + assert pricing.prices_from_endpoint_rates(rates) == { + "system.ai.claude-fable-5": ModelPrice( + input=Decimal("10.00006"), output=Decimal("50.00002") + ), + "system.ai.glm-5-3": ModelPrice(input=Decimal("1.4"), output=Decimal("4.3999998")), + } + + def test_reads_all_token_classes_including_cache(self): + rates = [ + self._rate( + "system.ai.claude-opus-4-8", + usd={ + "TOKEN_TYPE_INPUT": 5, + "TOKEN_TYPE_OUTPUT": 25, + "TOKEN_TYPE_CACHE_READ": 0.5, + "TOKEN_TYPE_CACHE_CREATION": 6.25, + "TOKEN_TYPE_CACHE_CREATION_1H": 10, + }, + ) + ] + + assert pricing.prices_from_endpoint_rates(rates) == { + "system.ai.claude-opus-4-8": ModelPrice( + input=Decimal("5"), + output=Decimal("25"), + cache_read=Decimal("0.5"), + cache_write_5m=Decimal("6.25"), + cache_write_1h=Decimal("10"), + ) + } + + def test_skips_models_without_a_dollar_group(self): + rates = [ + # No USD group: the org has no DBU-to-dollar conversion, so we can't price it. + self._rate("system.ai.glm-5-3", dbu={"TOKEN_TYPE_INPUT": 20}), + self._rate("system.ai.kimi-k3", usd={}), # dollar group present but empty + {"model_service": "system.ai.no-costs"}, # no costs field + { + "costs": [ + {"unit": "USD", "token_costs": [{"token_type": "TOKEN_TYPE_INPUT", "cost": 1}]} + ] + }, + "not-a-rate", + ] + + assert pricing.prices_from_endpoint_rates(rates) == {} + + def test_endpoint_prices_survive_the_cache(self, tmp_path): + path = tmp_path / "model-prices.json" + prices = pricing.prices_from_endpoint_rates( + [ + self._rate( + "system.ai.claude-haiku-4-5", + usd={"TOKEN_TYPE_INPUT": 1, "TOKEN_TYPE_OUTPUT": 5}, + ) + ] + ) + pricing.write_price_cache(path, prices) + + cached = pricing.read_price_cache(path) + + assert cached is not None + served = pricing.model_key("anthropic.claude-haiku-4-5-20251001-v1:0") + assert cached[0][served] == ModelPrice(input=Decimal("1"), output=Decimal("5")) + + +class TestPriceCachePath: + def test_is_stable_per_workspace_and_distinct_across_workspaces(self, tmp_path): + first = pricing.price_cache_path(tmp_path, "https://A.cloud.databricks.com/") + + assert first == pricing.price_cache_path(tmp_path, "https://a.cloud.databricks.com") + assert first != pricing.price_cache_path(tmp_path, "https://b.cloud.databricks.com") + assert first.parent == tmp_path + assert first.name.startswith("model-prices-") and first.suffix == ".json" diff --git a/tests/test_smart_routing_savings.py b/tests/test_smart_routing_savings.py new file mode 100644 index 000000000..726ba93b0 --- /dev/null +++ b/tests/test_smart_routing_savings.py @@ -0,0 +1,324 @@ +"""Tests for pricing the Claude mod's token sums into the smart-routing savings estimate.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from decimal import Decimal +from pathlib import Path + +import pytest + +from ucode.smart_routing import pricing, savings +from ucode.smart_routing.pricing import ModelPrice + +OPUS_ID = "system.ai.claude-opus-4-8" +SONNET_ID = "system.ai.claude-sonnet-5" +PRICES = { + OPUS_ID: ModelPrice( + input=Decimal("5"), + output=Decimal("25"), + cache_read=Decimal("0.5"), + cache_write_5m=Decimal("6.25"), + cache_write_1h=Decimal("10"), + ), + SONNET_ID: ModelPrice( + input=Decimal("2"), + output=Decimal("10"), + cache_read=Decimal("0.2"), + cache_write_5m=Decimal("2.5"), + cache_write_1h=Decimal("4"), + ), +} +# Main-agent request: $0.07505 on Opus (10*5 + 100k*0.5 + 1k*25), $0.03002 on Sonnet. +MAIN_USAGE = {"input_tokens": 10, "cache_read_input_tokens": 100_000, "output_tokens": 1_000} +# Subagent request: $0.122 on Sonnet (1k*2 + 20k*4 + 4k*10), $0.305 at Opus rates. +SUBAGENT_USAGE = { + "input_tokens": 1_000, + "cache_creation_input_tokens": 20_000, + "cache_creation": {"ephemeral_1h_input_tokens": 20_000, "ephemeral_5m_input_tokens": 0}, + "output_tokens": 4_000, +} + + +@pytest.fixture +def price_cache(tmp_path) -> Path: + path = tmp_path / "model-prices.json" + pricing.write_price_cache(path, PRICES, now=1.0) + return path + + +def entry( + agent: str, baseline: str | None, served: str, usage: dict, *, before_user_switch: bool = False +) -> dict: + found = { + "agent": agent, + "served": served, + "before_user_switch": before_user_switch, + "usage": usage, + } + return found if baseline is None else {**found, "baseline": baseline} + + +def write_usage(tmp_path: Path, entries: list, *, start_model: str | None = OPUS_ID) -> Path: + path = tmp_path / savings.USAGE_FILENAME + path.write_text( + json.dumps({"version": 1, "start_model": start_model, "entries": entries}), + encoding="utf-8", + ) + return path + + +def priced(usage: Path, price_cache: Path, *, first_prompt: bool = False) -> str | None: + line = savings.usage_output(usage, price_cache=price_cache, baseline_session_start=first_prompt) + assert "\n" not in line + return json.loads(line)["savings"] + + +class TestUsageSavings: + def test_prices_tokens_against_each_entrys_baseline(self, tmp_path, price_cache): + entries = [ + entry("main", OPUS_ID, OPUS_ID, MAIN_USAGE), + entry("a1", OPUS_ID, SONNET_ID, SUBAGENT_USAGE), + ] + + # Baseline $0.38005 - actual $0.19705. + assert priced(write_usage(tmp_path, entries), price_cache) == ( + "๐Ÿ’ฐ Est. saved with smart routing: $0.18 (48%)" + ) + + def test_an_entrys_baseline_is_the_main_model_when_its_request_ran(self, tmp_path, price_cache): + entries = [ + entry("main", OPUS_ID, OPUS_ID, MAIN_USAGE), + entry("a1", OPUS_ID, SONNET_ID, SUBAGENT_USAGE), + # After the user's `/model sonnet`, subagents on Sonnet saved nothing. + entry("main", SONNET_ID, SONNET_ID, MAIN_USAGE), + entry("a1", SONNET_ID, SONNET_ID, SUBAGENT_USAGE), + ] + + assert priced(write_usage(tmp_path, entries), price_cache) == ( + "๐Ÿ’ฐ Est. saved with smart routing: $0.18 (34%)" + ) + + def test_shows_a_net_cost_increase_honestly(self, tmp_path, price_cache): + entries = [ + entry("main", SONNET_ID, SONNET_ID, MAIN_USAGE), + entry("a1", SONNET_ID, OPUS_ID, SUBAGENT_USAGE), + ] + + # Baseline $0.15202 vs actual $0.33502. + assert priced(write_usage(tmp_path, entries), price_cache) == ( + "Smart routing cost $0.18 more (120%)" + ) + + def test_nothing_to_claim_until_a_model_was_swapped(self, tmp_path, price_cache): + entries = [entry("main", OPUS_ID, OPUS_ID, MAIN_USAGE)] + + assert priced(write_usage(tmp_path, entries), price_cache) is None + + def test_bills_entries_without_the_cache_write_split_as_one_hour_writes( + self, tmp_path, price_cache + ): + usage = { + "input_tokens": 1_000, + "cache_creation_input_tokens": 20_000, + "output_tokens": 4_000, + } + + # ug launches Claude with 1-hour caching, so unsplit writes bill at the 1-hour rate, as + # SUBAGENT_USAGE does: $0.122 on Sonnet vs $0.305 at Opus rates. + assert priced( + write_usage(tmp_path, [entry("a1", OPUS_ID, SONNET_ID, usage)]), price_cache + ) == ("๐Ÿ’ฐ Est. saved with smart routing: $0.18 (60%)") + + def test_honors_an_explicit_cache_write_split(self, tmp_path, price_cache): + usage = { + "input_tokens": 1_000, + "cache_creation_input_tokens": 20_000, + "cache_creation": {"ephemeral_1h_input_tokens": 0, "ephemeral_5m_input_tokens": 20_000}, + "output_tokens": 4_000, + } + + # The split says 5-minute writes: $0.092 on Sonnet vs $0.23 at Opus rates. + assert priced( + write_usage(tmp_path, [entry("a1", OPUS_ID, SONNET_ID, usage)]), price_cache + ) == ("๐Ÿ’ฐ Est. saved with smart routing: $0.14 (60%)") + + def test_skips_entries_without_tokens(self, tmp_path, price_cache): + entries = [ + entry("a1", OPUS_ID, SONNET_ID, SUBAGENT_USAGE), + entry("a2", "unpriced-baseline", "unpriced-served", {"input_tokens": 0}), + ] + + assert priced(write_usage(tmp_path, entries), price_cache) == ( + "๐Ÿ’ฐ Est. saved with smart routing: $0.18 (60%)" + ) + + def test_first_prompt_routing_prices_against_the_start_model(self, tmp_path, price_cache): + # Under first-prompt routing the router picks the main model too, so an entry's own + # baseline (the routed model) would hide the saving; until the user switches the main + # model its presence or absence is ignored. + entries = [ + entry("main", SONNET_ID, SONNET_ID, MAIN_USAGE, before_user_switch=True), + entry("a1", None, SONNET_ID, SUBAGENT_USAGE, before_user_switch=True), + ] + usage = write_usage(tmp_path, entries, start_model=OPUS_ID) + + # Baseline $0.38005 at Opus vs actual $0.15202 on Sonnet. + assert priced(usage, price_cache, first_prompt=True) == ( + "๐Ÿ’ฐ Est. saved with smart routing: $0.23 (60%)" + ) + assert priced(usage, price_cache) is None + + def test_first_prompt_routing_prices_entries_after_a_user_switch_as_their_own( + self, tmp_path, price_cache + ): + entries = [ + entry("main", SONNET_ID, SONNET_ID, MAIN_USAGE, before_user_switch=True), + entry("a1", SONNET_ID, SONNET_ID, SUBAGENT_USAGE, before_user_switch=True), + # The user's `/model opus`, then `/model sonnet`. + entry("main", OPUS_ID, OPUS_ID, MAIN_USAGE), + entry("main", SONNET_ID, SONNET_ID, MAIN_USAGE), + ] + usage = write_usage(tmp_path, entries, start_model=OPUS_ID) + + # Baseline $0.48512 - actual $0.25709. + assert priced(usage, price_cache, first_prompt=True) == ( + "๐Ÿ’ฐ Est. saved with smart routing: $0.23 (47%)" + ) + + @pytest.mark.parametrize("start_model", [None, ""]) + def test_first_prompt_routing_without_a_start_model_has_no_estimate( + self, tmp_path, price_cache, start_model + ): + entries = [entry("a1", OPUS_ID, SONNET_ID, SUBAGENT_USAGE)] + usage = write_usage(tmp_path, entries, start_model=start_model) + + assert priced(usage, price_cache, first_prompt=True) is None + + @pytest.mark.parametrize( + ("baseline", "served"), + [ + # An unpriced served model, or an unpriced baseline, would undercount. + (OPUS_ID, "system.ai.claude-mystery-1"), + ("system.ai.claude-mystery-1", SONNET_ID), + # Baseline-less or model-less entries. + (None, SONNET_ID), + ("", SONNET_ID), + (OPUS_ID, ""), + ], + ) + def test_an_unpriced_or_baseline_less_entry_voids_the_estimate( + self, tmp_path, price_cache, baseline, served + ): + entries = [ + entry("a0", OPUS_ID, SONNET_ID, SUBAGENT_USAGE), + entry("a1", baseline, served, SUBAGENT_USAGE), + ] + + assert priced(write_usage(tmp_path, entries), price_cache) is None + + @pytest.mark.parametrize( + "bad", ["nope", 7, {"baseline": OPUS_ID, "served": SONNET_ID}, {"baseline": OPUS_ID}] + ) + def test_a_malformed_entry_voids_the_estimate(self, tmp_path, price_cache, bad): + entries = [entry("a0", OPUS_ID, SONNET_ID, SUBAGENT_USAGE), bad] + + assert priced(write_usage(tmp_path, entries), price_cache) is None + + def test_without_prices_there_is_no_estimate(self, tmp_path, price_cache): + price_cache.unlink() + entries = [entry("a1", OPUS_ID, SONNET_ID, SUBAGENT_USAGE)] + + assert priced(write_usage(tmp_path, entries), price_cache) is None + + @pytest.mark.parametrize( + "text", + [None, "not json", "[]", json.dumps({"version": 2, "entries": []}), '{"version": 1}'], + ) + def test_an_unusable_file_has_no_estimate(self, tmp_path, price_cache, text): + usage = tmp_path / savings.USAGE_FILENAME + if text is not None: + usage.write_text(text, encoding="utf-8") + + assert priced(usage, price_cache) is None + + def test_a_file_in_another_format_is_an_error_not_an_empty_estimate( + self, tmp_path, price_cache + ): + usage = tmp_path / savings.USAGE_FILENAME + usage.write_text(json.dumps({"version": 2, "entries": []}), encoding="utf-8") + + with pytest.raises(ValueError, match="unsupported savings usage file"): + savings.usage_savings(usage, price_cache=price_cache, baseline_session_start=False) + + +class TestSavingsText: + def test_rounds_to_cents_and_whole_percent(self): + assert ( + savings.savings_text(Decimal("1234.565"), Decimal("2469.13")) + == "๐Ÿ’ฐ Est. saved with smart routing: $1,234.57 (50%)" + ) + + def test_negative_reads_as_a_cost_increase(self): + assert ( + savings.savings_text(Decimal("-0.04"), Decimal("0.40")) + == "Smart routing cost $0.04 more (10%)" + ) + + def test_sub_cent_amounts_are_not_shown_as_zero(self): + assert ( + savings.savings_text(Decimal("0.004"), Decimal("1")) + == "๐Ÿ’ฐ Est. saved with smart routing: <$0.01 (0%)" + ) + + +class TestPricerCommand: + @pytest.mark.parametrize("first_prompt", [False, True]) + def test_runs_the_module_on_the_given_python(self, tmp_path, first_prompt): + argv = savings.pricer_argv( + python="/py", price_cache=tmp_path / "prices.json", baseline_session_start=first_prompt + ) + + assert argv == [ + "/py", + "-P", + "-m", + savings.MODULE, + "--price-cache", + str(tmp_path / "prices.json"), + *(["--baseline-session-start"] if first_prompt else []), + ] + + @staticmethod + def run(price_cache: Path, usage: Path, *, first_prompt: bool = False): + argv = savings.pricer_argv( + python=sys.executable, price_cache=price_cache, baseline_session_start=first_prompt + ) + # An ASCII stdout locale must not cost the money bag. + return subprocess.run( + [*argv, "--usage", str(usage)], + capture_output=True, + env={**os.environ, "PYTHONIOENCODING": "ascii"}, + timeout=30, + check=False, + ) + + def test_prints_exactly_one_utf8_json_line(self, tmp_path, price_cache): + usage = write_usage(tmp_path, [entry("a1", OPUS_ID, SONNET_ID, SUBAGENT_USAGE)]) + + result = self.run(price_cache, usage) + + assert result.returncode == 0, result.stderr + assert result.stdout.count(b"\n") == 1 + assert json.loads(result.stdout.decode("utf-8")) == { + "savings": "๐Ÿ’ฐ Est. saved with smart routing: $0.18 (60%)" + } + + def test_never_fails_the_mod(self, tmp_path, price_cache): + result = self.run(price_cache, tmp_path / "missing.json", first_prompt=True) + + assert result.returncode == 0, result.stderr + assert json.loads(result.stdout) == {"savings": None} From 0ce99eea46c1b3e0e6bacd1d9bb8900e94f29617 Mon Sep 17 00:00:00 2001 From: "lilly.luo" Date: Tue, 13 Oct 2026 22:17:14 +0000 Subject: [PATCH 2/7] Gate smart routing savings behind explicit opt-in --- README.md | 36 ++- docs/claude-mod.md | 19 +- skills/smart-router/SKILL.md | 5 +- src/ucode/agents/claude_mods/register.ts | 277 ++++++++++++++++++++++- src/ucode/agents/claude_mods/savings.ts | 274 ---------------------- src/ucode/constants.py | 2 + src/ucode/smart_routing/config.py | 20 ++ src/ucode/smart_routing/savings.py | 16 +- src/ucode/smart_routing/v2.py | 27 +-- tests/README.md | 9 +- tests/conftest.py | 2 +- tests/integration/README.md | 6 + tests/native_claude/test_mod.py | 27 ++- tests/test_claude_mod.py | 41 ++-- tests/test_claude_smart_routing_v2.py | 43 +++- tests/test_smart_router.py | 11 +- tests/test_smart_routing_config.py | 45 ++++ tests/test_smart_routing_savings.py | 15 +- 18 files changed, 526 insertions(+), 349 deletions(-) delete mode 100644 src/ucode/agents/claude_mods/savings.ts diff --git a/README.md b/README.md index 4cc9d2f53..00c379d3f 100644 --- a/README.md +++ b/README.md @@ -269,14 +269,15 @@ Use `SMART_ROUTER_CONFIG_VERSION` at launch to select a smart-routing configurat Each version defines separate Claude Code and Codex flag mappings; the existing versions use identical mappings for both agents. The Claude-only version disables routing for Codex: -| Version | Routed agents | Subagent routing | First-prompt routing | Orchestrator | -| --- | --- | --- | --- | --- | -| `first_prompt_and_subagent_no_orch_v0` | Claude, Codex | On | On | Off | -| `subagent_only_v0` | Claude, Codex | On | Off | Off | -| `subagent_only_v1` | Claude, Codex | On | Off | Off | -| `subagent_orch_v0` | Claude, Codex | On | Off | On | -| `subagent_orch_v1` | Claude, Codex | On | Off | On | -| `subagent_orch_v0_claude_only` | Claude | On | Off | On | +| Version | Routed agents | Subagent routing | First-prompt routing | Orchestrator | Claude savings | +| --- | --- | --- | --- | --- | --- | +| `first_prompt_and_subagent_no_orch_v0` | Claude, Codex | On | On | Off | Off | +| `subagent_only_v0` | Claude, Codex | On | Off | Off | Off | +| `subagent_only_v1` | Claude, Codex | On | Off | Off | Off | +| `subagent_orch_v0` | Claude, Codex | On | Off | On | Off | +| `subagent_orch_v1` | Claude, Codex | On | Off | On | Off | +| `subagent_orch_v0_claude_only` | Claude | On | Off | On | Off | +| `subagent_orch_savings_v0` | Claude, Codex | On | Off | On | On | `first_prompt_and_subagent_no_orch_v0` is the customer configuration for first-prompt and subagent routing without orchestration: `ENABLE_SMART_ROUTING_V2=1`, @@ -294,10 +295,21 @@ and sets all three flags to `"0"` for Codex. Codex uses its ordinary launch and managed default model, without routing or orchestration. Explicit launch/session on/off controls retain their existing precedence over presets. +`subagent_orch_savings_v0` uses the routing and orchestration flags from `subagent_orch_v1` +and sets `ENABLE_SMART_ROUTING_SAVINGS=1`. Savings tracking, endpoint-price refresh, +pricing, and the above-prompt savings/orchestrator-version row are opt-in and currently +implemented for Claude. Every existing preset sets this flag to `"0"`, overriding an +inherited `"1"`. Without a selector, explicitly set `ENABLE_SMART_ROUTING_SAVINGS=1` +with routing enabled to use the feature. + +```bash +SMART_ROUTER_CONFIG_VERSION=subagent_orch_savings_v0 ug claude +``` + `SMART_ROUTER_NAME` still selects the router independently of the preset. -Smart-routed Claude and Codex sessions install `smart-router`. The `subagent_orch_v0` -and `subagent_orch_v1` versions also install and activate the bundled `smart-router-orchestrator` skill. +Smart-routed Claude and Codex sessions install `smart-router`. The `subagent_orch_v0`, +`subagent_orch_v1`, and `subagent_orch_savings_v0` versions also install and activate the bundled `smart-router-orchestrator` skill. For example: ```bash @@ -306,8 +318,8 @@ SMART_ROUTER_CONFIG_VERSION=subagent_orch_v0 ug claude The version takes precedence over conflicting legacy flags. Once the agent is selected, UG expands its mapping into -`ENABLE_SMART_ROUTING_V2`, `ENABLE_SMART_ROUTING_SUBAGENT_ONLY`, and -`ENABLE_SMART_ROUTER_ORCHESTRATOR` for the launched session. When the version is +`ENABLE_SMART_ROUTING_V2`, `ENABLE_SMART_ROUTING_SUBAGENT_ONLY`, +`ENABLE_SMART_ROUTER_ORCHESTRATOR`, and `ENABLE_SMART_ROUTING_SAVINGS` for the launched session. When the version is unset or empty, these legacy flags retain their existing behavior, including first-prompt routing through `ENABLE_SMART_ROUTING_V2=1`. Unknown versions are ignored: no preset is applied, the inherited environment is unchanged, and commands continue normally. diff --git a/docs/claude-mod.md b/docs/claude-mod.md index 24fbad1e6..a8a74b3a7 100644 --- a/docs/claude-mod.md +++ b/docs/claude-mod.md @@ -62,7 +62,7 @@ The source remains static; no additional launch snapshot file is created. ## Savings estimate -`hooks/savings.ts` draws a dim row above the prompt: the estimated savings from smart +`hooks/register.ts` draws a dim row above the prompt: the estimated savings from smart routing and, when the launch activates Smart Router Orchestrator, the skill's `metadata.version` from `SKILL.claude.md`. For each turn request, main agent and subagents, the mod sums token usage by the model that served it and its baseline. Main @@ -73,10 +73,21 @@ against the model the session started on. After each turn the mod writes those sums to `savings_usage`, a file in the launch's temporary directory. It then runs the `savings_pricer` argv (`python -P -m ucode.smart_routing.savings`), which prices the sums from the -workspace's cached endpoint rates and prints the row's text. Each routed launch +workspace's cached endpoint rates and prints the row's text. Each savings-enabled routed launch refreshes that cache on a background thread. An estimate is hidden, not undercounted, when any served or baseline model has no rate. The orchestrator version is passed as `orchestrator_version`. -The estimate is on by default; `ENABLE_SMART_ROUTING_SAVINGS=0` opts a launch out -and leaves the three options at their declared defaults. +Savings hooks live in `hooks/register.ts` alongside the routing controls: Claude +allows one hook module per plugin and its capability scanner follows engine API +calls only within that module. The shared `session.start` handler captures the +start model before first-prompt routing. Savings handlers are registered only +when the launch options explicitly enable the feature. + +The entire savings mod requires `ENABLE_SMART_ROUTING_SAVINGS=1`. Otherwise no usage +is tracked or written, no endpoint rates are fetched, no pricer runs, and no savings or +orchestrator-version row is shown; the three options stay at their declared defaults. +All existing versioned configurations set the flag to `"0"`. Select +`SMART_ROUTER_CONFIG_VERSION=subagent_orch_savings_v0` to enable it with subagent +routing and orchestration. Session routing off/on also disables/restores savings +from the launch baseline. diff --git a/skills/smart-router/SKILL.md b/skills/smart-router/SKILL.md index 439f8d4a8..dd775bd78 100644 --- a/skills/smart-router/SKILL.md +++ b/skills/smart-router/SKILL.md @@ -30,4 +30,7 @@ Continue in the root unless the user explicitly requests a subagent. Honor that native tools and normal harness model selection, without Smart Router Orchestrator; keep routing off. Existing children can finish. When turned on, apply the `smart-router-orchestrator` skill to further work only if the session was launched with `ENABLE_SMART_ROUTER_ORCHESTRATOR=1`. -Do not change the orchestration feature flag. Return the command's result. +For launches with `ENABLE_SMART_ROUTING_SAVINGS=1`, off/on also disables/restores +savings tracking and its display row. All earlier versioned presets disable savings; +`subagent_orch_savings_v0` enables it with subagent routing and orchestration. +Do not change the orchestration or savings launch baseline. Return the command's result. diff --git a/src/ucode/agents/claude_mods/register.ts b/src/ucode/agents/claude_mods/register.ts index b9ef57c1b..edc4d3063 100644 --- a/src/ucode/agents/claude_mods/register.ts +++ b/src/ucode/agents/claude_mods/register.ts @@ -8,8 +8,6 @@ import { RECIPE_HEADER, } from './constants.ts'; -import { registerSavings, savingsSessionStart } from './savings'; - type Environment = Record; async function readEnvironment($: EngineInterface): Promise { @@ -17,6 +15,7 @@ async function readEnvironment($: EngineInterface): Promise { ENABLE_SMART_ROUTING_V2: await $.env.get('ENABLE_SMART_ROUTING_V2'), ENABLE_SMART_ROUTING_SUBAGENT_ONLY: await $.env.get('ENABLE_SMART_ROUTING_SUBAGENT_ONLY'), ENABLE_SMART_ROUTER_ORCHESTRATOR: await $.env.get('ENABLE_SMART_ROUTER_ORCHESTRATOR'), + ENABLE_SMART_ROUTING_SAVINGS: await $.env.get('ENABLE_SMART_ROUTING_SAVINGS'), ANTHROPIC_CUSTOM_HEADERS: await $.env.get('ANTHROPIC_CUSTOM_HEADERS'), }; } @@ -26,6 +25,7 @@ async function writeEnvironment($: EngineInterface, env: Environment): Promise & { + cache_creation?: Partial>; +}; +type Entry = { + agent: string; + baseline: string; + served: string; + before_user_switch: boolean; + usage: Totals; +}; + +const num = (value: unknown): number => + typeof value === 'number' && Number.isFinite(value) ? value : 0; + +const str = (value: unknown): string | null => (typeof value === 'string' && value ? value : null); + +// Set from the launch's plugin options by register. +let pricer: string[] = []; +let usagePath: string | null = null; +let enabled = false; +let version: string | null = null; + +const entries = new Map(); +// The model the session started on, before any first-prompt routing. Set once. +let startModel: string | null = null; +// The model the main loop last named, so a subagent's baseline is a real model +// id ($.session.model() may return an alias). +let lastMainModel: string | null = null; +// The model the main loop's first request named (first-prompt routing's pick). +let firstMainModel: string | null = null; +// Whether a main request has since named another model, which only the user does. +let userSwitched = false; +let seeding: Promise | undefined; +let estimate: string | null = null; +let dirty = false; +let running = false; +let rerun = false; + +async function mainModel($: any): Promise { + try { + return str(await $.session.model()); + } catch { + return null; + } +} + +function accumulate( + agent: string, + baseline: string, + served: string, + beforeUserSwitch: boolean, + usage: any, +): void { + const key = JSON.stringify([agent, baseline, served, beforeUserSwitch]); + let entry = entries.get(key); + if (!entry) { + entry = { + agent, + baseline, + served, + before_user_switch: beforeUserSwitch, + usage: { + input_tokens: 0, + output_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + }, + }; + entries.set(key, entry); + } + for (const k of TOKEN_KEYS) entry.usage[k] += num(usage[k]); + // The 5m/1h cache-write split is only summed where the API reported it. + const split = usage.cache_creation; + if (split && typeof split === 'object') { + const sums = entry.usage.cache_creation ?? {}; + for (const k of CACHE_SPLIT_KEYS) if (k in split) sums[k] = (sums[k] ?? 0) + num(split[k]); + entry.usage.cache_creation = sums; + } +} + +// A mod reload resets module state while the session carries on, so the first +// hook of a fresh module instance picks up where the last one left off from the +// file it wrote. A missing or malformed file just means starting empty. +async function loadUsage($: any): Promise { + try { + const saved = JSON.parse(await $.fs.read(usagePath)); + if (!saved || saved.version !== USAGE_VERSION || !Array.isArray(saved.entries)) return; + for (const item of saved.entries) { + const agent = str(item?.agent); + const baseline = str(item?.baseline); + const served = str(item?.served); + if (agent && baseline && served && item.usage && typeof item.usage === 'object') { + accumulate(agent, baseline, served, item.before_user_switch === true, item.usage); + dirty = true; + } + } + startModel ??= str(saved.start_model); + lastMainModel ??= str(saved.main_model); + firstMainModel ??= str(saved.first_main_model); + userSwitched ||= saved.user_switched === true; + } catch {} +} + +// State is seeded once per module instance, by whichever hook runs first; +// concurrent hooks wait on the same load. +async function ready($: any): Promise { + if (!enabled || await $.env.get('ENABLE_SMART_ROUTING_SAVINGS') !== '1') return false; + seeding ??= loadUsage($); + await seeding; + return true; +} + +// The pricer's last stdout line; a failed run drops the estimate, which would be stale. +async function runPricer($: any): Promise { + try { + const run = await $.process.run([...pricer, '--usage', usagePath], { + timeoutMs: PRICER_TIMEOUT_MS, + }); + if (run?.exitCode === 0) { + const out = JSON.parse(String(run.stdout ?? '').trim().split('\n').pop() ?? ''); + if (out && typeof out === 'object') return str(out.savings); + } + } catch {} + return null; +} + +async function writeAndPrice($: any): Promise { + dirty = false; + try { + await $.fs.write( + usagePath, + JSON.stringify({ + version: USAGE_VERSION, + start_model: startModel, + main_model: lastMainModel, + first_main_model: firstMainModel, + user_switched: userSwitched, + entries: Array.from(entries.values()), + }), + ); + } catch (error) { + // Retry at the next turn.complete rather than pricing a file that is not there. + dirty = true; + throw error; + } + const next = await runPricer($); + if (next !== estimate) { + estimate = next; + // Claude Code only redraws the band when its props change, not when module state does. + $.ui.invalidate('ui.render'); + } +} + +// One pricer run at a time; a request that arrives mid-run just asks for another. +async function flush($: any): Promise { + if (!await ready($)) return; + if (running) { + rerun = true; + return; + } + running = true; + try { + do { + rerun = false; + await writeAndPrice($); + } while (rerun); + } catch { + } finally { + running = false; + } +} + +const line = (): string | null => + [estimate, version && `Smart Router Orchestrator v${version}`].filter(Boolean).join(' ยท ') || + null; + +function registerSavings(on: On, options: PluginOptions): void { + if (options.ENABLE_SMART_ROUTING_SAVINGS !== '1') return; + const raw = options as Record; + pricer = + Array.isArray(raw.savings_pricer) && raw.savings_pricer.every(arg => typeof arg === 'string') + ? (raw.savings_pricer as string[]) + : []; + usagePath = str(raw.savings_usage); + enabled = pricer.length > 0 && usagePath !== null; + version = str(raw.orchestrator_version); + + // Fires for every request to a model, subagents included (e.agentId is set). + // yield* forwards the streamed response untouched; only the finished result's + // usage is read. + on('turn.step', async function* ($, e, next) { + const request = e as any; + const agent = str(request.agentId); + const active = await ready($); + let baseline: string | null = null; + if (active) { + if (agent === null) { + lastMainModel = str(request.model) ?? lastMainModel; + firstMainModel ??= lastMainModel; + userSwitched ||= lastMainModel !== firstMainModel; + } else baseline = lastMainModel ?? (await mainModel($)); + } + // Read before the request runs: a concurrent main request may switch it. + const beforeUserSwitch = !userSwitched; + const result = yield* next(e); + try { + const usage = (result as any)?.usage; + const served = str(request.model) ?? str(usage?.model); + // Main is its own baseline: routing leaves the main model alone. + const base = agent === null ? served : baseline; + if (active && await ready($) && usage && served && base) { + accumulate(agent ?? 'main', base, served, beforeUserSwitch, usage); + dirty = true; + } + } catch {} + return result; + }); + + // Price at the end of each turn, off the turn's critical path. + on('turn.complete', async ($, e, next) => { + const result = await next(e); + if ((await ready($)) && dirty) void flush($); + return result; + }); + + on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => { + const theirs = await next(e); + if (!enabled || await $.env.get('ENABLE_SMART_ROUTING_SAVINGS') !== '1') return theirs; + const text = line(); + if (text === null) return theirs; + const { Box, Text } = $.ui.resolve(e); + const row = Text({ dimColor: true, children: [text] }); + return theirs ? Box({ flexDirection: 'column', children: [row, theirs] }) : row; + }); +} + +async function savingsSessionStart($: EngineInterface): Promise { + if (await ready($)) { + startModel ??= await mainModel($); + void flush($); + } +} diff --git a/src/ucode/agents/claude_mods/savings.ts b/src/ucode/agents/claude_mods/savings.ts deleted file mode 100644 index 2aaebe84f..000000000 --- a/src/ucode/agents/claude_mods/savings.ts +++ /dev/null @@ -1,274 +0,0 @@ -import type { On, PluginOptions } from 'claude-code'; - -// Smart-routing savings estimate, drawn above the prompt with the Smart Router -// Orchestrator skill version. Per turn request (main agent and subagents) it sums -// token usage by (agent, baseline, served, before_user_switch), where served is -// the model the request named (the model service billed; the id the API reports -// back is a gateway deployment's, which drifts) and baseline is what the request -// would have used without routing: the main model in effect for a subagent, served -// itself for main (routing never changes the main model mid-session). First-prompt -// routing switches the main model once, before the first request, so a later -// change is the user's; requests before it are flagged, and priced from -// start_model under that routing. -// Claude Code's side queries (title, compaction, helpers) are not turn steps, so -// they never count. A hooks module has no Node APIs and the price table lives in -// Python, so after each turn it writes the sums to the launch's savings_usage -// file and runs the savings_pricer command (ucode.smart_routing.savings) on it. -// Everything is best effort and swallows errors. - -const USAGE_VERSION = 1; -const PRICER_TIMEOUT_MS = 20_000; -const TOKEN_KEYS = [ - 'input_tokens', - 'output_tokens', - 'cache_creation_input_tokens', - 'cache_read_input_tokens', -] as const; -const CACHE_SPLIT_KEYS = ['ephemeral_5m_input_tokens', 'ephemeral_1h_input_tokens'] as const; - -type Totals = Record<(typeof TOKEN_KEYS)[number], number> & { - cache_creation?: Partial>; -}; -type Entry = { - agent: string; - baseline: string; - served: string; - before_user_switch: boolean; - usage: Totals; -}; - -const num = (value: unknown): number => - typeof value === 'number' && Number.isFinite(value) ? value : 0; - -const str = (value: unknown): string | null => (typeof value === 'string' && value ? value : null); - -// Set from the launch's plugin options by registerSavings. -let pricer: string[] = []; -let usagePath: string | null = null; -let enabled = false; -let version: string | null = null; - -const entries = new Map(); -// The model the session started on, before any first-prompt routing. Set once. -let startModel: string | null = null; -// The model the main loop last named, so a subagent's baseline is a real model -// id ($.session.model() may return an alias). -let lastMainModel: string | null = null; -// The model the main loop's first request named (first-prompt routing's pick). -let firstMainModel: string | null = null; -// Whether a main request has since named another model, which only the user does. -let userSwitched = false; -let seeding: Promise | undefined; -let estimate: string | null = null; -let dirty = false; -let running = false; -let rerun = false; - -async function mainModel($: any): Promise { - try { - return str(await $.session.model()); - } catch { - return null; - } -} - -function accumulate( - agent: string, - baseline: string, - served: string, - beforeUserSwitch: boolean, - usage: any, -): void { - const key = JSON.stringify([agent, baseline, served, beforeUserSwitch]); - let entry = entries.get(key); - if (!entry) { - entry = { - agent, - baseline, - served, - before_user_switch: beforeUserSwitch, - usage: { - input_tokens: 0, - output_tokens: 0, - cache_creation_input_tokens: 0, - cache_read_input_tokens: 0, - }, - }; - entries.set(key, entry); - } - for (const k of TOKEN_KEYS) entry.usage[k] += num(usage[k]); - // The 5m/1h cache-write split is only summed where the API reported it. - const split = usage.cache_creation; - if (split && typeof split === 'object') { - const sums = entry.usage.cache_creation ?? {}; - for (const k of CACHE_SPLIT_KEYS) if (k in split) sums[k] = (sums[k] ?? 0) + num(split[k]); - entry.usage.cache_creation = sums; - } -} - -// A mod reload resets module state while the session carries on, so the first -// hook of a fresh module instance picks up where the last one left off from the -// file it wrote. A missing or malformed file just means starting empty. -async function loadUsage($: any): Promise { - try { - const saved = JSON.parse(await $.fs.read(usagePath)); - if (!saved || saved.version !== USAGE_VERSION || !Array.isArray(saved.entries)) return; - for (const item of saved.entries) { - const agent = str(item?.agent); - const baseline = str(item?.baseline); - const served = str(item?.served); - if (agent && baseline && served && item.usage && typeof item.usage === 'object') { - accumulate(agent, baseline, served, item.before_user_switch === true, item.usage); - dirty = true; - } - } - startModel ??= str(saved.start_model); - lastMainModel ??= str(saved.main_model); - firstMainModel ??= str(saved.first_main_model); - userSwitched ||= saved.user_switched === true; - } catch {} -} - -// State is seeded once per module instance, by whichever hook runs first; -// concurrent hooks wait on the same load. -async function ready($: any): Promise { - if (!enabled) return false; - seeding ??= loadUsage($); - await seeding; - return true; -} - -// The pricer's last stdout line; a failed run drops the estimate, which would be stale. -async function runPricer($: any): Promise { - try { - const run = await $.process.run([...pricer, '--usage', usagePath], { - timeoutMs: PRICER_TIMEOUT_MS, - }); - if (run?.exitCode === 0) { - const out = JSON.parse(String(run.stdout ?? '').trim().split('\n').pop() ?? ''); - if (out && typeof out === 'object') return str(out.savings); - } - } catch {} - return null; -} - -async function writeAndPrice($: any): Promise { - dirty = false; - try { - await $.fs.write( - usagePath, - JSON.stringify({ - version: USAGE_VERSION, - start_model: startModel, - main_model: lastMainModel, - first_main_model: firstMainModel, - user_switched: userSwitched, - entries: Array.from(entries.values()), - }), - ); - } catch (error) { - // Retry at the next turn.complete rather than pricing a file that is not there. - dirty = true; - throw error; - } - const next = await runPricer($); - if (next !== estimate) { - estimate = next; - // Claude Code only redraws the band when its props change, not when module state does. - $.ui.invalidate('ui.render'); - } -} - -// One pricer run at a time; a request that arrives mid-run just asks for another. -async function flush($: any): Promise { - if (running) { - rerun = true; - return; - } - running = true; - try { - do { - rerun = false; - await writeAndPrice($); - } while (rerun); - } catch { - } finally { - running = false; - } -} - -const line = (): string | null => - [estimate, version && `Smart Router Orchestrator v${version}`].filter(Boolean).join(' ยท ') || - null; - -// The module's session.start hook, which register.ts installs (a second unmatched -// hook on the same event fails the module load). It runs before the first prompt, -// so the start model is captured before any routing happens. session.start can -// fire again (/clear), by when first-prompt routing may have switched the model, -// so the start model is only ever set once. The flush is not awaited: it spawns a -// process and Claude Code waits on this hook before the first prompt. -export async function savingsSessionStart($: any, e: any, next: any): Promise { - if (await ready($)) { - startModel ??= await mainModel($); - void flush($); - } - return next(e); -} - -export function registerSavings(on: On, options: PluginOptions): void { - const raw = options as Record; - pricer = - Array.isArray(raw.savings_pricer) && raw.savings_pricer.every(arg => typeof arg === 'string') - ? (raw.savings_pricer as string[]) - : []; - usagePath = str(raw.savings_usage); - enabled = pricer.length > 0 && usagePath !== null; - version = str(raw.orchestrator_version); - - // Fires for every request to a model, subagents included (e.agentId is set). - // yield* forwards the streamed response untouched; only the finished result's - // usage is read. - on('turn.step', async function* ($, e, next) { - const request = e as any; - const agent = str(request.agentId); - const active = await ready($); - let baseline: string | null = null; - if (active) { - if (agent === null) { - lastMainModel = str(request.model) ?? lastMainModel; - firstMainModel ??= lastMainModel; - userSwitched ||= lastMainModel !== firstMainModel; - } else baseline = lastMainModel ?? (await mainModel($)); - } - // Read before the request runs: a concurrent main request may switch it. - const beforeUserSwitch = !userSwitched; - const result = yield* next(e); - try { - const usage = (result as any)?.usage; - const served = str(request.model) ?? str(usage?.model); - // Main is its own baseline: routing leaves the main model alone. - const base = agent === null ? served : baseline; - if (active && usage && served && base) { - accumulate(agent ?? 'main', base, served, beforeUserSwitch, usage); - dirty = true; - } - } catch {} - return result; - }); - - // Price at the end of each turn, off the turn's critical path. - on('turn.complete', async ($, e, next) => { - const result = await next(e); - if ((await ready($)) && dirty) void flush($); - return result; - }); - - on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => { - const theirs = await next(e); - const text = line(); - if (text === null) return theirs; - const { Box, Text } = $.ui.resolve(e); - const row = Text({ dimColor: true, children: [text] }); - return theirs ? Box({ flexDirection: 'column', children: [row, theirs] }) : row; - }); -} diff --git a/src/ucode/constants.py b/src/ucode/constants.py index 3ef84ee25..bd1db2ef5 100644 --- a/src/ucode/constants.py +++ b/src/ucode/constants.py @@ -8,12 +8,14 @@ ENABLE_SMART_ROUTING_ENV_VAR = "ENABLE_SMART_ROUTING_V2" ENABLE_SUBAGENT_ROUTING_ENV_VAR = "ENABLE_SMART_ROUTING_SUBAGENT_ONLY" +ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR = "ENABLE_SMART_ROUTING_SAVINGS" ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR = "ENABLE_SMART_ROUTER_ORCHESTRATOR" SMART_ROUTER_CONFIG_VERSION_ENV_VAR = "SMART_ROUTER_CONFIG_VERSION" SMART_ROUTING_ENV_KEYS = ( ENABLE_SMART_ROUTING_ENV_VAR, ENABLE_SUBAGENT_ROUTING_ENV_VAR, ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR, + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, ) MODEL_PROVIDER_SERVICE_HEADER = "Databricks-Model-Provider-Service" diff --git a/src/ucode/smart_routing/config.py b/src/ucode/smart_routing/config.py index 37e6a91f5..acd275fde 100644 --- a/src/ucode/smart_routing/config.py +++ b/src/ucode/smart_routing/config.py @@ -10,6 +10,7 @@ AGENT_CODEX, ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR, ENABLE_SMART_ROUTING_ENV_VAR, + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, ENABLE_SUBAGENT_ROUTING_ENV_VAR, SMART_ROUTER_CONFIG_VERSION_ENV_VAR, SMART_ROUTING_ENV_KEYS, @@ -33,31 +34,45 @@ # Preserve subagent_orch_v0 for Claude, with all Codex routing flags disabled. SUBAGENT_ORCH_V0_CLAUDE_ONLY = "subagent_orch_v0_claude_only" +# Subagent orchestration with the Claude savings estimate enabled. +SUBAGENT_ORCH_SAVINGS_V0 = "subagent_orch_savings_v0" + _BASE_VERSIONS = { FIRST_PROMPT_AND_SUBAGENT_NO_ORCH_V0: { ENABLE_SMART_ROUTING_ENV_VAR: "1", ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "0", }, SUBAGENT_ONLY_V0: { ENABLE_SMART_ROUTING_ENV_VAR: "0", ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "0", }, SUBAGENT_ONLY_V1: { ENABLE_SMART_ROUTING_ENV_VAR: "1", ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "0", }, SUBAGENT_ORCH_V0: { ENABLE_SMART_ROUTING_ENV_VAR: "0", ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "0", }, SUBAGENT_ORCH_V1: { ENABLE_SMART_ROUTING_ENV_VAR: "1", ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "0", + }, + SUBAGENT_ORCH_SAVINGS_V0: { + ENABLE_SMART_ROUTING_ENV_VAR: "1", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "1", }, } @@ -83,6 +98,10 @@ AGENT_CLAUDE: _BASE_VERSIONS[SUBAGENT_ORCH_V1], AGENT_CODEX: _BASE_VERSIONS[SUBAGENT_ORCH_V1], }, + SUBAGENT_ORCH_SAVINGS_V0: { + AGENT_CLAUDE: _BASE_VERSIONS[SUBAGENT_ORCH_SAVINGS_V0], + AGENT_CODEX: _BASE_VERSIONS[SUBAGENT_ORCH_SAVINGS_V0], + }, # Codex opts out of routing and orchestration; Claude keeps subagent_orch_v0. SUBAGENT_ORCH_V0_CLAUDE_ONLY: { AGENT_CLAUDE: _BASE_VERSIONS[SUBAGENT_ORCH_V0], @@ -90,6 +109,7 @@ ENABLE_SMART_ROUTING_ENV_VAR: "0", ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "0", }, }, } diff --git a/src/ucode/smart_routing/savings.py b/src/ucode/smart_routing/savings.py index a25cbba25..724298d36 100644 --- a/src/ucode/smart_routing/savings.py +++ b/src/ucode/smart_routing/savings.py @@ -5,7 +5,7 @@ saved = sum(tokens * price(baseline)) - sum(tokens * price(model that served them)) -per token class. The mod (``agents/claude_mods/savings.ts``) sums each turn request's usage by +per token class. The mod (``agents/claude_mods/register.ts``) sums each turn request's usage by (agent, baseline, served) and runs this module on that file after every turn; it prints one JSON line with the segment the mod draws. A main-agent request is its own baseline, and a subagent's is the main model in effect when it ran, so switching models with ``/model`` never reprices earlier @@ -25,6 +25,8 @@ from decimal import ROUND_HALF_UP, Decimal from pathlib import Path +from ucode.constants import AGENT_CLAUDE, ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR +from ucode.smart_routing.config import resolve_environment from ucode.smart_routing.pricing import ( ModelPrice, TokenUsage, @@ -154,10 +156,14 @@ def main(argv: list[str] | None = None) -> int: parser.add_argument("--price-cache", type=Path, required=True) parser.add_argument("--baseline-session-start", action="store_true") args = parser.parse_args(argv) - line = usage_output( - args.usage, - price_cache=args.price_cache, - baseline_session_start=args.baseline_session_start, + line = ( + usage_output( + args.usage, + price_cache=args.price_cache, + baseline_session_start=args.baseline_session_start, + ) + if resolve_environment(agent=AGENT_CLAUDE).get(ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR) == "1" + else json.dumps({"savings": None}) ) # Write bytes so the money bag survives a non-UTF-8 stdout locale. sys.stdout.buffer.write((line + "\n").encode("utf-8")) diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 075141811..2fcef04f7 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -32,6 +32,7 @@ AGENT_CODEX, ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR, ENABLE_SMART_ROUTING_ENV_VAR, + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, ENABLE_SUBAGENT_ROUTING_ENV_VAR, LOOPBACK_HOST, SMART_ROUTER_CONFIG_VERSION_ENV_VAR, @@ -84,8 +85,6 @@ CLAUDE_ROUTE_SELECTION_TIMEOUT_S = 20.0 CLAUDE_ROUTED_AGENT_PREFIX = "ucode-route-" CLAUDE_ROUTING_PLUGIN_NAME = "ug-smart-router" -# Default on; "0" opts a launch out of the smart-routing savings estimate. -ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR = "ENABLE_SMART_ROUTING_SAVINGS" # Native plugin options for the savings mod, at the values Claude fills in for a launch that omits # them: no pricer (the estimate is off) and no orchestrator skill version to show. CLAUDE_SAVINGS_OPTION_DEFAULTS: dict[str, str | list[str]] = { @@ -357,9 +356,9 @@ def _claude_mod_launch_options(env: Mapping[str, str] | None = None) -> dict: def savings_enabled(env: Mapping[str, str] | None = None) -> bool: - """Whether a routed Claude launch shows the savings estimate (default on; ``0`` opts out).""" - source = os.environ if env is None else env - return source.get(ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, "1") != "0" + """Whether a routed Claude launch explicitly enables the savings feature.""" + source = resolve_environment(env, agent=AGENT_CLAUDE) + return source.get(ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR) == "1" def _claude_savings_launch_options( @@ -368,13 +367,14 @@ def _claude_savings_launch_options( """Options for the savings mod: the pricer it runs after each turn, the file it sums usage into, and the orchestrator skill version it shows when the launch activates that skill.""" options = dict(CLAUDE_SAVINGS_OPTION_DEFAULTS) - if savings_enabled(): - options["savings_pricer"] = savings.pricer_argv( - python=sys.executable, - price_cache=price_cache, - baseline_session_start=baseline_session_start, - ) - options["savings_usage"] = str(launch_dir / savings.USAGE_FILENAME) + if not savings_enabled(): + return options + options["savings_pricer"] = savings.pricer_argv( + python=sys.executable, + price_cache=price_cache, + baseline_session_start=baseline_session_start, + ) + options["savings_usage"] = str(launch_dir / savings.USAGE_FILENAME) if orchestrator.feature_enabled(agent=AGENT_CLAUDE): options["orchestrator_version"] = orchestrator.skill_version(agent=AGENT_CLAUDE) or "" return options @@ -507,7 +507,7 @@ def _write_routed_claude_plugin(plugin_dir: Path, model_ids: list[str]) -> None: orchestrator.add_claude_agents(plugin_dir) mods = Path(__file__).parents[1] / "agents" / "claude_mods" # constants.ts is imported by the entry point, so it ships beside it rather than as a module. - for name in ("register.ts", "constants.ts", "savings.ts"): + for name in ("register.ts", "constants.ts"): write_text_file(plugin_dir / "hooks" / name, (mods / name).read_text(encoding="utf-8")) write_json_file(plugin_dir / "hooks" / "hooks.json", {"modules": ["./register.ts"]}) @@ -761,6 +761,7 @@ def launch_claude( env[ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR] = ( "1" if orchestrator.feature_enabled(agent=AGENT_CLAUDE) else "0" ) + env[ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR] = "1" if savings_enabled() else "0" if route_first_prompt: env[ENABLE_SMART_ROUTING_ENV_VAR] = "1" else: diff --git a/tests/README.md b/tests/README.md index c57784852..fbba9891d 100644 --- a/tests/README.md +++ b/tests/README.md @@ -23,8 +23,13 @@ Savings estimate coverage: `test_smart_routing_savings.py` prices the mod's toke (per-entry and first-prompt baselines, unpriced or malformed entries, the pricer's single JSON line). `test_smart_routing_pricing.py` covers endpoint rates and the price cache. `test_claude_mod.py` and `test_claude_smart_routing_v2.py` cover the savings options, -opt-out and price refresh. In the native suite, the options reach the mod unchanged and -a real session records its start model. Token accumulation in `turn.step` and the +explicit opt-in (unset, empty, zero, and invalid values stay off), price refresh, +and gating the orchestrator-version row. Preset coverage in `test_smart_routing_config.py` +checks the savings flag, inherited-value precedence, restoration, and launch/session off +controls. `test_smart_router.py` checks that the skill toggle writes all managed flags. +`test_claude_mod.py` checks native literal setters against the flag registry. In the +native suite, options reach the mod unchanged, enabled sessions record their start +model, and disabled sessions create no usage file even with stale savings options. Token accumulation in `turn.step` and the above-prompt row are not covered automatically. Dedicated CUJs reuse `integration/utils` session/terminal mechanics and file-task diff --git a/tests/conftest.py b/tests/conftest.py index 11c76e4cc..e59e5b609 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -77,7 +77,7 @@ def reject_privileged_write(path, _desired_text): monkeypatch.delenv("ENABLE_SMART_ROUTER_ORCHESTRATOR", raising=False) monkeypatch.delenv("SMART_ROUTER_CONFIG_VERSION", raising=False) monkeypatch.delenv("SMART_ROUTER_NAME", raising=False) - # A developer's savings opt-out would otherwise drop the savings options from launch tests. + # Savings stays off unless a test explicitly enables it. monkeypatch.delenv("ENABLE_SMART_ROUTING_SAVINGS", raising=False) # On Windows, resolve_command swaps a bare program name for whatever `shutil.which` # finds on the developer's PATH (e.g. a real `codex.CMD`). Rebind only the compatibility diff --git a/tests/integration/README.md b/tests/integration/README.md index 5eab3067c..18e7ead15 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -1026,3 +1026,9 @@ responses, with only the existing verified native compatibility retries accepted Local Claude mod confirmation replaces the old saved-control-file/assistant-turn wait; Codex retains its CLI control-file checks. No extra TUI sessions or tasks are added. Clear/reload, compaction, and denied-write rollback are not independently checked. + +Savings gating and `subagent_orch_savings_v0` have component coverage in +`test_smart_routing_config.py`, `test_claude_mod.py`, and `test_claude_smart_routing_v2.py`. +Native component tests check opt-in usage-file creation and opt-out preservation. +Live savings accumulation, display, and session off/on are not covered by the +integration suite. diff --git a/tests/native_claude/test_mod.py b/tests/native_claude/test_mod.py index c1661218a..79918bdf9 100644 --- a/tests/native_claude/test_mod.py +++ b/tests/native_claude/test_mod.py @@ -60,12 +60,18 @@ def test_mod_validation_and_event_forwarding(native_binary, tmp_path): ) assert result.returncode == 0, result.stdout + result.stderr if command == "validate": - assert "./register.ts hooks: session.start" in result.stdout + hooks = next( + line for line in result.stdout.splitlines() if "./register.ts hooks:" in line + ) + assert "session.start" in hooks -def test_native_options_are_launch_local(native_binary, tmp_path): +@pytest.mark.parametrize( + "version", [config.SUBAGENT_ORCH_V0_CLAUDE_ONLY, config.SUBAGENT_ORCH_SAVINGS_V0] +) +def test_native_options_are_launch_local(native_binary, tmp_path, version): env = { - "SMART_ROUTER_CONFIG_VERSION": config.SUBAGENT_ORCH_V0_CLAUDE_ONLY, + "SMART_ROUTER_CONFIG_VERSION": version, "SMART_ROUTER_NAME": 'custom "quoted" ฮป', } savings_options = _savings_options(tmp_path / "a") @@ -86,7 +92,9 @@ def test_savings_mod_records_the_start_model(native_binary, tmp_path): root = tmp_path / "session" options = _savings_options(root) usage = Path(options["savings_usage"]) - with Session(native_binary, root, {}, options=options) as session: + with Session( + native_binary, root, {"ENABLE_SMART_ROUTING_SAVINGS": "1"}, options=options + ) as session: session.observe() # The mod writes its sums off the session.start critical path. deadline = time.monotonic() + 10 @@ -106,3 +114,14 @@ def test_savings_mod_records_the_start_model(native_binary, tmp_path): "user_switched": False, "entries": [], } + + +@pytest.mark.parametrize("flag", [None, "", "0"]) +def test_disabled_savings_does_not_create_usage_file(native_binary, tmp_path, flag): + root = tmp_path / "session" + options = _savings_options(root) + env = {} if flag is None else {"ENABLE_SMART_ROUTING_SAVINGS": flag} + # Stale savings options must not enable the feature without an explicit flag. + with Session(native_binary, root, env, options=options) as session: + session.observe() + assert not Path(options["savings_usage"]).exists() diff --git a/tests/test_claude_mod.py b/tests/test_claude_mod.py index d211f8fd4..d022feae1 100644 --- a/tests/test_claude_mod.py +++ b/tests/test_claude_mod.py @@ -16,13 +16,12 @@ def test_plugin_contains_mod_entry_point(tmp_path): assert json.loads((tmp_path / "hooks/hooks.json").read_text()) == {"modules": ["./register.ts"]} mods = Path(v2.__file__).parents[1] / "agents" / "claude_mods" # constants.ts is imported by the entry point, so it ships beside it unmodified. - for name in ("register.ts", "constants.ts", "savings.ts"): + for name in ("register.ts", "constants.ts"): assert (tmp_path / "hooks" / name).read_bytes() == (mods / name).read_bytes() assert sorted(path.name for path in (tmp_path / "hooks").iterdir()) == [ "constants.ts", "hooks.json", "register.ts", - "savings.ts", ] @@ -44,6 +43,7 @@ def test_launch_options_resolve_versions_without_mutating_input(version): "ENABLE_SMART_ROUTING_V2": "0", "ENABLE_SMART_ROUTING_SUBAGENT_ONLY": "0", "ENABLE_SMART_ROUTER_ORCHESTRATOR": "0", + "ENABLE_SMART_ROUTING_SAVINGS": "0", "SMART_ROUTER_NAME": " custom_recipe ", } before = env.copy() @@ -64,23 +64,28 @@ def test_launch_options_preserve_empty_and_absent_values(): "ENABLE_SMART_ROUTING_V2": "", "ENABLE_SMART_ROUTING_SUBAGENT_ONLY": "", "ENABLE_SMART_ROUTER_ORCHESTRATOR": "", - "unset": ["ENABLE_SMART_ROUTING_SUBAGENT_ONLY", "ENABLE_SMART_ROUTER_ORCHESTRATOR"], + "ENABLE_SMART_ROUTING_SAVINGS": "", + "unset": [ + "ENABLE_SMART_ROUTING_SUBAGENT_ONLY", + "ENABLE_SMART_ROUTER_ORCHESTRATOR", + "ENABLE_SMART_ROUTING_SAVINGS", + ], "recipe": "task_v3", } -def test_savings_is_on_by_default_and_opts_out_with_zero(monkeypatch): - assert v2.savings_enabled() - monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, "0") - assert not v2.savings_enabled() - monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, "1") - assert v2.savings_enabled() +@pytest.mark.parametrize("flag", [None, "", "0", "1", "true", "2"]) +def test_savings_requires_explicit_one(monkeypatch, flag): + if flag is not None: + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, flag) + assert v2.savings_enabled() == (flag == "1") @pytest.mark.parametrize("first_prompt", [False, True]) def test_savings_options_give_the_mod_a_pricer_and_a_launch_local_usage_file( - tmp_path, first_prompt + tmp_path, monkeypatch, first_prompt ): + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, "1") price_cache = tmp_path / "model-prices.json" assert v2._claude_savings_launch_options( @@ -106,9 +111,7 @@ def test_opting_out_of_savings_leaves_its_options_at_their_defaults(tmp_path, mo @pytest.mark.parametrize("savings_flag", ["1", "0"]) -def test_savings_options_show_the_orchestrator_skill_version_when_it_is_active( - tmp_path, monkeypatch, savings_flag -): +def test_savings_gates_the_orchestrator_version_too(tmp_path, monkeypatch, savings_flag): monkeypatch.setenv("ENABLE_SMART_ROUTER_ORCHESTRATOR", "1") monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, savings_flag) @@ -116,5 +119,13 @@ def test_savings_options_show_the_orchestrator_skill_version_when_it_is_active( tmp_path, price_cache=tmp_path / "model-prices.json", baseline_session_start=False ) - assert options["orchestrator_version"] == orchestrator.skill_version(agent=AGENT_CLAUDE) - assert options["orchestrator_version"] + expected = orchestrator.skill_version(agent=AGENT_CLAUDE) if savings_flag == "1" else "" + assert options["orchestrator_version"] == expected + + +def test_native_mod_covers_every_managed_flag(): + source = (Path(v2.__file__).parents[1] / "agents/claude_mods/register.ts").read_text() + for key in SMART_ROUTING_ENV_KEYS: + assert f"$.env.get('{key}')" in source + assert f"$.env.set('{key}', env.{key})" in source + assert f"{key}: enabled ? original('{key}') : '0'" in source diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 383cdc305..f059a9760 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -446,6 +446,7 @@ def fake_run(_argv, **kwargs): def test_subagent_only_launch_skips_first_prompt_routing( self, tmp_path, monkeypatch, savings_price_refresh ): + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, "1") user_settings = tmp_path / "settings.json" user_settings.write_text(json.dumps({"model": "opus"})) monkeypatch.delenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, raising=False) @@ -579,6 +580,7 @@ def _launch(monkeypatch, tmp_path, *, first_prompt: bool) -> dict: def capture(argv) -> None: captured["launch_dir"] = Path(argv[argv.index("--plugin-dir") + 1]).parent settings = json.loads(Path(argv[argv.index("--settings") + 1]).read_text()) + captured["env"] = settings["env"] captured["options"] = settings["pluginConfigs"][v2.CLAUDE_ROUTING_PLUGIN_NAME][ "options" ] @@ -603,7 +605,16 @@ def popen(argv, **_kwargs) -> Mock: binary="claude", user_settings_path=user_settings, launch_model="opus", - compose_settings=lambda _args: ({}, []), + compose_settings=lambda _args: ( + { + "env": { + v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "0" + if v2.savings_enabled() + else "1" + } + }, + [], + ), launch_model_args=claude._launch_model_args, model_name=claude._maybe_add_1m_suffix, ) @@ -614,6 +625,7 @@ def popen(argv, **_kwargs) -> Mock: def test_first_prompt_routing_prices_against_the_start_model( self, monkeypatch, tmp_path, first_prompt ): + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, "1") captured = self._launch(monkeypatch, tmp_path, first_prompt=first_prompt) assert captured["options"]["savings_pricer"] == savings.pricer_argv( @@ -625,10 +637,13 @@ def test_first_prompt_routing_prices_against_the_start_model( captured["launch_dir"] / savings.USAGE_FILENAME ) - def test_opting_out_skips_the_pricer_and_the_price_refresh( - self, monkeypatch, tmp_path, savings_price_refresh + @pytest.mark.parametrize("flag", [None, "", "0", "true"]) + def test_disabled_savings_skips_the_pricer_and_the_price_refresh( + self, monkeypatch, tmp_path, savings_price_refresh, flag ): - monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, "0") + if flag is not None: + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, flag) + monkeypatch.setenv("ENABLE_SMART_ROUTER_ORCHESTRATOR", "1") captured = self._launch(monkeypatch, tmp_path, first_prompt=False) @@ -637,6 +652,26 @@ def test_opting_out_skips_the_pricer_and_the_price_refresh( } == v2.CLAUDE_SAVINGS_OPTION_DEFAULTS savings_price_refresh.assert_not_called() + @pytest.mark.parametrize( + "version, enabled", + [ + ("subagent_orch_v1", False), + ("subagent_orch_savings_v0", True), + ], + ) + def test_preset_controls_all_savings_launch_work( + self, monkeypatch, tmp_path, savings_price_refresh, version, enabled + ): + monkeypatch.setenv("SMART_ROUTER_CONFIG_VERSION", version) + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, "0" if enabled else "1") + captured = self._launch(monkeypatch, tmp_path, first_prompt=False) + assert captured["env"][v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR] == str(int(enabled)) + assert captured["options"][v2.ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR] == str(int(enabled)) + assert bool(captured["options"]["savings_pricer"]) == enabled + assert bool(captured["options"]["savings_usage"]) == enabled + assert bool(captured["options"]["orchestrator_version"]) == enabled + assert savings_price_refresh.called == enabled + class TestSavingsPriceRefresh: def test_prices_routable_and_main_models_by_system_ai_name(self): diff --git a/tests/test_smart_router.py b/tests/test_smart_router.py index 504cf7991..4be82c82d 100644 --- a/tests/test_smart_router.py +++ b/tests/test_smart_router.py @@ -12,6 +12,7 @@ from typer.testing import CliRunner from ucode import cli, config_io, skills +from ucode.constants import SMART_ROUTING_ENV_KEYS from ucode.skills import SMART_ROUTER_SKILL from ucode.smart_routing import session_env, v2 @@ -70,15 +71,7 @@ def test_skill_toggles_with_launch_installation_despite_shadowed_path(tmp_path, ) assert result.returncode == 0, result.stdout + result.stderr assert f"Smart Router is {'off' if action == 'disable' else 'on'}" in result.stdout - expected = ( - { - v2.ENABLE_SMART_ROUTING_ENV_VAR: "0", - v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", - v2.ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", - } - if action == "disable" - else {} - ) + expected = dict.fromkeys(SMART_ROUTING_ENV_KEYS, "0") if action == "disable" else {} assert json.loads(session_path.read_text()) == expected diff --git a/tests/test_smart_routing_config.py b/tests/test_smart_routing_config.py index 4b081402b..fd60cb321 100644 --- a/tests/test_smart_routing_config.py +++ b/tests/test_smart_routing_config.py @@ -9,6 +9,7 @@ AGENT_CODEX, ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR, ENABLE_SMART_ROUTING_ENV_VAR, + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, ENABLE_SUBAGENT_ROUTING_ENV_VAR, SMART_ROUTER_CONFIG_VERSION_ENV_VAR, ) @@ -19,26 +20,37 @@ ENABLE_SMART_ROUTING_ENV_VAR: "1", ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "0", }, config.SUBAGENT_ONLY_V0: { ENABLE_SMART_ROUTING_ENV_VAR: "0", ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "0", }, config.SUBAGENT_ONLY_V1: { ENABLE_SMART_ROUTING_ENV_VAR: "1", ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "0", }, config.SUBAGENT_ORCH_V0: { ENABLE_SMART_ROUTING_ENV_VAR: "0", ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "0", }, config.SUBAGENT_ORCH_V1: { ENABLE_SMART_ROUTING_ENV_VAR: "1", ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "0", + }, + config.SUBAGENT_ORCH_SAVINGS_V0: { + ENABLE_SMART_ROUTING_ENV_VAR: "1", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "1", }, } @@ -51,6 +63,7 @@ def test_absent_or_unknown_selector_is_a_no_op(selector): ENABLE_SMART_ROUTING_ENV_VAR: "1", ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "0", "INHERITED_SETTING": "preserved", } expected = applied.copy() @@ -67,6 +80,7 @@ def test_absent_or_unknown_selector_is_a_no_op(selector): @pytest.mark.parametrize("ENABLE_SMART_ROUTING_V2", [None, "0", "1"]) @pytest.mark.parametrize("ENABLE_SMART_ROUTING_SUBAGENT_ONLY", [None, "0", "1"]) @pytest.mark.parametrize("ENABLE_SMART_ROUTER_ORCHESTRATOR", [None, "0", "1"]) +@pytest.mark.parametrize("ENABLE_SMART_ROUTING_SAVINGS", [None, "0", "1"]) @pytest.mark.parametrize( "SMART_ROUTER_CONFIG_VERSION", [ @@ -76,12 +90,14 @@ def test_absent_or_unknown_selector_is_a_no_op(selector): config.SUBAGENT_ONLY_V1, config.SUBAGENT_ORCH_V0, config.SUBAGENT_ORCH_V1, + config.SUBAGENT_ORCH_SAVINGS_V0, ], ) def test_smart_routing_config_cartesian_grid( ENABLE_SMART_ROUTING_V2, ENABLE_SMART_ROUTING_SUBAGENT_ONLY, ENABLE_SMART_ROUTER_ORCHESTRATOR, + ENABLE_SMART_ROUTING_SAVINGS, SMART_ROUTER_CONFIG_VERSION, ): source = { @@ -90,6 +106,7 @@ def test_smart_routing_config_cartesian_grid( (ENABLE_SMART_ROUTING_ENV_VAR, ENABLE_SMART_ROUTING_V2), (ENABLE_SUBAGENT_ROUTING_ENV_VAR, ENABLE_SMART_ROUTING_SUBAGENT_ONLY), (ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR, ENABLE_SMART_ROUTER_ORCHESTRATOR), + (ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR, ENABLE_SMART_ROUTING_SAVINGS), ) if value is not None } @@ -180,6 +197,7 @@ def test_agent_config_selection_and_session_precedence(monkeypatch, agent, enabl "ENABLE_SMART_ROUTING_V2": "0", "ENABLE_SMART_ROUTING_SUBAGENT_ONLY": "0", "ENABLE_SMART_ROUTER_ORCHESTRATOR": "0", + "ENABLE_SMART_ROUTING_SAVINGS": "0", }, ) source = {"SMART_ROUTER_CONFIG_VERSION": "subagent_orch_v0"} @@ -187,6 +205,7 @@ def test_agent_config_selection_and_session_precedence(monkeypatch, agent, enabl "ENABLE_SMART_ROUTING_V2": "0", "ENABLE_SMART_ROUTING_SUBAGENT_ONLY": enabled, "ENABLE_SMART_ROUTER_ORCHESTRATOR": enabled, + "ENABLE_SMART_ROUTING_SAVINGS": "0", } assert config.resolve_environment(source, agent=agent) == expected assert v2.smart_routing_enabled(source, default=True, agent=agent) == (enabled == "1") @@ -200,3 +219,29 @@ def test_agent_config_selection_and_session_precedence(monkeypatch, agent, enabl session_env.SESSION_ENV_VAR: str(control), "ENABLE_SMART_ROUTER_ORCHESTRATOR": "0", } + + +@pytest.mark.parametrize("agent", [AGENT_CLAUDE, AGENT_CODEX]) +def test_savings_preset_launch_and_session_off_restore_the_baseline(agent, tmp_path): + source = { + SMART_ROUTER_CONFIG_VERSION_ENV_VAR: config.SUBAGENT_ORCH_SAVINGS_V0, + ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "0", + } + original = source.copy() + previous = v2.override_smart_routing(False, source, agent=agent) + assert all(source[key] == "0" for key in config.SMART_ROUTING_ENV_KEYS) + v2.restore_smart_routing_env(previous, source) + assert source == original + + controls = tmp_path / "session.json" + controls.write_text('{"ENABLE_SMART_ROUTING_SAVINGS":"0"}') + source[session_env.SESSION_ENV_VAR] = str(controls) + effective = session_env.effective_environment(source, agent=agent) + assert effective[ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR] == "0" + assert effective[ENABLE_SUBAGENT_ROUTING_ENV_VAR] == "1" + + +def test_savings_alone_does_not_enable_routing(): + source = {ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR: "1"} + assert not v2.smart_routing_enabled(source, default=False, agent=AGENT_CLAUDE) + assert not v2.first_prompt_routing_enabled(source, agent=AGENT_CLAUDE) diff --git a/tests/test_smart_routing_savings.py b/tests/test_smart_routing_savings.py index 726ba93b0..ea7fec4e5 100644 --- a/tests/test_smart_routing_savings.py +++ b/tests/test_smart_routing_savings.py @@ -293,15 +293,19 @@ def test_runs_the_module_on_the_given_python(self, tmp_path, first_prompt): ] @staticmethod - def run(price_cache: Path, usage: Path, *, first_prompt: bool = False): + def run(price_cache: Path, usage: Path, *, first_prompt: bool = False, flag: str | None = "1"): argv = savings.pricer_argv( python=sys.executable, price_cache=price_cache, baseline_session_start=first_prompt ) + env = {**os.environ, "PYTHONIOENCODING": "ascii"} + env.pop("ENABLE_SMART_ROUTING_SAVINGS", None) + if flag is not None: + env["ENABLE_SMART_ROUTING_SAVINGS"] = flag # An ASCII stdout locale must not cost the money bag. return subprocess.run( [*argv, "--usage", str(usage)], capture_output=True, - env={**os.environ, "PYTHONIOENCODING": "ascii"}, + env=env, timeout=30, check=False, ) @@ -322,3 +326,10 @@ def test_never_fails_the_mod(self, tmp_path, price_cache): assert result.returncode == 0, result.stderr assert json.loads(result.stdout) == {"savings": None} + + @pytest.mark.parametrize("flag", [None, "", "0", "true"]) + def test_disabled_pricer_returns_no_estimate(self, tmp_path, price_cache, flag): + usage = write_usage(tmp_path, [entry("a1", OPUS_ID, SONNET_ID, SUBAGENT_USAGE)]) + result = self.run(price_cache, usage, flag=flag) + assert result.returncode == 0, result.stderr + assert json.loads(result.stdout) == {"savings": None} From a5a32946be1f75be2116e3c0f196c444765f92b8 Mon Sep 17 00:00:00 2001 From: Josh Joseph Date: Tue, 13 Oct 2026 22:24:46 +0000 Subject: [PATCH 3/7] Redraw the savings row after /smart-router on|off Co-authored-by: Isaac --- src/ucode/agents/claude_mods/register.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/ucode/agents/claude_mods/register.ts b/src/ucode/agents/claude_mods/register.ts index edc4d3063..576bdfaf4 100644 --- a/src/ucode/agents/claude_mods/register.ts +++ b/src/ucode/agents/claude_mods/register.ts @@ -105,6 +105,8 @@ export function register(on: On, options: PluginOptions): void { if (!launchedEnabled) return {text: 'Smart Router was not enabled for this session.'}; try { await setRouting($, options, action === 'on', true); + // The savings row hides while routing is off, but the toggle changes no band props. + if (options.ENABLE_SMART_ROUTING_SAVINGS === '1') $.ui.invalidate('ui.render'); const guidance = options.ENABLE_SMART_ROUTER_ORCHESTRATOR !== '1' ? '' : action === 'off' ? ' Earlier Smart Router Orchestrator instructions are superseded. Stop automatic delegation; honor explicit subagent requests using native model selection. Existing children can finish.' : ' Apply the smart-router-orchestrator skill to further work.'; From 58b4e479dabc5d6a716f5ab4be9cbcb05fd2c95c Mon Sep 17 00:00:00 2001 From: Josh Joseph Date: Wed, 14 Oct 2026 00:08:40 +0000 Subject: [PATCH 4/7] Expect the savings flag in the Codex smart-router off snapshot Co-authored-by: Isaac --- tests/integration/test_ug_smart_routing_hooks.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_ug_smart_routing_hooks.py b/tests/integration/test_ug_smart_routing_hooks.py index afb6f00fd..db0b352a6 100644 --- a/tests/integration/test_ug_smart_routing_hooks.py +++ b/tests/integration/test_ug_smart_routing_hooks.py @@ -167,6 +167,7 @@ def _toggle_with_skill( "ENABLE_SMART_ROUTING_V2": "0", "ENABLE_SMART_ROUTING_SUBAGENT_ONLY": "0", "ENABLE_SMART_ROUTER_ORCHESTRATOR": "0", + "ENABLE_SMART_ROUTING_SAVINGS": "0", } ) assert json.loads(controls[0].read_text()) != expected From 754ec73a016e447b739da22e4f6a6aba80bdb8e3 Mon Sep 17 00:00:00 2001 From: Josh Joseph Date: Wed, 14 Oct 2026 17:22:43 +0000 Subject: [PATCH 5/7] Style the savings row like the status band and place it under it Label the version as Unity Gateway Smart Router, show the amount in a green (saved) or yellow (cost more) pill, and render the row after the other AbovePrompt handlers so it sits under the band in either order. The pricer now prints the estimate's parts instead of preformatted text. Co-authored-by: Isaac --- docs/claude-mod.md | 12 +++-- src/ucode/agents/claude_mods/register.ts | 58 ++++++++++++++++------ src/ucode/smart_routing/savings.py | 16 +++--- tests/test_smart_routing_savings.py | 63 ++++++++++-------------- 4 files changed, 83 insertions(+), 66 deletions(-) diff --git a/docs/claude-mod.md b/docs/claude-mod.md index 602d75fc4..93828783f 100644 --- a/docs/claude-mod.md +++ b/docs/claude-mod.md @@ -64,9 +64,12 @@ The source remains static; no additional launch snapshot file is created. ## Savings estimate -`hooks/register.ts` draws a dim row above the prompt: the estimated savings from smart -routing and, when the launch activates Smart Router Orchestrator, the skill's -`metadata.version` from `SKILL.claude.md`. For each turn request, main agent and +`hooks/register.ts` draws a row under the status band, styled like it: the estimated +savings from smart routing, with the amount in a pill (green for a saving, yellow when +routing cost more), and, when the launch activates Smart Router Orchestrator, the skill's +`metadata.version` from `SKILL.claude.md`, labeled `Unity Gateway Smart Router v`. +The row goes after the other `AbovePrompt` handlers' output, so it stays under the band +whichever order Claude runs them in. For each turn request, main agent and subagents, the mod sums token usage by the model that served it and its baseline. Main is its own baseline, and a subagent's baseline is the main model in effect when it ran. Under first-prompt routing, requests before the user changes the main model are priced @@ -75,7 +78,8 @@ against the model the session started on. After each turn the mod writes those sums to `savings_usage`, a file in the launch's temporary directory. It then runs the `savings_pricer` argv (`python -P -m ucode.smart_routing.savings`), which prices the sums from the -workspace's cached endpoint rates and prints the row's text. Each savings-enabled routed launch +workspace's cached endpoint rates and prints the estimate's parts (direction, amount, +percent) as one JSON line. Each savings-enabled routed launch refreshes that cache on a background thread. An estimate is hidden, not undercounted, when any served or baseline model has no rate. The orchestrator version is passed as `orchestrator_version`. diff --git a/src/ucode/agents/claude_mods/register.ts b/src/ucode/agents/claude_mods/register.ts index 4ae2c97ec..8732251e3 100644 --- a/src/ucode/agents/claude_mods/register.ts +++ b/src/ucode/agents/claude_mods/register.ts @@ -116,9 +116,10 @@ export function register(on: On, options: PluginOptions): void { registerStatusBand(on); } -// Smart-routing savings estimate, drawn above the prompt with the Smart Router -// Orchestrator skill version. Per turn request (main agent and subagents) it sums -// token usage by (agent, baseline, served, before_user_switch), where served is +// Smart-routing savings estimate, drawn above the prompt under the status band, with +// the orchestrator skill version labeled as the Unity Gateway Smart Router's. Per +// turn request (main agent and subagents) it sums token usage by +// (agent, baseline, served, before_user_switch), where served is // the model the request named (the model service billed; the id the API reports // back is a gateway deployment's, which drifts) and baseline is what the request // would have used without routing: the main model in effect for a subagent, served @@ -152,6 +153,8 @@ type Entry = { before_user_switch: boolean; usage: Totals; }; +// The pricer's segment: a saving, or a cost increase when routing chose pricier models. +type Estimate = { saved: boolean; amount: string; percent: string }; const num = (value: unknown): number => typeof value === 'number' && Number.isFinite(value) ? value : 0; @@ -175,7 +178,7 @@ let firstMainModel: string | null = null; // Whether a main request has since named another model, which only the user does. let userSwitched = false; let seeding: Promise | undefined; -let estimate: string | null = null; +let estimate: Estimate | null = null; let dirty = false; let running = false; let rerun = false; @@ -254,15 +257,22 @@ async function ready($: any): Promise { return true; } +function estimateOf(value: any): Estimate | null { + if (!value || typeof value !== 'object' || typeof value.saved !== 'boolean') return null; + const amount = str(value.amount); + const percent = str(value.percent); + return amount && percent ? { saved: value.saved, amount, percent } : null; +} + // The pricer's last stdout line; a failed run drops the estimate, which would be stale. -async function runPricer($: any): Promise { +async function runPricer($: any): Promise { try { const run = await $.process.run([...pricer, '--usage', usagePath], { timeoutMs: PRICER_TIMEOUT_MS, }); if (run?.exitCode === 0) { const out = JSON.parse(String(run.stdout ?? '').trim().split('\n').pop() ?? ''); - if (out && typeof out === 'object') return str(out.savings); + if (out && typeof out === 'object') return estimateOf(out.savings); } } catch {} return null; @@ -288,7 +298,7 @@ async function writeAndPrice($: any): Promise { throw error; } const next = await runPricer($); - if (next !== estimate) { + if (JSON.stringify(next) !== JSON.stringify(estimate)) { estimate = next; // Claude Code only redraws the band when its props change, not when module state does. $.ui.invalidate('ui.render'); @@ -314,9 +324,26 @@ async function flush($: any): Promise { } } -const line = (): string | null => - [estimate, version && `Smart Router Orchestrator v${version}`].filter(Boolean).join(' ยท ') || - null; +// Styled like the status band: a bold label, the amount in a pill, a dim italic trailer. +function savingsRow({ Box, Text }: any): unknown { + const children: unknown[] = []; + if (estimate) { + const color = estimate.saved ? 'green' : 'yellow'; + children.push( + Text({ + bold: true, + children: [estimate.saved ? '๐Ÿ’ฐ Est. saved with smart routing' : 'Smart routing cost'], + }), + Text({ bold: true, color: 'black', backgroundColor: color, children: [` ${estimate.amount} `] }), + Text({ color, children: [`${estimate.saved ? '' : 'more '}(${estimate.percent})`] }), + ); + } + if (version) { + const label = `Unity Gateway Smart Router v${version}`; + children.push(Text({ dimColor: true, italic: true, children: [estimate ? `โ€” ${label}` : label] })); + } + return children.length ? Box({ flexDirection: 'row', columnGap: 1, children }) : null; +} function registerSavings(on: On, options: PluginOptions): void { if (options.ENABLE_SMART_ROUTING_SAVINGS !== '1') return; @@ -370,11 +397,12 @@ function registerSavings(on: On, options: PluginOptions): void { on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => { const theirs = await next(e); if (!enabled || await $.env.get('ENABLE_SMART_ROUTING_SAVINGS') !== '1') return theirs; - const text = line(); - if (text === null) return theirs; - const { Box, Text } = $.ui.resolve(e); - const row = Text({ dimColor: true, children: [text] }); - return theirs ? Box({ flexDirection: 'column', children: [row, theirs] }) : row; + const ui = $.ui.resolve(e); + const row = savingsRow(ui); + if (row === null) return theirs; + // Going after the inner handlers' output puts the row under the status band + // whichever handler Claude Code runs first. + return theirs ? ui.Box({ flexDirection: 'column', children: [theirs, row] }) : row; }); } diff --git a/src/ucode/smart_routing/savings.py b/src/ucode/smart_routing/savings.py index 724298d36..5b367a104 100644 --- a/src/ucode/smart_routing/savings.py +++ b/src/ucode/smart_routing/savings.py @@ -21,7 +21,6 @@ import argparse import json -import sys from decimal import ROUND_HALF_UP, Decimal from pathlib import Path @@ -122,16 +121,14 @@ def usage_savings( return baseline_total - actual_total, baseline_total -def savings_text(saved: Decimal, baseline: Decimal) -> str: - """A money-bag estimate when positive, an honest cost increase when not.""" +def savings_segment(saved: Decimal, baseline: Decimal) -> dict[str, object]: + """The estimate's parts, for the mod to style: a saving, or an honest cost increase.""" percent = (abs(saved) / baseline * 100).quantize(Decimal(1), rounding=ROUND_HALF_UP) magnitude = abs(saved) amount = ( "<$0.01" if magnitude < _CENT else f"${magnitude.quantize(_CENT, rounding=ROUND_HALF_UP):,}" ) - if saved >= 0: - return f"๐Ÿ’ฐ Est. saved with smart routing: {amount} ({percent}%)" - return f"Smart routing cost {amount} more ({percent}%)" + return {"saved": saved >= 0, "amount": amount, "percent": f"{percent}%"} def usage_output(usage_path: Path, *, price_cache: Path, baseline_session_start: bool) -> str: @@ -142,10 +139,10 @@ def usage_output(usage_path: Path, *, price_cache: Path, baseline_session_start: computed = usage_savings( usage_path, price_cache=price_cache, baseline_session_start=baseline_session_start ) - savings = None if computed is None else savings_text(*computed) + savings = None if computed is None else savings_segment(*computed) except Exception: # noqa: BLE001 pass - return json.dumps({"savings": savings}, ensure_ascii=False) + return json.dumps({"savings": savings}) def main(argv: list[str] | None = None) -> int: @@ -165,8 +162,7 @@ def main(argv: list[str] | None = None) -> int: if resolve_environment(agent=AGENT_CLAUDE).get(ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR) == "1" else json.dumps({"savings": None}) ) - # Write bytes so the money bag survives a non-UTF-8 stdout locale. - sys.stdout.buffer.write((line + "\n").encode("utf-8")) + print(line) return 0 diff --git a/tests/test_smart_routing_savings.py b/tests/test_smart_routing_savings.py index ea7fec4e5..48cda2901 100644 --- a/tests/test_smart_routing_savings.py +++ b/tests/test_smart_routing_savings.py @@ -71,7 +71,15 @@ def write_usage(tmp_path: Path, entries: list, *, start_model: str | None = OPUS return path -def priced(usage: Path, price_cache: Path, *, first_prompt: bool = False) -> str | None: +def saved(amount: str, percent: str) -> dict: + return {"saved": True, "amount": amount, "percent": percent} + + +def cost_more(amount: str, percent: str) -> dict: + return {"saved": False, "amount": amount, "percent": percent} + + +def priced(usage: Path, price_cache: Path, *, first_prompt: bool = False) -> dict | None: line = savings.usage_output(usage, price_cache=price_cache, baseline_session_start=first_prompt) assert "\n" not in line return json.loads(line)["savings"] @@ -85,9 +93,7 @@ def test_prices_tokens_against_each_entrys_baseline(self, tmp_path, price_cache) ] # Baseline $0.38005 - actual $0.19705. - assert priced(write_usage(tmp_path, entries), price_cache) == ( - "๐Ÿ’ฐ Est. saved with smart routing: $0.18 (48%)" - ) + assert priced(write_usage(tmp_path, entries), price_cache) == saved("$0.18", "48%") def test_an_entrys_baseline_is_the_main_model_when_its_request_ran(self, tmp_path, price_cache): entries = [ @@ -98,9 +104,7 @@ def test_an_entrys_baseline_is_the_main_model_when_its_request_ran(self, tmp_pat entry("a1", SONNET_ID, SONNET_ID, SUBAGENT_USAGE), ] - assert priced(write_usage(tmp_path, entries), price_cache) == ( - "๐Ÿ’ฐ Est. saved with smart routing: $0.18 (34%)" - ) + assert priced(write_usage(tmp_path, entries), price_cache) == saved("$0.18", "34%") def test_shows_a_net_cost_increase_honestly(self, tmp_path, price_cache): entries = [ @@ -109,9 +113,7 @@ def test_shows_a_net_cost_increase_honestly(self, tmp_path, price_cache): ] # Baseline $0.15202 vs actual $0.33502. - assert priced(write_usage(tmp_path, entries), price_cache) == ( - "Smart routing cost $0.18 more (120%)" - ) + assert priced(write_usage(tmp_path, entries), price_cache) == cost_more("$0.18", "120%") def test_nothing_to_claim_until_a_model_was_swapped(self, tmp_path, price_cache): entries = [entry("main", OPUS_ID, OPUS_ID, MAIN_USAGE)] @@ -131,7 +133,7 @@ def test_bills_entries_without_the_cache_write_split_as_one_hour_writes( # SUBAGENT_USAGE does: $0.122 on Sonnet vs $0.305 at Opus rates. assert priced( write_usage(tmp_path, [entry("a1", OPUS_ID, SONNET_ID, usage)]), price_cache - ) == ("๐Ÿ’ฐ Est. saved with smart routing: $0.18 (60%)") + ) == saved("$0.18", "60%") def test_honors_an_explicit_cache_write_split(self, tmp_path, price_cache): usage = { @@ -144,7 +146,7 @@ def test_honors_an_explicit_cache_write_split(self, tmp_path, price_cache): # The split says 5-minute writes: $0.092 on Sonnet vs $0.23 at Opus rates. assert priced( write_usage(tmp_path, [entry("a1", OPUS_ID, SONNET_ID, usage)]), price_cache - ) == ("๐Ÿ’ฐ Est. saved with smart routing: $0.14 (60%)") + ) == saved("$0.14", "60%") def test_skips_entries_without_tokens(self, tmp_path, price_cache): entries = [ @@ -152,9 +154,7 @@ def test_skips_entries_without_tokens(self, tmp_path, price_cache): entry("a2", "unpriced-baseline", "unpriced-served", {"input_tokens": 0}), ] - assert priced(write_usage(tmp_path, entries), price_cache) == ( - "๐Ÿ’ฐ Est. saved with smart routing: $0.18 (60%)" - ) + assert priced(write_usage(tmp_path, entries), price_cache) == saved("$0.18", "60%") def test_first_prompt_routing_prices_against_the_start_model(self, tmp_path, price_cache): # Under first-prompt routing the router picks the main model too, so an entry's own @@ -167,9 +167,7 @@ def test_first_prompt_routing_prices_against_the_start_model(self, tmp_path, pri usage = write_usage(tmp_path, entries, start_model=OPUS_ID) # Baseline $0.38005 at Opus vs actual $0.15202 on Sonnet. - assert priced(usage, price_cache, first_prompt=True) == ( - "๐Ÿ’ฐ Est. saved with smart routing: $0.23 (60%)" - ) + assert priced(usage, price_cache, first_prompt=True) == saved("$0.23", "60%") assert priced(usage, price_cache) is None def test_first_prompt_routing_prices_entries_after_a_user_switch_as_their_own( @@ -185,9 +183,7 @@ def test_first_prompt_routing_prices_entries_after_a_user_switch_as_their_own( usage = write_usage(tmp_path, entries, start_model=OPUS_ID) # Baseline $0.48512 - actual $0.25709. - assert priced(usage, price_cache, first_prompt=True) == ( - "๐Ÿ’ฐ Est. saved with smart routing: $0.23 (47%)" - ) + assert priced(usage, price_cache, first_prompt=True) == saved("$0.23", "47%") @pytest.mark.parametrize("start_model", [None, ""]) def test_first_prompt_routing_without_a_start_model_has_no_estimate( @@ -255,24 +251,19 @@ def test_a_file_in_another_format_is_an_error_not_an_empty_estimate( savings.usage_savings(usage, price_cache=price_cache, baseline_session_start=False) -class TestSavingsText: +class TestSavingsSegment: def test_rounds_to_cents_and_whole_percent(self): - assert ( - savings.savings_text(Decimal("1234.565"), Decimal("2469.13")) - == "๐Ÿ’ฐ Est. saved with smart routing: $1,234.57 (50%)" + assert savings.savings_segment(Decimal("1234.565"), Decimal("2469.13")) == saved( + "$1,234.57", "50%" ) def test_negative_reads_as_a_cost_increase(self): - assert ( - savings.savings_text(Decimal("-0.04"), Decimal("0.40")) - == "Smart routing cost $0.04 more (10%)" + assert savings.savings_segment(Decimal("-0.04"), Decimal("0.40")) == cost_more( + "$0.04", "10%" ) def test_sub_cent_amounts_are_not_shown_as_zero(self): - assert ( - savings.savings_text(Decimal("0.004"), Decimal("1")) - == "๐Ÿ’ฐ Est. saved with smart routing: <$0.01 (0%)" - ) + assert savings.savings_segment(Decimal("0.004"), Decimal("1")) == saved("<$0.01", "0%") class TestPricerCommand: @@ -301,7 +292,7 @@ def run(price_cache: Path, usage: Path, *, first_prompt: bool = False, flag: str env.pop("ENABLE_SMART_ROUTING_SAVINGS", None) if flag is not None: env["ENABLE_SMART_ROUTING_SAVINGS"] = flag - # An ASCII stdout locale must not cost the money bag. + # The line must survive an ASCII stdout locale. return subprocess.run( [*argv, "--usage", str(usage)], capture_output=True, @@ -310,16 +301,14 @@ def run(price_cache: Path, usage: Path, *, first_prompt: bool = False, flag: str check=False, ) - def test_prints_exactly_one_utf8_json_line(self, tmp_path, price_cache): + def test_prints_exactly_one_json_line(self, tmp_path, price_cache): usage = write_usage(tmp_path, [entry("a1", OPUS_ID, SONNET_ID, SUBAGENT_USAGE)]) result = self.run(price_cache, usage) assert result.returncode == 0, result.stderr assert result.stdout.count(b"\n") == 1 - assert json.loads(result.stdout.decode("utf-8")) == { - "savings": "๐Ÿ’ฐ Est. saved with smart routing: $0.18 (60%)" - } + assert json.loads(result.stdout.decode("utf-8")) == {"savings": saved("$0.18", "60%")} def test_never_fails_the_mod(self, tmp_path, price_cache): result = self.run(price_cache, tmp_path / "missing.json", first_prompt=True) From a2628c749043022e1173bd97a69da70db710adfc Mon Sep 17 00:00:00 2001 From: Josh Joseph Date: Fri, 16 Oct 2026 23:18:22 +0000 Subject: [PATCH 6/7] Keep the status band's turn.complete hook in its own file Claude allows one unmatched hook per event, and the status band already has the turn.complete one, so the savings hook matches every turn on isAborted instead of the band exporting its hide logic to register.ts. Co-authored-by: Isaac --- src/ucode/agents/claude_mods/register.ts | 24 +++++++++---------- .../claude_mods/smart-routing-status.ts | 16 +++++++------ 2 files changed, 21 insertions(+), 19 deletions(-) diff --git a/src/ucode/agents/claude_mods/register.ts b/src/ucode/agents/claude_mods/register.ts index 311d7bb85..cb6e54aa2 100644 --- a/src/ucode/agents/claude_mods/register.ts +++ b/src/ucode/agents/claude_mods/register.ts @@ -1,5 +1,5 @@ import type { EngineInterface, On, PluginOptions } from 'claude-code'; -import { hideStatusBandAfterTurn, registerStatusBand, showStatusBand } from './smart-routing-status.ts'; +import { registerStatusBand, showStatusBand } from './smart-routing-status.ts'; import { registerSubagentRouting } from './subagent-routing.ts'; import { CLI_AGENT, @@ -100,14 +100,6 @@ export function register(on: On, options: PluginOptions): void { return next(e); }); - // Claude permits only one unfiltered handler per event. - on('turn.complete', async ($, e, next) => { - if (hideStatusBandAfterTurn()) $.ui.invalidate('ui.render'); - const result = await next(e); - if ((await ready($)) && dirty) void flush($); - return result; - }); - // The installed smart-router skill already supplies the slash command. on('command.run', {command: 'smart-router'}, async ($, e) => { const action = e.args.trim(); @@ -126,9 +118,9 @@ export function register(on: On, options: PluginOptions): void { registerSubagentRouting(on, options); } -// Smart-routing savings estimate, drawn above the prompt under the status band, with -// the orchestrator skill version labeled as the Unity Gateway Smart Router's. Per -// turn request (main agent and subagents) it sums token usage by +// Smart-routing savings estimate, drawn above the prompt (under the status band while +// it shows), with the orchestrator skill version labeled as the Unity Gateway Smart +// Router's. Per turn request (main agent and subagents) it sums token usage by // (agent, baseline, served, before_user_switch), where served is // the model the request named (the model service billed; the id the API reports // back is a gateway deployment's, which drifts) and baseline is what the request @@ -397,6 +389,14 @@ function registerSavings(on: On, options: PluginOptions): void { return result; }); + // Price at the end of each turn, off the turn's critical path. The status band holds the + // event's one unmatched hook, so this one matches every turn instead. + on('turn.complete', { isAborted: [false, true] }, async ($, e, next) => { + const result = await next(e); + if ((await ready($)) && dirty) void flush($); + return result; + }); + on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => { const theirs = await next(e); if (!enabled || await $.env.get('ENABLE_SMART_ROUTING_SAVINGS') !== '1') return theirs; diff --git a/src/ucode/agents/claude_mods/smart-routing-status.ts b/src/ucode/agents/claude_mods/smart-routing-status.ts index 461d85450..2421ada4c 100644 --- a/src/ucode/agents/claude_mods/smart-routing-status.ts +++ b/src/ucode/agents/claude_mods/smart-routing-status.ts @@ -12,13 +12,6 @@ export function showStatusBand(): void { band.isHideArmed = false } -export function hideStatusBandAfterTurn(): boolean { - if (!band.isHideArmed) return false - band.isVisible = false - band.isHideArmed = false - return true -} - type RoutingState = { enabled: boolean; firstPrompt: boolean } // Claude's capability scanner requires literal names in $.env.get calls. @@ -56,6 +49,15 @@ export const registerStatusBand = (on: On): void => { return next(e) }) + on('turn.complete', ($, e, next) => { + if (band.isHideArmed) { + band.isVisible = false + band.isHideArmed = false + $.ui.invalidate('ui.render') + } + return next(e) + }) + // The prompt hint row always carries the state as a plain tail, so it stays visible after // the band hides and flips as soon as a toggle succeeds. on('ui.render', { component: 'PromptHint' }, async ($, e, next) => { From 67b1465e5c431c140c0c60b7b928cde428177389 Mon Sep 17 00:00:00 2001 From: Josh Joseph Date: Fri, 16 Oct 2026 23:17:36 +0000 Subject: [PATCH 7/7] Show the savings row only once there is an estimate The status band and prompt hint row already show the routing state, so a version-only row just repeated the band. The orchestrator version now rides along with the estimate. Adds a mod test for both states. Co-authored-by: Isaac --- README.md | 2 +- docs/claude-mod.md | 18 ++-- src/ucode/agents/claude_mods/register.ts | 28 +++--- tests/README.md | 8 +- tests/native_claude/mod_tests/savings.test.ts | 89 +++++++++++++++++++ 5 files changed, 119 insertions(+), 26 deletions(-) create mode 100644 tests/native_claude/mod_tests/savings.test.ts diff --git a/README.md b/README.md index 405727404..05e66adab 100644 --- a/README.md +++ b/README.md @@ -320,7 +320,7 @@ on/off controls retain their existing precedence over presets. `subagent_orch_savings_v0` uses the routing and orchestration flags from `subagent_orch_v1` and sets `ENABLE_SMART_ROUTING_SAVINGS=1`. Savings tracking, endpoint-price refresh, -pricing, and the above-prompt savings/orchestrator-version row are opt-in and currently +pricing, and the above-prompt savings row are opt-in and currently implemented for Claude. Every existing preset sets this flag to `"0"`, overriding an inherited `"1"`. Without a selector, explicitly set `ENABLE_SMART_ROUTING_SAVINGS=1` with routing enabled to use the feature. diff --git a/docs/claude-mod.md b/docs/claude-mod.md index caebb4ca1..8f7aa8af1 100644 --- a/docs/claude-mod.md +++ b/docs/claude-mod.md @@ -79,12 +79,14 @@ The source remains static; no additional launch snapshot file is created. ## Savings estimate -`hooks/register.ts` draws a row under the status band, styled like it: the estimated -savings from smart routing, with the amount in a pill (green for a saving, yellow when -routing cost more), and, when the launch activates Smart Router Orchestrator, the skill's -`metadata.version` from `SKILL.claude.md`, labeled `Unity Gateway Smart Router v`. -The row goes after the other `AbovePrompt` handlers' output, so it stays under the band -whichever order Claude runs them in. For each turn request, main agent and +`hooks/register.ts` draws a row above the prompt, styled like the status band: the +estimated savings from smart routing, with the amount in a pill (green for a saving, +yellow when routing cost more), and, when the launch activates Smart Router Orchestrator, +the skill's `metadata.version` from `SKILL.claude.md`, labeled +`Unity Gateway Smart Router v`. The row appears only once there is an estimate; +the band and the prompt hint row already show the routing state. It goes after the other +`AbovePrompt` handlers' output, so while the band shows it stays under it whichever order +Claude runs them in. For each turn request, main agent and subagents, the mod sums token usage by the model that served it and its baseline. Main is its own baseline, and a subagent's baseline is the main model in effect when it ran. Under first-prompt routing, requests before the user changes the main model are priced @@ -106,8 +108,8 @@ start model before first-prompt routing. Savings handlers are registered only when the launch options explicitly enable the feature. The entire savings mod requires `ENABLE_SMART_ROUTING_SAVINGS=1`. Otherwise no usage -is tracked or written, no endpoint rates are fetched, no pricer runs, and no savings or -orchestrator-version row is shown; the three options stay at their declared defaults. +is tracked or written, no endpoint rates are fetched, no pricer runs, and no savings row is +shown; the three options stay at their declared defaults. All existing versioned configurations set the flag to `"0"`. Select `SMART_ROUTER_CONFIG_VERSION=subagent_orch_savings_v0` to enable it with subagent routing and orchestration. Session routing off/on also disables/restores savings diff --git a/src/ucode/agents/claude_mods/register.ts b/src/ucode/agents/claude_mods/register.ts index cb6e54aa2..91df2d0d1 100644 --- a/src/ucode/agents/claude_mods/register.ts +++ b/src/ucode/agents/claude_mods/register.ts @@ -327,24 +327,24 @@ async function flush($: any): Promise { } // Styled like the status band: a bold label, the amount in a pill, a dim italic trailer. +// Nothing until there is an estimate: the band and prompt hint already show routing state. function savingsRow({ Box, Text }: any): unknown { - const children: unknown[] = []; - if (estimate) { - const color = estimate.saved ? 'green' : 'yellow'; + if (!estimate) return null; + const color = estimate.saved ? 'green' : 'yellow'; + const children: unknown[] = [ + Text({ + bold: true, + children: [estimate.saved ? '๐Ÿ’ฐ Est. saved with smart routing' : 'Smart routing cost'], + }), + Text({ bold: true, color: 'black', backgroundColor: color, children: [` ${estimate.amount} `] }), + Text({ color, children: [`${estimate.saved ? '' : 'more '}(${estimate.percent})`] }), + ]; + if (version) { children.push( - Text({ - bold: true, - children: [estimate.saved ? '๐Ÿ’ฐ Est. saved with smart routing' : 'Smart routing cost'], - }), - Text({ bold: true, color: 'black', backgroundColor: color, children: [` ${estimate.amount} `] }), - Text({ color, children: [`${estimate.saved ? '' : 'more '}(${estimate.percent})`] }), + Text({ dimColor: true, italic: true, children: [`โ€” Unity Gateway Smart Router v${version}`] }), ); } - if (version) { - const label = `Unity Gateway Smart Router v${version}`; - children.push(Text({ dimColor: true, italic: true, children: [estimate ? `โ€” ${label}` : label] })); - } - return children.length ? Box({ flexDirection: 'row', columnGap: 1, children }) : null; + return Box({ flexDirection: 'row', columnGap: 1, children }); } function registerSavings(on: On, options: PluginOptions): void { diff --git a/tests/README.md b/tests/README.md index 5ebea53ca..5840e41d6 100644 --- a/tests/README.md +++ b/tests/README.md @@ -54,9 +54,11 @@ Savings estimate coverage: `test_smart_routing_savings.py` prices the mod's toke JSON line). `test_smart_routing_pricing.py` covers endpoint rates and the price cache. `test_claude_mod.py` and `test_claude_smart_routing_v2.py` cover the savings options, explicit opt-in (unset, empty, zero, and invalid values stay off), price refresh, -and gating the orchestrator-version row. Preset coverage in `test_smart_routing_config.py` -checks the savings flag, inherited-value precedence, restoration, and launch/session off -controls. `test_smart_router.py` checks that the skill toggle writes all managed flags. +and gating the orchestrator version. `native_claude/mod_tests/savings.test.ts` checks that +the savings row shows only while the pricer has an estimate. Preset coverage in +`test_smart_routing_config.py` checks the savings flag, inherited-value precedence, +restoration, and launch/session off controls. `test_smart_router.py` checks that the skill +toggle writes all managed flags. `test_claude_mod.py` checks native literal setters against the flag registry. In the native suite, options reach the mod unchanged, enabled sessions record their start model, and disabled sessions create no usage file even with stale savings options. Token accumulation in `turn.step` and the diff --git a/tests/native_claude/mod_tests/savings.test.ts b/tests/native_claude/mod_tests/savings.test.ts new file mode 100644 index 000000000..7f35ef9e9 --- /dev/null +++ b/tests/native_claude/mod_tests/savings.test.ts @@ -0,0 +1,89 @@ +import type { On } from 'claude-code'; +import { expect, test } from 'claude-code/testing'; + +const USAGE = '/launch/savings-usage.json'; +const PRICER = ['/ug/python', '-P', '-m', 'ucode.smart_routing.savings']; +const FLAGS = { + ENABLE_SMART_ROUTING_V2: '0', + ENABLE_SMART_ROUTING_SUBAGENT_ONLY: '1', + ENABLE_SMART_ROUTER_ORCHESTRATOR: '1', + ENABLE_SMART_ROUTING_SAVINGS: '1', +}; +const OPTIONS = { + ...FLAGS, + unset: [], + recipe: 'task_v3', + savings_pricer: PRICER, + savings_usage: USAGE, + orchestrator_version: '1.4.0', +}; +const ABOVE_PROMPT = { plugin: 'ug-smart-router', surface: 'terminal', component: 'AbovePrompt', props: {} } as const; +const start = { surface: 'terminal', isInteractive: true, cwd: '/work' } as const; + +// A subagent routed from Opus to Sonnet, as an earlier module instance wrote it. +const SAVED_USAGE = { + version: 1, + start_model: 'claude-opus-4-8', + main_model: 'claude-opus-4-8', + first_main_model: 'claude-opus-4-8', + user_switched: false, + entries: [ + { + agent: 'a1', + baseline: 'claude-opus-4-8', + served: 'claude-sonnet-5', + before_user_switch: true, + usage: { input_tokens: 1000, output_tokens: 4000, cache_creation_input_tokens: 0, cache_read_input_tokens: 0 }, + }, + ], +}; + +// The engine beneath the mod: the launch environment, the saved sums, another plugin's +// row, and a pricer that prints whatever estimate the test last set. +function engine(on: On) { + const env: Record = { ...FLAGS }; + const pricer: { savings: unknown } = { savings: null }; + on('env.get', ($, e) => ({ value: env[e.name] })); + on('env.set', ($, e) => { + env[e.name] = e.value; + return { value: undefined }; + }); + on('session.start', ($, e) => ({ cwd: e.cwd })); + on('session.model', () => ({ value: 'claude-opus-4-8' })); + on('fs.read', ($, e) => (e.path === USAGE ? { value: JSON.stringify(SAVED_USAGE) } : { deny: 'missing' })); + on('fs.write', () => ({ value: undefined })); + on('process.run', () => ({ + value: { exitCode: 0, stdout: JSON.stringify({ savings: pricer.savings }), stderr: '' }, + })); + on('ui.render', { component: 'AbovePrompt' }, ($, e) => $.ui.resolve(e).Text({ children: ['theirs'] })); + return { pricer }; +} + +test('the savings row shows only while the pricer has an estimate', { options: OPTIONS }, async ($, on) => { + const { pricer } = engine(on); + const band = await $.ui.mount(ABOVE_PROMPT); + + // No estimate yet (no rates, or nothing rerouted): no row, not even the version. + await $.session.start(start); + let drawn = JSON.stringify(await band.drawn()); + expect(drawn).toContain('theirs'); + expect(drawn).not.toContain('Est. saved'); + expect(drawn).not.toContain('Unity Gateway Smart Router v'); + + // A reload's session.start prices again. + pricer.savings = { saved: true, amount: '$0.18', percent: '48%' }; + await $.session.start(start); + drawn = JSON.stringify(await band.drawn()); + expect(drawn).toContain('theirs'); + expect(drawn).toContain('๐Ÿ’ฐ Est. saved with smart routing'); + expect(drawn).toContain(' $0.18 '); + expect(drawn).toContain('(48%)'); + expect(drawn).toContain('โ€” Unity Gateway Smart Router v1.4.0'); + + // A failed or empty pricing drops the stale estimate, and the row with it. + pricer.savings = null; + await $.session.start(start); + drawn = JSON.stringify(await band.drawn()); + expect(drawn).toContain('theirs'); + expect(drawn).not.toContain('Est. saved'); +});