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 a9734e93f..93828783f 100644 --- a/docs/claude-mod.md +++ b/docs/claude-mod.md @@ -61,3 +61,39 @@ 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/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 +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 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`. + +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 27edf8b49..8732251e3 100644 --- a/src/ucode/agents/claude_mods/register.ts +++ b/src/ucode/agents/claude_mods/register.ts @@ -16,6 +16,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'), }; } @@ -25,6 +26,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; +}; +// 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; + +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: Estimate | 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; +} + +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 { + 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 estimateOf(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 (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'); + } +} + +// 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; + } +} + +// 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; + 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 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; + }); +} + +async function savingsSessionStart($: EngineInterface): Promise { + if (await ready($)) { + startModel ??= await mainModel($); + void flush($); + } +} 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/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/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/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..5b367a104 --- /dev/null +++ b/src/ucode/smart_routing/savings.py @@ -0,0 +1,170 @@ +"""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/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 +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 +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, + 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_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):,}" + ) + return {"saved": saved >= 0, "amount": amount, "percent": f"{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_segment(*computed) + except Exception: # noqa: BLE001 + pass + return json.dumps({"savings": savings}) + + +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, + ) + if resolve_environment(agent=AGENT_CLAUDE).get(ENABLE_SMART_ROUTING_SAVINGS_ENV_VAR) == "1" + else json.dumps({"savings": None}) + ) + print(line) + 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 7d083bd42..7831c5fa0 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 @@ -31,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, @@ -40,6 +42,7 @@ from ucode.databricks import ( AnthropicModelCatalog, build_auth_token_argv, + fetch_endpoint_rates, get_databricks_token, list_anthropic_model_catalog, list_anthropic_models, @@ -56,7 +59,7 @@ install_skill, uninstall_agent_skill, ) -from ucode.smart_routing import claude_routing, orchestrator, routing +from ucode.smart_routing import claude_routing, orchestrator, pricing, routing, savings from ucode.smart_routing.claude_hooks import ( FIRST_PROMPT_SOCKET_ENV, sync_first_prompt_hook, @@ -82,6 +85,13 @@ CLAUDE_ROUTE_SELECTION_TIMEOUT_S = 20.0 CLAUDE_ROUTED_AGENT_PREFIX = "ucode-route-" CLAUDE_ROUTING_PLUGIN_NAME = "ug-smart-router" +# 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." @@ -345,6 +355,89 @@ 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 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( + 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 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 + + +# 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( @@ -372,6 +465,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", + }, }, }, ) @@ -649,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: @@ -687,9 +800,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 2a4e3e03a..e55b57494 100644 --- a/tests/README.md +++ b/tests/README.md @@ -23,6 +23,19 @@ CUJ 4's preset journeys check Claude's ON/OFF status band, routing scope, and updates after off/on toggles. Checks ignore terminal styling and trailing UI controls; offline cases reject stale states, wrong routing scope, and command confirmations. +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, +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 and transcript readers, not its config stubs or pytest fixtures. Prompt/model correlation remains CUJ-specific. Concurrent runs may read the same CUJ workspace. diff --git a/tests/conftest.py b/tests/conftest.py index ab2f4f3d3..e59e5b609 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) + # 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 # 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 b2301e728..cf834c198 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -438,7 +438,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 @@ -1028,3 +1031,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/integration/test_ug_smart_routing_hooks.py b/tests/integration/test_ug_smart_routing_hooks.py index 49c1bb822..da4b54590 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 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..79918bdf9 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") @@ -40,18 +60,68 @@ 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") 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, {"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 + 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": [], + } + + +@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 584b5fa8d..635172759 100644 --- a/tests/test_claude_mod.py +++ b/tests/test_claude_mod.py @@ -1,12 +1,13 @@ """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): @@ -25,6 +26,17 @@ def test_plugin_contains_mod_entry_point(tmp_path): ] +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 = { @@ -32,6 +44,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() @@ -52,6 +65,68 @@ 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", } + + +@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, 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( + 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_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) + + options = v2._claude_savings_launch_options( + tmp_path, price_cache=tmp_path / "model-prices.json", baseline_session_start=False + ) + + 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 e5fbc5ba8..f059a9760 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,10 @@ 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 + ): + 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) @@ -500,9 +516,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 +552,220 @@ 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["env"] = settings["env"] + 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: ( + { + "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, + ) + 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 + ): + 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( + 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 + ) + + @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 + ): + 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) + + assert { + key: captured["options"][key] for key in v2.CLAUDE_SAVINGS_OPTION_DEFAULTS + } == 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): + 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 f19ff8e56..ce52675fa 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 @@ -115,6 +116,47 @@ def test_orchestrator_injects_the_installed_agent_instructions( assert output["hookSpecificOutput"]["additionalContext"] == context +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_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_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..48cda2901 --- /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 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"] + + +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) == saved("$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) == saved("$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) == 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)] + + 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 + ) == saved("$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 + ) == saved("$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) == 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 + # 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) == 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( + 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) == saved("$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 TestSavingsSegment: + def test_rounds_to_cents_and_whole_percent(self): + 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_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_segment(Decimal("0.004"), Decimal("1")) == saved("<$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, 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 + # The line must survive an ASCII stdout locale. + return subprocess.run( + [*argv, "--usage", str(usage)], + capture_output=True, + env=env, + timeout=30, + check=False, + ) + + 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": 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) + + 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}