diff --git a/public/dashboard.html b/public/dashboard.html index f68be3a..90c033d 100644 --- a/public/dashboard.html +++ b/public/dashboard.html @@ -39,9 +39,18 @@

Supported models

try { const res = await fetch('/api/stats'); const s = await res.json(); + // Share of all prompt volume served from cache. Copilot bills per + // premium request rather than per token, so that count — not a dollar + // estimate — is what actually tracks consumption. + const cacheRate = s.total_input_tokens > 0 + ? Math.round(s.total_cache_read_tokens / s.total_input_tokens * 100) + '%' + : '—'; const cards = [ ['Requests', s.request_count], + ['Premium requests', (s.premium_requests ?? 0).toFixed(2)], ['Input tokens', s.total_input_tokens], + ['Cache hit rate', cacheRate], + ['Cache writes', s.total_cache_creation_tokens], ['Output tokens', s.total_output_tokens], ['Bytes sent', s.bytes_sent], ['Bytes received', s.bytes_received], diff --git a/public/requests.html b/public/requests.html index 706ec14..ae580c8 100644 --- a/public/requests.html +++ b/public/requests.html @@ -22,18 +22,44 @@ .caret { color: #8b949e; width: 1em; text-align: center; } .detail-row td { background: #010409; padding: 0; } .detail { padding: 0.75rem 1rem; display: grid; gap: 0.75rem; } - .detail h3 { margin: 0 0 0.25rem; font-size: 0.7rem; text-transform: uppercase; letter-spacing: 0.05em; color: #8b949e; } + .detail h3 { margin: 0; font-size: 0.7rem; text-transform: uppercase; letter-spacing: 0.05em; color: #8b949e; } + .body-head { display: flex; align-items: center; justify-content: space-between; gap: 0.5rem; margin-bottom: 0.25rem; } + .copy-btn { background: transparent; border: 1px solid #30363d; border-radius: 5px; padding: 0.25rem 0.4rem; color: #8b949e; line-height: 0; cursor: pointer; } + .copy-btn:hover { color: #e6edf3; border-color: #8b949e; background: #21262d; } + .copy-btn.copied { color: #3fb950; border-color: #3fb950; } pre { margin: 0; max-height: 360px; overflow: auto; background: #0d1117; border: 1px solid #30363d; border-radius: 6px; padding: 0.75rem; font-size: 0.8rem; white-space: pre-wrap; word-break: break-word; } .muted { color: #8b949e; font-style: italic; padding: 0.75rem 1rem; } + .cache-rate { color: #8b949e; font-size: 0.9em; } + .unknown { color: #8b949e; } + .idle-hot { color: #d29922; font-weight: 600; } + .sess { color: #58a6ff; cursor: pointer; font-family: ui-monospace, monospace; } + .sess:hover { text-decoration: underline; } + .fail-kind { color: #f85149; font-size: 0.75em; opacity: 0.85; margin-top: 2px; } + tr.row-failed td { background: #2d1214; } + tr.row-failed:hover td { background: #3d181b; } + .toolbar { display: flex; gap: 0.75rem; align-items: center; margin-bottom: 0.75rem; font-size: 0.85rem; color: #8b949e; } + .toolbar label { display: flex; gap: 0.35rem; align-items: center; cursor: pointer; }

Requests ← Dashboard

+
+ + + +
- - + + + + + + + + + @@ -56,7 +82,84 @@

Requests ← Dashboard

if (s == null) return null; try { return JSON.stringify(JSON.parse(s), null, 2); } catch (e) { return s; } } + // Compact token counts: 103673 -> "103.7k". Raw six-digit numbers in a + // dense table are hard to compare at a glance. + function compactTokens(n) { + if (n == null) return '0'; + if (n < 1000) return String(n); + if (n < 1000000) return (n / 1000).toFixed(1) + 'k'; + return (n / 1000000).toFixed(2) + 'M'; + } + // "103.7k (98%)" — total prompt size plus how much of it came from cache. + // The hit rate is the early warning for a broken cache prefix: it should + // sit high and stable, and a sudden drop means the prompt stopped + // matching and every turn is paying full price again. + function formatInput(r) { + const total = r.input_tokens ?? 0; + const read = r.cache_read_input_tokens ?? 0; + if (total === 0) return '0'; + if (read === 0) return compactTokens(total); + const pct = Math.round(read / total * 100); + return `${compactTokens(total)} (${pct}%)`; + } + // `output_tokens` only becomes authoritative when the upstream's closing + // `message_delta` arrives. A stream cut short before that leaves the + // opening placeholder from `message_start` — 3, for a turn that went on to + // emit 35899 — so show a dash rather than a number that reads as fact. + function formatOutput(r) { + if (r.output_tokens_final === false) { + return ''; + } + return r.output_tokens; + } + // Inline SVGs rather than emoji: clipboard/check glyphs render wildly + // differently across platforms, and these inherit the button's color. + const ICON_COPY = ''; + const ICON_DONE = ''; + + function bodyBlock(title, content) { + return `
+
+

${title}

+ +
+
${escapeHtml(content ?? '(none)')}
+
`; + } + + let lastSignature = ''; + + // Longest upstream silence. This is the number that assigns blame when a + // stream stalls: a large value means the upstream went quiet, a small one + // means it kept sending and the stall was downstream of the proxy. + // Highlighted past 60s, which is where intermediaries start dropping + // idle connections. + function formatIdle(r) { + const ms = r.upstream_idle_max_ms; + if (ms == null) return '—'; + const s = ms / 1000; + const txt = s < 10 ? s.toFixed(1) + 's' : Math.round(s) + 's'; + return s >= 60 ? `${txt}` : txt; + } function render(data) { + // Records are immutable once stored, so an identical id list means + // nothing on screen would change. Skipping the rebuild entirely keeps + // text selection and scroll position untouched while the user reads. + const signature = `${data.page}|${data.total}|${(data.items || []).map(r => r.id).join(',')}`; + if (signature === lastSignature) return; + lastSignature = signature; + + // The 5s auto-refresh rebuilds this table wholesale, which resets the + // scroll offset of any expanded body the user is in the middle of + // reading. Snapshot those offsets and put them back afterwards. + const scrollTops = new Map(); + document.querySelectorAll('tr.detail-row').forEach(tr => { + const id = tr.getAttribute('data-detail'); + tr.querySelectorAll('pre').forEach((pre, i) => { + if (pre.scrollTop > 0) scrollTops.set(`${id}:${i}`, pre.scrollTop); + }); + }); + const rows = (data.items || []).map(r => { const ok = r.status_code >= 200 && r.status_code < 400; const model = r.translated_model ? `${r.model} → ${r.translated_model}` : r.model; @@ -65,23 +168,63 @@

Requests ← Dashboard

const hasBody = r.request_body != null || r.response_body != null; const detail = hasBody ? `
-

Request body

${escapeHtml(pretty(r.request_body) ?? '(none)')}
-

Response body

${escapeHtml(pretty(r.response_body) ?? '(none)')}
+ ${bodyBlock('Request body', pretty(r.request_body))} + ${bodyBlock('Response body', pretty(r.response_body))}
` : `
No body captured. Run the proxy with debug enabled (--debug) to record request/response bodies.
`; - return ` + return ` + - - - + + + + + + - `; + `; }).join(''); document.getElementById('rows').innerHTML = rows; + + // Restore the offsets captured above. A finished request's body never + // changes, so the position stays meaningful across refreshes. + document.querySelectorAll('tr.detail-row').forEach(tr => { + const id = tr.getAttribute('data-detail'); + tr.querySelectorAll('pre').forEach((pre, i) => { + const top = scrollTops.get(`${id}:${i}`); + if (top != null) pre.scrollTop = top; + }); + }); + + document.querySelectorAll('.sess').forEach(el => { + el.onclick = (e) => { + e.stopPropagation(); + sessionFilter = el.getAttribute('data-sess').slice(0, 8); + page = 1; lastSignature = ''; load(); + }; + }); + document.querySelectorAll('.copy-btn').forEach(btn => { + btn.onclick = async (e) => { + // The detail row is a sibling of the clickable request row, but stop + // propagation anyway so a future layout change cannot make copying + // collapse the row. + e.stopPropagation(); + const pre = btn.closest('.body-block').querySelector('pre'); + try { + await navigator.clipboard.writeText(pre.textContent); + btn.innerHTML = ICON_DONE; + btn.classList.add('copied'); + setTimeout(() => { btn.innerHTML = ICON_COPY; btn.classList.remove('copied'); }, 1200); + } catch (err) { + btn.title = 'Copy failed: ' + err; + } + }; + }); + document.querySelectorAll('tr.req-row').forEach(tr => { tr.onclick = () => { const id = tr.getAttribute('data-id'); @@ -96,11 +239,30 @@

Requests ← Dashboard

document.getElementById('prev').disabled = page <= 1; document.getElementById('next').disabled = page >= (data.total_pages || 1); } + let sessionFilter = null; async function load() { - const res = await fetch(`/api/requests?page=${page}&per_page=${perPage}`); + const failed = document.getElementById('failedOnly').checked; + const sq = sessionFilter ? `&session=${encodeURIComponent(sessionFilter)}` : ''; + document.getElementById('sessFilter').innerHTML = sessionFilter + ? `session ${escapeHtml(sessionFilter)}clear` + : ''; + const res = await fetch(`/api/requests?page=${page}&per_page=${perPage}${failed ? '&failed=true' : ''}${sq}`); const data = await res.json(); + document.getElementById('filterHint').textContent = + failed ? `${data.total} failed request(s) — non-2xx, interrupted streams, and client disconnects` : ''; render(data); } + document.addEventListener('click', (e) => { + if (e.target && e.target.id === 'clearSess') { + e.preventDefault(); + sessionFilter = null; page = 1; lastSignature = ''; load(); + } + }); + document.getElementById('failedOnly').onchange = () => { + page = 1; + lastSignature = ''; // force a rebuild; the filter changed, not the data + load(); + }; document.getElementById('prev').onclick = () => { if (page > 1) { page--; load(); } }; document.getElementById('next').onclick = () => { page++; load(); }; load(); diff --git a/scripts/__pycache__/replay.cpython-313.pyc b/scripts/__pycache__/replay.cpython-313.pyc new file mode 100644 index 0000000..27145d0 Binary files /dev/null and b/scripts/__pycache__/replay.cpython-313.pyc differ diff --git a/scripts/replay.py b/scripts/replay.py new file mode 100644 index 0000000..8d25714 --- /dev/null +++ b/scripts/replay.py @@ -0,0 +1,281 @@ +#!/usr/bin/env python3 +"""Replay a request captured by the ghc-proxy dashboard and time the upstream SSE stream. + +Unlike Claude Code, this has no stall watchdog: it waits for the upstream to +finish no matter how long it goes silent, so it can answer two questions the +client can never answer for itself -- how long the silence actually lasts, and +what the upstream eventually sends. + + python scripts/replay.py --list + python scripts/replay.py --id --target upstream + python scripts/replay.py --id --target proxy --max-tokens 4000 +""" +import argparse +import json +import os +import sys +import time +import uuid + +import requests + +DASHBOARD = "http://127.0.0.1:8314" +CFG_DIR = os.path.join(os.environ.get("APPDATA", ""), "ghc-tunnel") +GITHUB_API = "https://api.github.com" + + +def load_config(): + """Reads the few config.yaml values that appear in Copilot request headers. + + Parsed by hand rather than with PyYAML so the script has no dependency the + proxy itself doesn't already imply. + """ + cfg = {} + path = os.path.join(CFG_DIR, "config.yaml") + with open(path, encoding="utf-8") as fh: + for line in fh: + line = line.strip() + if not line or line.startswith("#") or ":" not in line: + continue + k, _, v = line.partition(":") + cfg[k.strip()] = v.strip().strip('"').strip("'") + return cfg + + +def base_url(cfg): + acct = cfg.get("account_type", "individual") + if acct == "individual": + return "https://api.githubcopilot.com" + return f"https://api.{acct}.githubcopilot.com" + + +def copilot_token(cfg): + """Exchanges the stored GitHub token for a short-lived Copilot token. + + Mirrors src/auth.rs::fetch_copilot_token. The proxy holds its own copy in + memory and never exposes it, so the exchange is repeated here. + """ + with open(os.path.join(CFG_DIR, "github_token.txt"), encoding="utf-8") as fh: + gh = fh.read().strip() + r = requests.get( + f"{GITHUB_API}/copilot_internal/v2/token", + headers={ + "Authorization": f"token {gh}", + "Editor-Version": f"vscode/{cfg.get('vscode_version', '1.123.0')}", + "User-Agent": "GithubCopilot/1.155.0", + }, + timeout=30, + ) + r.raise_for_status() + return r.json()["token"] + + +def upstream_headers(cfg, tok, machine_id, beta): + """Replicates SharedState::copilot_headers so the replay looks identical.""" + rid = str(uuid.uuid4()) + h = { + "Authorization": f"Bearer {tok}", + "Content-Type": "application/json", + "Copilot-Integration-Id": "vscode-chat", + "Editor-Version": f"vscode/{cfg.get('vscode_version', '1.123.0')}", + "Editor-Plugin-Version": f"copilot-chat/{cfg.get('copilot_version', '0.48.1')}", + "User-Agent": f"GitHubCopilotChat/{cfg.get('copilot_version', '0.48.1')}", + "OpenAI-Intent": "conversation-panel", + "openai-organization": "github-copilot", + "vscode-machineid": machine_id, + "vscode-sessionid": str(uuid.uuid4()), + "X-GitHub-Api-Version": cfg.get("api_version", "2025-05-01"), + "X-Interaction-Type": "conversation-panel", + "X-Request-Id": rid, + "X-Agent-Task-Id": rid, + "X-VSCode-User-Agent-Library-Version": "electron-fetch", + "anthropic-version": "2023-06-01", + } + if beta: + h["anthropic-beta"] = beta + return h + + +def derive_beta(payload): + """Builds the `anthropic-beta` header the proxy would have sent. + + Mirrors server.rs::apply_anthropic_beta. `context_management` in the body + is rejected with a misleading `Extra inputs are not permitted` 400 unless + its beta is requested, so it has to be derived rather than hardcoded. + """ + betas = ["claude-code-20250219", "context-1m-2025-08-07"] + if "context_management" in payload: + betas.append("context-management-2025-06-27") + return ",".join(betas) + + +def fetch_records(limit=60): + r = requests.get(f"{DASHBOARD}/api/requests?per_page={limit}", timeout=30) + r.raise_for_status() + return r.json()["items"] + + +def parse_sse(raw_events): + """Splits a list of (t, bytes) chunks into (t, event_name, data) triples. + + A `ping` is two lines -- `event: ping` then its own `data:` line. Counting + that data line as content is what made the first run of this experiment + report a 15s maximum gap when the real gap was 455s. + """ + buf = b"" + out = [] + for t, chunk in raw_events: + buf += chunk + while True: + idx = buf.find(b"\n\n") + idx_crlf = buf.find(b"\r\n\r\n") + if idx < 0 and idx_crlf < 0: + break + if idx < 0 or (0 <= idx_crlf < idx): + idx, width = idx_crlf, 4 + else: + width = 2 + block, buf = buf[:idx], buf[idx + width:] + name, data = None, None + for line in block.split(b"\n"): + line = line.rstrip(b"\r") + if line.startswith(b"event:"): + name = line[6:].strip().decode("utf-8", "replace") + elif line.startswith(b"data:"): + data = line[5:].strip().decode("utf-8", "replace") + if name or data: + out.append((t, name, data)) + return out + + +def run(payload, url, headers, label, dump): + print(f"\n=== {label} ===") + print(f"POST {url}") + body = json.dumps(payload).encode("utf-8") + print(f"payload : {len(body)} bytes") + print(f"model : {payload.get('model')}") + print(f"max_tokens : {payload.get('max_tokens')}") + print(f"messages : {len(payload.get('messages') or [])}") + print(f"tools : {len(payload.get('tools') or [])}") + print("waiting (no client-side stall watchdog -- will wait indefinitely)...\n", flush=True) + + t0 = time.monotonic() + raw = [] + # `timeout` is the connect/read-inactivity timeout. Deliberately generous: + # the whole point is to outlast a silence that kills the real client. + resp = requests.post(url, headers=headers, data=body, stream=True, timeout=(30, 1800)) + t_headers = time.monotonic() - t0 + print(f"HTTP {resp.status_code} (headers after {t_headers:.1f}s)") + print(f"content-type : {resp.headers.get('content-type')}") + if resp.status_code >= 300: + print("--- error body ---") + print(resp.text[:4000]) + return + + for chunk in resp.iter_content(chunk_size=None): + if chunk: + raw.append((time.monotonic() - t0, chunk)) + total = time.monotonic() - t0 + + events = parse_sse(raw) + content_ts = [t for t, name, _ in events if name and name != "ping"] + all_ts = [t for t, _ in raw] + + def max_gap(ts): + if not ts: + return total, 0.0 + pts = [0.0] + ts + gaps = [(pts[i + 1] - pts[i], pts[i]) for i in range(len(pts) - 1)] + gaps.append((total - ts[-1], ts[-1])) + return max(gaps) + + gap_content, gap_at = max_gap(content_ts) + gap_bytes, _ = max_gap(all_ts) + + counts, tool_json, stop_reason, usage = {}, 0, None, None + for _, name, data in events: + if name: + counts[name] = counts.get(name, 0) + 1 + if not data: + continue + try: + j = json.loads(data) + except ValueError: + continue + d = j.get("delta") or {} + if d.get("type") == "input_json_delta": + tool_json += len(d.get("partial_json") or "") + if d.get("stop_reason"): + stop_reason = d["stop_reason"] + if j.get("usage"): + usage = j["usage"] + if (j.get("message") or {}).get("usage"): + usage = j["message"]["usage"] + + print("\n--- 汇总 ---") + print(f"总耗时 {total:.1f}s") + print(f"TTFB (首字节) {all_ts[0]:.1f}s" if all_ts else "TTFB n/a") + print(f"最长静默(内容事件) {gap_content:.1f}s 起于 t={gap_at:.1f}s") + print(f"最长静默(任意字节) {gap_bytes:.1f}s") + print(f"ping 事件 {counts.get('ping', 0)}") + print(f"事件计数 {counts}") + print(f"tool 参数长度 {tool_json} 字符") + print(f"stop_reason {stop_reason}") + print(f"usage {usage}") + verdict = "会被 abort" if gap_content > 300 else "不会被 abort" + print(f"\n>>> Claude Code 300s 阈值: {verdict} (最长内容静默 {gap_content:.1f}s)") + + if dump: + with open(dump, "w", encoding="utf-8") as fh: + for t, name, data in events: + fh.write(f"{t:9.3f} {name or '-':24s} {(data or '')[:400]}\n") + print(f"\n逐事件时序已写入 {dump}") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--list", action="store_true", help="list recent records") + ap.add_argument("--id", help="record id to replay") + ap.add_argument("--target", choices=["upstream", "proxy"], default="upstream") + ap.add_argument("--max-tokens", type=int, help="override max_tokens") + ap.add_argument("--dump", help="write per-event timeline to this file") + args = ap.parse_args() + + records = fetch_records() + if args.list or not args.id: + print(f"{'id':38s} {'time':9s} {'in':>8s} {'out':>7s} {'idle':>7s} {'status':>6s} kind") + for it in records: + print("%-38s %s %8d %7d %6.1fs %6s %s" % ( + it.get("id", "?"), it["timestamp"][11:19], it["input_tokens"], + it["output_tokens"], (it.get("upstream_idle_max_ms") or 0) / 1000, + it.get("status_code"), it.get("failure_kind") or "")) + return + + rec = next((r for r in records if str(r.get("id")) == args.id), None) + if rec is None: + sys.exit(f"record {args.id} not found in the last {len(records)} records") + if not rec.get("request_body"): + sys.exit("record has no request_body -- is debug enabled in config.yaml?") + + payload = json.loads(rec["request_body"]) + if args.max_tokens: + payload["max_tokens"] = args.max_tokens + payload["stream"] = True + + cfg = load_config() + label = f"replay {args.id} -> {args.target}" + if args.target == "proxy": + url = f"{DASHBOARD}/v1/messages" + headers = {"Content-Type": "application/json", "anthropic-version": "2023-06-01"} + else: + with open(os.path.join(CFG_DIR, "machine_id.txt"), encoding="utf-8") as fh: + machine_id = fh.read().strip() + beta = derive_beta(payload) + print(f"anthropic-beta: {beta}") + headers = upstream_headers(cfg, copilot_token(cfg), machine_id, beta) + url = f"{base_url(cfg)}/v1/messages" + run(payload, url, headers, label, args.dump) + + +if __name__ == "__main__": + main() diff --git a/src/server.rs b/src/server.rs index fc90242..b1752d2 100644 --- a/src/server.rs +++ b/src/server.rs @@ -8,6 +8,7 @@ use crate::state::SharedState; use crate::store::RequestRecord; use crate::translate; use crate::util; +use crate::util::TokenUsage; use axum::{ body::{Body, Bytes}, extract::{DefaultBodyLimit, Path, Query, Request, State}, @@ -335,12 +336,30 @@ fn model_rates(model: &str) -> (f64, f64) { } } +/// Cache writes carry a premium over fresh input, and cache reads a steep +/// discount. These are Anthropic's published multipliers for 5-minute +/// ephemeral caching, which is the tier Copilot passes through. +const CACHE_WRITE_RATE_MULTIPLIER: f64 = 1.25; +const CACHE_READ_RATE_MULTIPLIER: f64 = 0.1; + /// Calculate estimated cost in USD based on token counts and model. /// Rates are approximate public list prices per 1K tokens and are only used for /// the dashboard/metrics estimate, never for billing. -fn calculate_cost(model: &str, input_tokens: u64, output_tokens: u64) -> f64 { +/// +/// `usage.input_tokens` is the true total prompt size, so the cached buckets +/// are carved back out of it and repriced: billing every cached token at the +/// full input rate would overstate a Claude Code turn several times over, +/// while billing only the uncached remainder — as this did before — understates +/// it by orders of magnitude. +fn calculate_cost(model: &str, usage: &TokenUsage) -> f64 { let (input_rate, output_rate) = model_rates(model); - (input_tokens as f64 * input_rate + output_tokens as f64 * output_rate) / 1000.0 + let cached = usage.cache_read_input_tokens + usage.cache_creation_input_tokens; + let uncached = usage.input_tokens.saturating_sub(cached); + (uncached as f64 * input_rate + + usage.cache_creation_input_tokens as f64 * input_rate * CACHE_WRITE_RATE_MULTIPLIER + + usage.cache_read_input_tokens as f64 * input_rate * CACHE_READ_RATE_MULTIPLIER + + usage.output_tokens as f64 * output_rate) + / 1000.0 } /// Checks if a request is eligible for prompt caching (system prompt is large enough). @@ -368,25 +387,328 @@ fn is_prompt_cache_eligible(req: &Value) -> bool { } } -/// Detect if a response used prompt caching by checking for cache tokens. -#[allow(dead_code)] -fn extract_prompt_cache_hit(response: &Value) -> Option { - let usage = response.get("usage")?; - let cache_read = usage - .get("cache_read_input_tokens") - .and_then(|t| t.as_u64()) - .unwrap_or(0); - let cache_creation = usage - .get("cache_creation_input_tokens") - .and_then(|t| t.as_u64()) - .unwrap_or(0); - - if cache_read > 0 { - Some(true) // Cache hit (reading from cache) - } else if cache_creation > 0 { - Some(false) // Cache write (creating new cache) +/// Ratio of `part` to `whole` as a **0–1 fraction** rounded to two decimals. +/// +/// Every rate the audit API reports goes through here so they all share one +/// scale. Mixing a 0–1 fraction with a 0–100 percentage under two +/// similarly-named JSON keys is the kind of thing that silently corrupts a +/// dashboard downstream. Returns `0.0` for a zero denominator instead of the +/// `NaN` that would serialize as `null`. +fn ratio_2dp(part: f64, whole: f64) -> f64 { + if whole <= 0.0 { + return 0.0; + } + (part / whole * 100.0).round() / 100.0 +} + +/// Extracts the client session id from a request's `metadata.user_id`. +/// +/// Claude Code sends that field as a JSON *string* containing another JSON +/// object (`{"device_id":…,"session_id":…}`), so it takes two parses. Several +/// client instances routinely share one proxy, and without this their records +/// interleave in the dashboard with nothing to tell them apart. +/// +/// Every unexpected shape degrades to `None` — this is diagnostic metadata and +/// must never be able to fail a request. +fn extract_session_id(req: &Value) -> Option { + let raw = req.get("metadata")?.get("user_id")?.as_str()?; + let parsed: Value = serde_json::from_str(raw).ok()?; + let id = parsed.get("session_id")?.as_str()?; + (!id.is_empty()).then(|| id.to_string()) +} + +/// Records a request that failed before producing a usable response. +/// +/// Every early `return` on an inference path routes through here. Without it +/// the request vanishes entirely: records were only ever added on the success +/// path, so a connection error or a 429 left nothing in the dashboard and +/// nothing to diagnose from. +#[allow(clippy::too_many_arguments)] +fn record_failure( + state: &SharedState, + endpoint: &str, + model: &str, + translated: Option<&str>, + status: u16, + kind: &str, + req_size: usize, + req_body: Option, + resp_body: Option, + start: Instant, + session_id: Option, +) { + state.store.add(RequestRecord { + id: uuid::Uuid::new_v4().to_string(), + timestamp: now_iso(), + endpoint: endpoint.to_string(), + model: model.to_string(), + translated_model: translated.filter(|t| *t != model).map(String::from), + status_code: status, + request_size: req_size, + response_size: resp_body.as_ref().map(|s| s.len()).unwrap_or(0), + input_tokens: 0, + output_tokens: 0, + output_tokens_final: None, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + premium_multiplier: None, + upstream_idle_max_ms: None, + keepalive_probes: None, + duration: elapsed_secs(start), + request_body: req_body, + // Kept regardless of debug mode: an error body is small and is the + // entire reason this record exists. + response_body: resp_body, + message_count: None, + tool_count: None, + tool_names: None, + stop_reason: None, + tools_called: None, + is_agent_initiated: None, + prompt_cache_hit: None, + session_id, + failure_kind: Some(kind.to_string()), + estimated_cost_usd: None, + }); +} + +/// Guarantees a streaming request is recorded on every exit path, including +/// the one that used to lose it entirely. +/// +/// When a client hangs up, axum drops the response body, which drops the +/// `async_stream` generator — so a `store.add` written after the loop never +/// runs. Holding the record in a value with a `Drop` impl moves the write onto +/// an unconditional path. `RequestStore::add` takes a synchronous +/// `std::sync::Mutex`, so it is safe to call from `Drop`; anything async (such +/// as the premium multiplier) must be resolved before the stream starts. +/// Running state folded out of the upstream Anthropic SSE stream on the +/// passthrough path (`stream_anthropic_direct`). +/// +/// Kept as a plain value rather than as generator locals so the whole +/// event-interpretation rulebook is testable without standing up a stream. +#[derive(Debug, Default)] +struct DirectStreamState { + usage: TokenUsage, + /// Whether `message_delta` — the only event carrying an authoritative + /// `output_tokens` — has arrived. Until it does, the count taken from + /// `message_start` is an opening placeholder, so a record finalized early + /// must not present it as the turn's real output. + usage_final: bool, + stop_reason: Option, + tools_called: Vec, + saw_message_stop: bool, +} + +impl DirectStreamState { + /// Folds one decoded SSE event into the running state. + fn observe(&mut self, v: &Value) { + match v.get("type").and_then(|t| t.as_str()) { + Some("message_start") => { + // The input buckets are only ever stated here, and + // `input_tokens` alone is the uncached remainder — for a + // cached Claude Code turn that is single digits against a + // six-figure prompt. + if let Some(u) = v.get("message").and_then(|m| m.get("usage")) { + self.usage = util::anthropic_usage(u); + } + } + Some("message_delta") => { + // `message_delta` restates the running output count and + // nothing else; adopting it wholesale would zero the input + // buckets captured at message_start. + if let Some(o) = v + .get("usage") + .and_then(|u| u.get("output_tokens")) + .and_then(|t| t.as_u64()) + { + self.usage.output_tokens = o; + self.usage_final = true; + } + if let Some(sr) = v + .get("delta") + .and_then(|d| d.get("stop_reason")) + .and_then(|s| s.as_str()) + { + self.stop_reason = Some(sr.to_string()); + } + } + Some("content_block_start") => { + if let Some(block) = v.get("content_block") { + if block.get("type").and_then(|t| t.as_str()) == Some("tool_use") { + if let Some(name) = block.get("name").and_then(|n| n.as_str()) { + if !self.tools_called.iter().any(|t| t == name) { + self.tools_called.push(name.to_string()); + } + } + } + } + } + Some("message_stop") => self.saw_message_stop = true, + _ => {} + } + } +} + +struct StreamRecorder { + state: SharedState, + /// Taken by the first finalize, so `Drop` cannot record a second time. + rec: Option, + model_for_cost: String, + start: Instant, + /// Running counters, kept here rather than as generator locals so `Drop` + /// can still see how far the stream got. + usage: TokenUsage, + /// Mirrors `DirectStreamState::usage_final`, for the same reason: `Drop` + /// must be able to tell a record cut short mid-stream — whose output count + /// is still the `message_start` placeholder — from one that saw the real + /// total in `message_delta`. + usage_final: bool, + resp_size: usize, + idle: util::IdleTracker, + /// Shared with the keepalive layer wrapping this stream, so the record can + /// state how many probes actually went out during an upstream silence. + probes: std::sync::Arc, + /// Raw upstream bytes seen so far, for the same reason: a client + /// disconnect is exactly when you want to read what did arrive, and a + /// generator local would be gone by then. Only filled in debug mode. + debug_raw: Vec, +} + +impl StreamRecorder { + #[allow(clippy::too_many_arguments)] + fn new( + state: SharedState, + endpoint: &str, + model: String, + translated: String, + req_size: usize, + req_body: Option, + premium_multiplier: Option, + start: Instant, + probes: std::sync::Arc, + session_id: Option, + ) -> Self { + let rec = RequestRecord { + id: uuid::Uuid::new_v4().to_string(), + timestamp: now_iso(), + endpoint: endpoint.to_string(), + translated_model: (translated != model).then(|| translated.clone()), + model, + status_code: 0, + request_size: req_size, + response_size: 0, + input_tokens: 0, + output_tokens: 0, + output_tokens_final: None, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + premium_multiplier, + upstream_idle_max_ms: None, + keepalive_probes: None, + duration: 0.0, + request_body: req_body, + response_body: None, + message_count: None, + tool_count: None, + tool_names: None, + stop_reason: None, + tools_called: None, + is_agent_initiated: None, + prompt_cache_hit: None, + session_id, + failure_kind: None, + estimated_cost_usd: None, + }; + StreamRecorder { + state, + rec: Some(rec), + model_for_cost: translated, + start, + usage: TokenUsage::default(), + usage_final: false, + resp_size: 0, + idle: util::IdleTracker::new(Instant::now()), + probes, + debug_raw: Vec::new(), + } + } + + /// Mutable access to the pending record for path-specific fields the + /// generic finalize does not cover (tool names, message counts, …). + fn record_mut(&mut self) -> Option<&mut RequestRecord> { + self.rec.as_mut() + } + + /// Folds the running counters into the record and hands it to the store. + /// Idempotent — the second call (typically from `Drop`) is a no-op. + fn finalize( + &mut self, + status: u16, + failure_kind: Option<&str>, + stop_reason: Option, + response_body: Option, + ) { + let Some(mut rec) = self.rec.take() else { + return; + }; + rec.status_code = status; + rec.response_size = self.resp_size; + rec.input_tokens = self.usage.input_tokens; + rec.output_tokens = self.usage.output_tokens; + rec.output_tokens_final = Some(self.usage_final); + rec.cache_read_input_tokens = self.usage.cache_read_input_tokens; + rec.cache_creation_input_tokens = self.usage.cache_creation_input_tokens; + rec.upstream_idle_max_ms = Some(self.idle.max_idle_ms_including_now()); + rec.keepalive_probes = Some(self.probes.load(std::sync::atomic::Ordering::Relaxed)); + rec.duration = elapsed_secs(self.start); + rec.prompt_cache_hit = cache_disposition(&self.usage); + rec.estimated_cost_usd = Some(calculate_cost(&self.model_for_cost, &self.usage)); + rec.stop_reason = stop_reason; + rec.response_body = response_body; + rec.failure_kind = failure_kind.map(String::from); + self.state.store.add(rec); + } +} + +impl Drop for StreamRecorder { + fn drop(&mut self) { + // Only reachable when the generator was dropped before finalize ran, + // i.e. the client went away mid-stream. 499 follows nginx's + // "client closed request". Whatever had already arrived from the + // upstream is preserved — that partial body is the main evidence for + // diagnosing why the client gave up. + let partial = (!self.debug_raw.is_empty()) + .then(|| String::from_utf8_lossy(&self.debug_raw).into_owned()); + self.finalize( + 499, + Some(crate::store::failure::CLIENT_DISCONNECTED), + None, + partial, + ); + } +} + +/// Whether an upstream response should be forwarded to the client as an SSE +/// stream. +/// +/// A non-2xx upstream returns a JSON error body, not SSE. Wrapping one in a +/// 200 `text/event-stream` produces a stream that never yields a single event, +/// which clients report as a stalled or hung response instead of the actual +/// auth/quota failure underneath. Every streaming path must gate on this. +fn is_streamable_status(status: u16) -> bool { + (200..300).contains(&status) +} + +/// Whether a request touched the prompt cache: `Some(true)` when part of the +/// prompt was served from cache, `Some(false)` on a write-only turn (a cold +/// cache being populated), and `None` when caching was not involved at all. +fn cache_disposition(usage: &TokenUsage) -> Option { + if usage.cache_read_input_tokens > 0 { + Some(true) + } else if usage.cache_creation_input_tokens > 0 { + Some(false) } else { - None // No caching + None } } @@ -674,17 +996,10 @@ async fn chat_completions(State(state): State, body: Bytes) -> Resp log_debug_response(&state, "/v1/chat/completions", &text); if status.is_success() { let parsed: Value = serde_json::from_str(&text).unwrap_or(Value::Null); - let usage = parsed.get("usage").cloned().unwrap_or(json!({})); - let input_tokens = usage - .get("prompt_tokens") - .and_then(|t| t.as_u64()) - .unwrap_or(0); - let output_tokens = usage - .get("completion_tokens") - .and_then(|t| t.as_u64()) - .unwrap_or(0); + let usage = util::openai_usage(&parsed.get("usage").cloned().unwrap_or(json!({}))); let (tool_count, tool_names) = extract_tools_from_request(&req); - let cost = calculate_cost(&translated, input_tokens, output_tokens); + let cost = calculate_cost(&translated, &usage); + let premium_multiplier = state.model_premium_multiplier(&translated).await; state.store.add(RequestRecord { id: uuid::Uuid::new_v4().to_string(), timestamp: now_iso(), @@ -694,8 +1009,14 @@ async fn chat_completions(State(state): State, body: Bytes) -> Resp status_code: status.as_u16(), request_size: req_size, response_size: resp_size, - input_tokens, - output_tokens, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + output_tokens_final: None, + cache_read_input_tokens: usage.cache_read_input_tokens, + cache_creation_input_tokens: usage.cache_creation_input_tokens, + premium_multiplier, + upstream_idle_max_ms: None, + keepalive_probes: None, duration: elapsed_secs(start), request_body: capture_json(&state, &req), response_body: capture_str(&state, &text), @@ -705,7 +1026,9 @@ async fn chat_completions(State(state): State, body: Bytes) -> Resp stop_reason: None, // OpenAI responses don't have stop_reason in JSON tools_called: None, is_agent_initiated: Some(agent), - prompt_cache_hit: None, + session_id: None, + prompt_cache_hit: cache_disposition(&usage), + failure_kind: None, estimated_cost_usd: Some(cost), }); Json(parsed).into_response() @@ -821,17 +1144,10 @@ async fn responses(State(state): State, body: Bytes) -> Response { log_debug_response(&state, "/v1/responses", &text); if status.is_success() { let parsed: Value = serde_json::from_str(&text).unwrap_or(Value::Null); - let usage = parsed.get("usage").cloned().unwrap_or(json!({})); - let input_tokens = usage - .get("input_tokens") - .and_then(|t| t.as_u64()) - .unwrap_or(0); - let output_tokens = usage - .get("output_tokens") - .and_then(|t| t.as_u64()) - .unwrap_or(0); + let usage = util::responses_usage(&parsed.get("usage").cloned().unwrap_or(json!({}))); let (tool_count, tool_names) = extract_tools_from_request(&req); - let cost = calculate_cost(&translated, input_tokens, output_tokens); + let cost = calculate_cost(&translated, &usage); + let premium_multiplier = state.model_premium_multiplier(&translated).await; state.store.add(RequestRecord { id: uuid::Uuid::new_v4().to_string(), timestamp: now_iso(), @@ -841,8 +1157,14 @@ async fn responses(State(state): State, body: Bytes) -> Response { status_code: status.as_u16(), request_size: req_size, response_size: resp_size, - input_tokens, - output_tokens, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + output_tokens_final: None, + cache_read_input_tokens: usage.cache_read_input_tokens, + cache_creation_input_tokens: usage.cache_creation_input_tokens, + premium_multiplier, + upstream_idle_max_ms: None, + keepalive_probes: None, duration: elapsed_secs(start), request_body: capture_json(&state, &req), response_body: capture_str(&state, &text), @@ -852,7 +1174,9 @@ async fn responses(State(state): State, body: Bytes) -> Response { stop_reason: None, tools_called: None, is_agent_initiated: Some(agent), - prompt_cache_hit: None, + session_id: None, + prompt_cache_hit: cache_disposition(&usage), + failure_kind: None, estimated_cost_usd: Some(cost), }); Json(parsed).into_response() @@ -994,7 +1318,22 @@ async fn messages_direct( .await; let upstream = match upstream { Ok(r) => r, - Err(e) => return anthropic_error(StatusCode::GATEWAY_TIMEOUT, e.to_string()), + Err(e) => { + record_failure( + &state, + "/v1/messages", + &original_model, + Some(&translated), + 504, + crate::store::failure::CONNECT, + req_size, + capture_json(&state, &sanitized), + Some(e.to_string()), + start, + extract_session_id(&sanitized), + ); + return anthropic_error(StatusCode::GATEWAY_TIMEOUT, e.to_string()); + } }; let status = upstream.status(); // Inspect 400 responses so we can transparently recover from the @@ -1013,16 +1352,32 @@ async fn messages_direct( continue; } } + record_failure( + &state, + "/v1/messages", + &original_model, + Some(&translated), + status.as_u16(), + crate::store::failure::UPSTREAM_STATUS, + req_size, + capture_json(&state, &sanitized), + Some(text.clone()), + start, + extract_session_id(&sanitized), + ); return passthrough_error(status, text); } return stream_anthropic_direct( state.clone(), upstream, - original_model, - translated, - req_size, - capture_json(&state, &sanitized), - start, + RequestMeta { + req_body: capture_json(&state, &sanitized), + session_id: extract_session_id(&sanitized), + original_model, + translated, + req_size, + start, + }, ) .await; } @@ -1030,6 +1385,19 @@ async fn messages_direct( let resp = util::post_with_retry(&state, &url, headers.clone(), payload, "/v1/messages").await; let Some(resp) = resp else { + record_failure( + &state, + "/v1/messages", + &original_model, + Some(&translated), + 504, + crate::store::failure::CONNECT, + req_size, + capture_json(&state, &sanitized), + None, + start, + extract_session_id(&sanitized), + ); return anthropic_error( StatusCode::GATEWAY_TIMEOUT, "Upstream connection error".into(), @@ -1044,14 +1412,7 @@ async fn messages_direct( &serde_json::to_string(&parsed).unwrap_or_default(), ); let usage = parsed.get("usage").cloned().unwrap_or(json!({})); - let input_tokens = usage - .get("input_tokens") - .and_then(|t| t.as_u64()) - .unwrap_or(0); - let output_tokens = usage - .get("output_tokens") - .and_then(|t| t.as_u64()) - .unwrap_or(0); + let usage = util::anthropic_usage(&usage); let (tool_count, tool_names) = extract_tools_from_request(&req); let tools_called: Vec = parsed .get("content") @@ -1080,8 +1441,14 @@ async fn messages_direct( status_code: status.as_u16(), request_size: req_size, response_size: serde_json::to_vec(&parsed).map(|v| v.len()).unwrap_or(0), - input_tokens, - output_tokens, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + output_tokens_final: None, + cache_read_input_tokens: usage.cache_read_input_tokens, + cache_creation_input_tokens: usage.cache_creation_input_tokens, + premium_multiplier: state.model_premium_multiplier(&translated).await, + upstream_idle_max_ms: None, + keepalive_probes: None, duration: elapsed_secs(start), request_body: capture_json(&state, &sanitized), response_body: capture_json(&state, &parsed), @@ -1091,8 +1458,10 @@ async fn messages_direct( stop_reason, tools_called: (!tools_called.is_empty()).then_some(tools_called), is_agent_initiated: Some(agent), - prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&translated, input_tokens, output_tokens)), + session_id: None, + prompt_cache_hit: cache_disposition(&usage), + failure_kind: None, + estimated_cost_usd: Some(calculate_cost(&translated, &usage)), }); return Json(parsed).into_response(); } @@ -1119,6 +1488,19 @@ async fn messages_direct( continue; } } + record_failure( + &state, + "/v1/messages", + &original_model, + Some(&translated), + status.as_u16(), + crate::store::failure::UPSTREAM_STATUS, + req_size, + capture_json(&state, &sanitized), + Some(text.clone()), + start, + extract_session_id(&sanitized), + ); return passthrough_error(status, text); } anthropic_error(StatusCode::BAD_GATEWAY, "Exhausted retries".into()) @@ -1169,6 +1551,7 @@ async fn messages_translated( translated, req_size, start, + extract_session_id(&openai_req), ) .await; } @@ -1177,6 +1560,19 @@ async fn messages_translated( util::post_with_retry(&state, &url, headers, payload, "/v1/messages (translated)") .await; let Some(resp) = resp else { + record_failure( + &state, + "/v1/messages", + &original_model, + Some(&translated), + 504, + crate::store::failure::CONNECT, + req_size, + capture_json(&state, &openai_req), + None, + start, + extract_session_id(&openai_req), + ); return anthropic_error( StatusCode::GATEWAY_TIMEOUT, "Upstream connection error".into(), @@ -1191,15 +1587,7 @@ async fn messages_translated( "/v1/messages", &serde_json::to_string(&parsed).unwrap_or_default(), ); - let usage = parsed.get("usage").cloned().unwrap_or(json!({})); - let input_tokens = usage - .get("prompt_tokens") - .and_then(|t| t.as_u64()) - .unwrap_or(0); - let output_tokens = usage - .get("completion_tokens") - .and_then(|t| t.as_u64()) - .unwrap_or(0); + let usage = util::openai_usage(&parsed.get("usage").cloned().unwrap_or(json!({}))); let (tool_count, tool_names) = extract_tools_from_request(&openai_req); state.store.add(RequestRecord { id: uuid::Uuid::new_v4().to_string(), @@ -1212,8 +1600,14 @@ async fn messages_translated( response_size: serde_json::to_vec(&anthropic_resp) .map(|v| v.len()) .unwrap_or(0), - input_tokens, - output_tokens, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + output_tokens_final: None, + cache_read_input_tokens: usage.cache_read_input_tokens, + cache_creation_input_tokens: usage.cache_creation_input_tokens, + premium_multiplier: state.model_premium_multiplier(&translated).await, + upstream_idle_max_ms: None, + keepalive_probes: None, duration: elapsed_secs(start), request_body: capture_json(&state, &openai_req), response_body: capture_json(&state, &parsed), @@ -1223,8 +1617,10 @@ async fn messages_translated( stop_reason: None, tools_called: None, is_agent_initiated: Some(agent), - prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&translated, input_tokens, output_tokens)), + session_id: None, + prompt_cache_hit: cache_disposition(&usage), + failure_kind: None, + estimated_cost_usd: Some(calculate_cost(&translated, &usage)), }); return Json(anthropic_resp).into_response(); } @@ -1241,6 +1637,19 @@ async fn messages_translated( } } } + record_failure( + &state, + "/v1/messages", + &original_model, + Some(&translated), + status.as_u16(), + crate::store::failure::UPSTREAM_STATUS, + req_size, + capture_json(&state, &openai_req), + Some(text.clone()), + start, + extract_session_id(&openai_req), + ); return passthrough_error(status, text); } anthropic_error(StatusCode::BAD_GATEWAY, "Exhausted retries".into()) @@ -1438,16 +1847,9 @@ async fn gemini_generate( if status.is_success() { let parsed: Value = serde_json::from_str(&text).unwrap_or(Value::Null); let gemini_resp = gemini::openai_to_gemini(&parsed); - let usage = parsed.get("usage").cloned().unwrap_or(json!({})); - let input_tokens = usage - .get("prompt_tokens") - .and_then(|t| t.as_u64()) - .unwrap_or(0); - let output_tokens = usage - .get("completion_tokens") - .and_then(|t| t.as_u64()) - .unwrap_or(0); - let cost = calculate_cost(&translated, input_tokens, output_tokens); + let usage = util::openai_usage(&parsed.get("usage").cloned().unwrap_or(json!({}))); + let cost = calculate_cost(&translated, &usage); + let premium_multiplier = state.model_premium_multiplier(&translated).await; state.store.add(RequestRecord { id: uuid::Uuid::new_v4().to_string(), timestamp: now_iso(), @@ -1457,8 +1859,14 @@ async fn gemini_generate( status_code: status.as_u16(), request_size: req_size, response_size: resp_size, - input_tokens, - output_tokens, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + output_tokens_final: None, + cache_read_input_tokens: usage.cache_read_input_tokens, + cache_creation_input_tokens: usage.cache_creation_input_tokens, + premium_multiplier, + upstream_idle_max_ms: None, + keepalive_probes: None, duration: elapsed_secs(start), request_body: capture_json(&state, &openai_req), response_body: capture_str(&state, &text), @@ -1468,7 +1876,9 @@ async fn gemini_generate( stop_reason: None, tools_called: None, is_agent_initiated: Some(agent), - prompt_cache_hit: None, + session_id: None, + prompt_cache_hit: cache_disposition(&usage), + failure_kind: None, estimated_cost_usd: Some(cost), }); Json(gemini_resp).into_response() @@ -1535,7 +1945,7 @@ async fn stream_gemini( }; let status = upstream.status().as_u16(); // Surface a non-2xx upstream (JSON error, not SSE) as a normal error. - if !(200..300).contains(&status) { + if !is_streamable_status(status) { let text = upstream.text().await.unwrap_or_default(); log_debug_response(&state, "/v1beta/models", &text); log_error( @@ -1550,22 +1960,26 @@ async fn stream_gemini( ); } let model_json = Value::String(translated.clone()); + // Shared with the keepalive wrapper so the record can report how + // many probes actually went out during an upstream silence. + let probe_count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)); let stream = async_stream::stream! { use futures_util::StreamExt; let mut byte_stream = upstream.bytes_stream(); let mut lines = util::SseLineBuffer::new(); - let mut input_tokens = 0u64; - let mut output_tokens = 0u64; + let mut usage = TokenUsage::default(); let mut resp_size = 0usize; let mut finish: Option = None; let mut debug_raw: Vec = Vec::new(); let mut interrupted: Option = None; + let mut idle = util::IdleTracker::new(std::time::Instant::now()); loop { let chunk = match byte_stream.next().await { Some(Ok(chunk)) => chunk, Some(Err(e)) => { interrupted = Some(e.to_string()); break; } None => break, }; + idle.mark_now(); if state.is_debug() { debug_raw.extend_from_slice(&chunk); } for line in lines.push(&chunk) { let Some(data) = util::sse_data(&line) else { continue }; @@ -1573,8 +1987,7 @@ async fn stream_gemini( if let Ok(v) = serde_json::from_str::(data) { if let Some(u) = v.get("usage") { if !u.is_null() { - input_tokens = u.get("prompt_tokens").and_then(|t| t.as_u64()).unwrap_or(input_tokens); - output_tokens = u.get("completion_tokens").and_then(|t| t.as_u64()).unwrap_or(output_tokens); + usage.merge_stream_update(util::openai_usage(u)); } } if let Some(choice) = v.get("choices").and_then(|c| c.as_array()).and_then(|a| a.first()) { @@ -1600,8 +2013,8 @@ async fn stream_gemini( tracing::warn!("[/v1beta/models] upstream stream interrupted: {reason}"); finish = Some("OTHER".to_string()); } - let usage = json!({"prompt_tokens": input_tokens, "completion_tokens": output_tokens}); - let ev = gemini::gemini_stream_final_chunk(finish.as_deref(), &usage, &model_json); + let usage_json = json!({"prompt_tokens": usage.input_tokens, "completion_tokens": usage.output_tokens}); + let ev = gemini::gemini_stream_final_chunk(finish.as_deref(), &usage_json, &model_json); let payload = format!("data: {}\n\n", serde_json::to_string(&ev).unwrap_or_default()); resp_size += payload.len(); yield Ok(Bytes::from(payload)); @@ -1616,8 +2029,14 @@ async fn stream_gemini( status_code: if interrupted.is_some() { 502 } else { status }, request_size: req_size, response_size: resp_size, - input_tokens, - output_tokens, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + output_tokens_final: None, + cache_read_input_tokens: usage.cache_read_input_tokens, + cache_creation_input_tokens: usage.cache_creation_input_tokens, + premium_multiplier: state.model_premium_multiplier(&translated).await, + upstream_idle_max_ms: Some(idle.max_idle_ms_including_now()), + keepalive_probes: None, duration: elapsed_secs(start), request_body: req_body, response_body: if state.is_debug() { Some(debug_resp) } else { None }, @@ -1627,11 +2046,13 @@ async fn stream_gemini( stop_reason: interrupted.is_some().then(|| "stream_interrupted".to_string()), tools_called: None, is_agent_initiated: None, - prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&translated, input_tokens, output_tokens)), + session_id: None, + prompt_cache_hit: cache_disposition(&usage), + failure_kind: None, + estimated_cost_usd: Some(calculate_cost(&translated, &usage)), }); }; - build_sse_response(stream) + build_sse_response(stream, COMMENT_KEEPALIVE_PROBE, probe_count.clone()) } // --------------------------------------------------------------------------- @@ -1677,11 +2098,7 @@ async fn embeddings(State(state): State, body: Bytes) -> Response { log_debug_response(&state, "/v1/embeddings", &text); if status.is_success() { let parsed: Value = serde_json::from_str(&text).unwrap_or(Value::Null); - let usage = parsed.get("usage").cloned().unwrap_or(json!({})); - let input_tokens = usage - .get("prompt_tokens") - .and_then(|t| t.as_u64()) - .unwrap_or(0); + let usage = util::openai_usage(&parsed.get("usage").cloned().unwrap_or(json!({}))); state.store.add(RequestRecord { id: uuid::Uuid::new_v4().to_string(), timestamp: now_iso(), @@ -1691,8 +2108,15 @@ async fn embeddings(State(state): State, body: Bytes) -> Response { status_code: status.as_u16(), request_size: req_size, response_size: resp_size, - input_tokens, + input_tokens: usage.input_tokens, output_tokens: 0, + output_tokens_final: None, + cache_read_input_tokens: usage.cache_read_input_tokens, + cache_creation_input_tokens: usage.cache_creation_input_tokens, + // Embeddings are not billed as premium requests. + premium_multiplier: None, + upstream_idle_max_ms: None, + keepalive_probes: None, duration: elapsed_secs(start), request_body: capture_json(&state, &req), response_body: capture_str(&state, &text), @@ -1702,8 +2126,10 @@ async fn embeddings(State(state): State, body: Bytes) -> Response { stop_reason: None, tools_called: None, is_agent_initiated: Some(false), - prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&model, input_tokens, 0)), + session_id: None, + prompt_cache_hit: cache_disposition(&usage), + failure_kind: None, + estimated_cost_usd: Some(calculate_cost(&model, &usage)), }); Json(parsed).into_response() } else { @@ -1821,7 +2247,7 @@ async fn stream_openai( // A non-2xx upstream (e.g. GitHub Models returning 401/403 as JSON when // the token lacks the `models: read` permission) is not an SSE stream — // surface it as a normal error instead of forwarding a broken "stream". - if !(200..300).contains(&status) { + if !is_streamable_status(status) { let text = upstream.text().await.unwrap_or_default(); log_debug_response(&state, endpoint, &text); // Newer OpenAI-family models reject `max_tokens` in favour of @@ -1845,16 +2271,19 @@ async fn stream_openai( }; let status = upstream.status().as_u16(); let model = translated.clone(); + // Shared with the keepalive wrapper so the record can report how + // many probes actually went out during an upstream silence. + let probe_count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)); let stream = async_stream::stream! { use futures_util::StreamExt; let mut byte_stream = upstream.bytes_stream(); let mut lines = util::SseLineBuffer::new(); - let mut input_tokens = 0u64; - let mut output_tokens = 0u64; + let mut usage = TokenUsage::default(); let mut resp_size = 0usize; let mut debug_raw: Vec = Vec::new(); let mut interrupted: Option = None; let mut saw_done = false; + let mut idle = util::IdleTracker::new(std::time::Instant::now()); loop { let chunk = match byte_stream.next().await { Some(Ok(chunk)) => chunk, @@ -1865,6 +2294,7 @@ async fn stream_openai( Some(Err(e)) => { interrupted = Some(e.to_string()); break; } None => break, }; + idle.mark_now(); if state.is_debug() { debug_raw.extend_from_slice(&chunk); } for line in lines.push(&chunk) { let Some(data) = util::sse_data(&line) else { continue }; @@ -1877,8 +2307,7 @@ async fn stream_openai( match serde_json::from_str::(data) { Ok(v) => { if let Some(u) = v.get("usage") { - input_tokens = u.get("prompt_tokens").and_then(|t| t.as_u64()).unwrap_or(input_tokens); - output_tokens = u.get("completion_tokens").and_then(|t| t.as_u64()).unwrap_or(output_tokens); + usage.merge_stream_update(util::openai_usage(u)); } resp_size += data.len(); yield Ok(Bytes::from(format!("data: {data}\n\n"))); @@ -1904,8 +2333,7 @@ async fn stream_openai( } else if !data.is_empty() { if let Ok(v) = serde_json::from_str::(data) { if let Some(u) = v.get("usage") { - input_tokens = u.get("prompt_tokens").and_then(|t| t.as_u64()).unwrap_or(input_tokens); - output_tokens = u.get("completion_tokens").and_then(|t| t.as_u64()).unwrap_or(output_tokens); + usage.merge_stream_update(util::openai_usage(u)); } } resp_size += data.len(); @@ -1939,8 +2367,14 @@ async fn stream_openai( status_code: if interrupted.is_some() { 502 } else { status }, request_size: req_size, response_size: resp_size, - input_tokens, - output_tokens, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + output_tokens_final: None, + cache_read_input_tokens: usage.cache_read_input_tokens, + cache_creation_input_tokens: usage.cache_creation_input_tokens, + premium_multiplier: state.model_premium_multiplier(&translated).await, + upstream_idle_max_ms: Some(idle.max_idle_ms_including_now()), + keepalive_probes: None, duration: elapsed_secs(start), request_body: req_body, response_body: if state.is_debug() { Some(debug_resp) } else { None }, @@ -1950,11 +2384,13 @@ async fn stream_openai( stop_reason: interrupted.is_some().then(|| "stream_interrupted".to_string()), tools_called: None, is_agent_initiated: None, - prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&translated, input_tokens, output_tokens)), + session_id: None, + prompt_cache_hit: cache_disposition(&usage), + failure_kind: None, + estimated_cost_usd: Some(calculate_cost(&translated, &usage)), }); }; - build_sse_response(stream) + build_sse_response(stream, COMMENT_KEEPALIVE_PROBE, probe_count.clone()) } /// Streams an OpenAI Responses SSE stream back to the client verbatim while @@ -1987,7 +2423,7 @@ async fn stream_responses( // A non-2xx upstream returns a JSON error body, not an SSE stream. Forward // it as a normal error response instead of wrapping it in a 200 "stream", // which would make clients treat an auth/quota failure as a valid answer. - if !(200..300).contains(&status) { + if !is_streamable_status(status) { let text = upstream.text().await.unwrap_or_default(); log_debug_response(&state, "/v1/responses", &text); log_error("/v1/responses", &req, &text, status); @@ -1996,22 +2432,26 @@ async fn stream_responses( text, ); } + // Shared with the keepalive wrapper so the record can report how + // many probes actually went out during an upstream silence. + let probe_count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)); let stream = async_stream::stream! { use futures_util::StreamExt; let mut byte_stream = upstream.bytes_stream(); let mut lines = util::SseLineBuffer::new(); - let mut input_tokens = 0u64; - let mut output_tokens = 0u64; + let mut usage = TokenUsage::default(); let mut resp_size = 0usize; let mut debug_raw: Vec = Vec::new(); let mut interrupted: Option = None; let mut completed = false; + let mut idle = util::IdleTracker::new(std::time::Instant::now()); loop { let chunk = match byte_stream.next().await { Some(Ok(chunk)) => chunk, Some(Err(e)) => { interrupted = Some(e.to_string()); break; } None => break, }; + idle.mark_now(); resp_size += chunk.len(); // Verbatim passthrough of raw bytes. Bytes are never re-decoded on // this path, so a multi-byte character split across chunks reaches @@ -2024,9 +2464,8 @@ async fn stream_responses( if let Ok(v) = serde_json::from_str::(data) { if v.get("type").and_then(|t| t.as_str()) == Some("response.completed") { completed = true; - let usage = v.get("response").and_then(|r| r.get("usage")).cloned().unwrap_or(json!({})); - input_tokens = usage.get("input_tokens").and_then(|t| t.as_u64()).unwrap_or(0); - output_tokens = usage.get("output_tokens").and_then(|t| t.as_u64()).unwrap_or(0); + let raw = v.get("response").and_then(|r| r.get("usage")).cloned().unwrap_or(json!({})); + usage = util::responses_usage(&raw); } } } @@ -2060,8 +2499,14 @@ async fn stream_responses( status_code: if interrupted.is_some() { 502 } else { status }, request_size: req_size, response_size: resp_size, - input_tokens, - output_tokens, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + output_tokens_final: None, + cache_read_input_tokens: usage.cache_read_input_tokens, + cache_creation_input_tokens: usage.cache_creation_input_tokens, + premium_multiplier: state.model_premium_multiplier(&translated).await, + upstream_idle_max_ms: Some(idle.max_idle_ms_including_now()), + keepalive_probes: None, duration: elapsed_secs(start), request_body: req_body, response_body: if state.is_debug() { Some(debug_resp) } else { None }, @@ -2072,80 +2517,123 @@ async fn stream_responses( .then(|| "stream_interrupted".to_string()), tools_called: None, is_agent_initiated: None, - prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&translated, input_tokens, output_tokens)), + session_id: None, + prompt_cache_hit: cache_disposition(&usage), + failure_kind: None, + estimated_cost_usd: Some(calculate_cost(&translated, &usage)), }); }; - build_sse_response(stream) + build_sse_response(stream, COMMENT_KEEPALIVE_PROBE, probe_count.clone()) } /// Streams a direct Anthropic SSE response back to the client verbatim. -async fn stream_anthropic_direct( - state: SharedState, - upstream: reqwest::Response, +/// Identity and bookkeeping for one request, threaded unchanged into whichever +/// record it ends up producing. +struct RequestMeta { original_model: String, translated: String, req_size: usize, req_body: Option, + session_id: Option, start: Instant, +} + +async fn stream_anthropic_direct( + state: SharedState, + upstream: reqwest::Response, + meta: RequestMeta, ) -> Response { + let RequestMeta { + original_model, + translated, + req_size, + req_body, + session_id, + start, + } = meta; let status = upstream.status().as_u16(); + // Without this the error body is wrapped in a 200 `text/event-stream`, so + // the client waits on a stream that never produces an event and reports a + // stalled response instead of the rate limit or auth failure that actually + // happened. The other three streaming paths already gate on this. + if !is_streamable_status(status) { + let text = upstream.text().await.unwrap_or_default(); + log_debug_response(&state, "/v1/messages", &text); + log_error( + "/v1/messages", + &json!({"model": &translated}), + &text, + status, + ); + record_failure( + &state, + "/v1/messages", + &translated, + Some(&translated), + status, + crate::store::failure::UPSTREAM_STATUS, + req_size, + req_body.clone(), + Some(text.clone()), + start, + session_id.clone(), + ); + return passthrough_error( + StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY), + text, + ); + } + // Resolved before the stream starts: the recorder's Drop path is + // synchronous and cannot await. + let premium_multiplier = state.model_premium_multiplier(&translated).await; + // Shared with the keepalive wrapper so the record can report how + // many probes actually went out during an upstream silence. + let probe_count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)); + let mut recorder = StreamRecorder::new( + state.clone(), + "/v1/messages", + original_model, + translated, + req_size, + req_body, + premium_multiplier, + start, + probe_count.clone(), + session_id.clone(), + ); let stream = async_stream::stream! { use futures_util::StreamExt; let mut byte_stream = upstream.bytes_stream(); - let mut input_tokens = 0u64; - let mut output_tokens = 0u64; - let mut resp_size = 0usize; let mut lines = util::SseLineBuffer::new(); - let mut debug_raw: Vec = Vec::new(); let mut interrupted: Option = None; - let mut saw_message_stop = false; - let mut stop_reason: Option = None; - let mut tools_called: Vec = Vec::new(); + let mut st = DirectStreamState::default(); loop { let chunk = match byte_stream.next().await { Some(Ok(chunk)) => chunk, Some(Err(e)) => { interrupted = Some(e.to_string()); break; } None => break, }; - resp_size += chunk.len(); + recorder.idle.mark_now(); + recorder.resp_size += chunk.len(); // Verbatim passthrough: the client receives the exact upstream // bytes, so characters split across chunks are never mangled. yield Ok::(Bytes::copy_from_slice(&chunk)); - if state.is_debug() { debug_raw.extend_from_slice(&chunk); } + if state.is_debug() { recorder.debug_raw.extend_from_slice(&chunk); } for line in lines.push(&chunk) { let Some(data) = util::sse_data(&line) else { continue }; if data.is_empty() { continue; } let Ok(v) = serde_json::from_str::(data) else { continue }; - match v.get("type").and_then(|t| t.as_str()) { - Some("message_start") => { - input_tokens = v.get("message").and_then(|m| m.get("usage")).and_then(|u| u.get("input_tokens")).and_then(|t| t.as_u64()).unwrap_or(0); - } - Some("message_delta") => { - output_tokens = v.get("usage").and_then(|u| u.get("output_tokens")).and_then(|t| t.as_u64()).unwrap_or(output_tokens); - if let Some(sr) = v.get("delta").and_then(|d| d.get("stop_reason")).and_then(|s| s.as_str()) { - stop_reason = Some(sr.to_string()); - } - } - Some("content_block_start") => { - if let Some(block) = v.get("content_block") { - if block.get("type").and_then(|t| t.as_str()) == Some("tool_use") { - if let Some(name) = block.get("name").and_then(|n| n.as_str()) { - if !tools_called.iter().any(|t| t == name) { - tools_called.push(name.to_string()); - } - } - } - } - } - Some("message_stop") => saw_message_stop = true, - _ => {} - } + st.observe(&v); + // Mirrored onto the recorder because a client disconnect drops + // this generator: `Drop` can only report what the recorder + // itself holds. + recorder.usage = st.usage; + recorder.usage_final = st.usage_final; } } // Anthropic clients wait for `message_stop`. Without it they either // hang or record the partial text as a completed assistant turn, which // then poisons every following request in the conversation. - if interrupted.is_some() || !saw_message_stop { + if interrupted.is_some() || !st.saw_message_stop { let reason = interrupted .clone() .unwrap_or_else(|| "upstream closed the stream before message_stop".to_string()); @@ -2158,35 +2646,22 @@ async fn stream_anthropic_direct( } }); yield Ok(Bytes::from(format!("event: error\ndata: {err}\n\n"))); - stop_reason = Some("stream_interrupted".to_string()); + st.stop_reason = Some("stream_interrupted".to_string()); } - let debug_resp = String::from_utf8_lossy(&debug_raw).into_owned(); + let debug_resp = String::from_utf8_lossy(&recorder.debug_raw).into_owned(); log_debug_response(&state, "/v1/messages", &debug_resp); - state.store.add(RequestRecord { - id: uuid::Uuid::new_v4().to_string(), - timestamp: now_iso(), - endpoint: "/v1/messages".to_string(), - model: original_model.clone(), - translated_model: (translated != original_model).then_some(translated.clone()), - status_code: if interrupted.is_some() { 502 } else { status }, - request_size: req_size, - response_size: resp_size, - input_tokens, - output_tokens, - duration: elapsed_secs(start), - request_body: req_body, - response_body: if state.is_debug() { Some(debug_resp) } else { None }, - message_count: None, - tool_count: None, - tool_names: None, - stop_reason, - tools_called: (!tools_called.is_empty()).then_some(tools_called), - is_agent_initiated: None, - prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&translated, input_tokens, output_tokens)), - }); + if let Some(r) = recorder.record_mut() { + r.tools_called = (!st.tools_called.is_empty()).then_some(st.tools_called); + } + recorder.finalize( + if interrupted.is_some() { 502 } else { status }, + (interrupted.is_some() || !st.saw_message_stop) + .then_some(crate::store::failure::STREAM_INTERRUPTED), + st.stop_reason, + state.is_debug().then_some(debug_resp), + ); }; - build_sse_response(stream) + build_sse_response(stream, ANTHROPIC_KEEPALIVE_PROBE, probe_count.clone()) } /// Streams an OpenAI chat-completions SSE stream translated into Anthropic @@ -2201,6 +2676,7 @@ async fn stream_anthropic_translated( translated: String, req_size: usize, start: Instant, + session_id: Option, ) -> Response { let req_body = state .is_debug() @@ -2218,7 +2694,7 @@ async fn stream_anthropic_translated( }; let status = upstream.status().as_u16(); // Surface a non-2xx upstream (JSON error, not SSE) as a normal error. - if !(200..300).contains(&status) { + if !is_streamable_status(status) { let text = upstream.text().await.unwrap_or_default(); log_debug_response(&state, "/v1/messages", &text); log_error( @@ -2227,28 +2703,45 @@ async fn stream_anthropic_translated( &text, status, ); + record_failure( + &state, + "/v1/messages", + &translated, + Some(&translated), + status, + crate::store::failure::UPSTREAM_STATUS, + req_size, + req_body.clone(), + Some(text.clone()), + start, + session_id.clone(), + ); return passthrough_error( StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY), text, ); } + // Shared with the keepalive wrapper so the record can report how + // many probes actually went out during an upstream silence. + let probe_count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)); let stream = async_stream::stream! { use futures_util::StreamExt; let mut byte_stream = upstream.bytes_stream(); let mut lines = util::SseLineBuffer::new(); let mut conv = AnthropicStreamState::new(); let mut chunks: Vec = Vec::new(); - let mut input_tokens = 0u64; - let mut output_tokens = 0u64; + let mut usage = TokenUsage::default(); let mut resp_size = 0usize; let mut debug_raw: Vec = Vec::new(); let mut interrupted: Option = None; + let mut idle = util::IdleTracker::new(std::time::Instant::now()); loop { let chunk = match byte_stream.next().await { Some(Ok(chunk)) => chunk, Some(Err(e)) => { interrupted = Some(e.to_string()); break; } None => break, }; + idle.mark_now(); if state.is_debug() { debug_raw.extend_from_slice(&chunk); } for line in lines.push(&chunk) { let Some(data) = util::sse_data(&line) else { continue }; @@ -2258,8 +2751,7 @@ async fn stream_anthropic_translated( continue; }; if let Some(u) = v.get("usage") { - input_tokens = u.get("prompt_tokens").and_then(|t| t.as_u64()).unwrap_or(input_tokens); - output_tokens = u.get("completion_tokens").and_then(|t| t.as_u64()).unwrap_or(output_tokens); + usage.merge_stream_update(util::openai_usage(u)); } chunks.push(v.clone()); for event in conv.process(&v) { @@ -2297,11 +2789,10 @@ async fn stream_anthropic_translated( yield Ok(Bytes::from(format!("event: error\ndata: {err}\n\n"))); } // Fall back to merged usage if streaming chunks did not carry usage. - if input_tokens == 0 && output_tokens == 0 { + if usage.input_tokens == 0 && usage.output_tokens == 0 { let merged = anthropic::merge_chat_chunks(&chunks); - if let Some(usage) = merged.get("usage") { - input_tokens = usage.get("prompt_tokens").and_then(|t| t.as_u64()).unwrap_or(0); - output_tokens = usage.get("completion_tokens").and_then(|t| t.as_u64()).unwrap_or(0); + if let Some(raw) = merged.get("usage") { + usage = util::openai_usage(raw); } } let debug_resp = String::from_utf8_lossy(&debug_raw).into_owned(); @@ -2315,8 +2806,14 @@ async fn stream_anthropic_translated( status_code: if interrupted.is_some() { 502 } else { status }, request_size: req_size, response_size: resp_size, - input_tokens, - output_tokens, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + output_tokens_final: None, + cache_read_input_tokens: usage.cache_read_input_tokens, + cache_creation_input_tokens: usage.cache_creation_input_tokens, + premium_multiplier: state.model_premium_multiplier(&translated).await, + upstream_idle_max_ms: Some(idle.max_idle_ms_including_now()), + keepalive_probes: None, duration: elapsed_secs(start), request_body: req_body, response_body: if state.is_debug() { Some(debug_resp) } else { None }, @@ -2326,11 +2823,13 @@ async fn stream_anthropic_translated( stop_reason: synthesized_stop.then(|| stop_reason.to_string()), tools_called: None, is_agent_initiated: None, - prompt_cache_hit: None, - estimated_cost_usd: Some(calculate_cost(&translated, input_tokens, output_tokens)), + session_id: None, + prompt_cache_hit: cache_disposition(&usage), + failure_kind: None, + estimated_cost_usd: Some(calculate_cost(&translated, &usage)), }); }; - build_sse_response(stream) + build_sse_response(stream, ANTHROPIC_KEEPALIVE_PROBE, probe_count.clone()) } /// How long a stream may stay silent before a keepalive comment is emitted. @@ -2349,8 +2848,64 @@ const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15); /// before anything has been sent — because the verbatim passthrough paths yield /// raw upstream chunks and splicing a comment into a half-written event would /// corrupt it. +/// Keepalive probe for the Anthropic Messages protocol. +/// +/// Anthropic's own streaming API emits `ping` events during long silences, and +/// clients reset their idle watchdog on *events*. An SSE comment is discarded +/// by the parser per spec, so it keeps the TCP connection and intermediaries +/// alive while never reaching the application — which is how a connection that +/// is provably still flowing bytes gets aborted as `Response stalled +/// mid-stream`. Measured against the real upstream: GitHub Copilot's Anthropic +/// endpoint never emits a ping of its own, so this probe is the only keepalive +/// signal on the link. +const ANTHROPIC_KEEPALIVE_PROBE: &[u8] = b"event: ping\ndata: {\"type\":\"ping\"}\n\n"; + +/// Keepalive probe for protocols that have no ping event of their own (OpenAI +/// chat completions, OpenAI Responses, Gemini). A comment is the only thing +/// guaranteed not to be mistaken for content by those parsers. +const COMMENT_KEEPALIVE_PROBE: &[u8] = b": keepalive\n\n"; + fn keepalive( stream: S, + probe: &'static [u8], + counter: std::sync::Arc, +) -> impl futures_util::Stream> +where + S: futures_util::Stream> + Send + 'static, +{ + keepalive_with_interval(stream, SSE_KEEPALIVE_INTERVAL, probe, counter) +} + +/// Byte offset just past the last **complete** SSE event in `buf`, or `None` +/// when `buf` does not yet contain a whole event. +fn last_event_boundary(buf: &[u8]) -> Option { + let lf = buf.windows(2).rposition(|w| w == b"\n\n").map(|i| i + 2); + let crlf = buf + .windows(4) + .rposition(|w| w == b"\r\n\r\n") + .map(|i| i + 4); + lf.max(crlf) +} + +/// Keepalive with an injectable interval so the stall behavior is testable +/// without waiting out the production interval. +/// +/// Upstream chunks are re-aligned to event boundaries: bytes that do not yet +/// complete an event are held back rather than forwarded. That is what makes +/// the keepalive unconditional — everything already sent downstream ends on a +/// blank line, so a comment can never land inside a half-written event. +/// +/// The previous implementation forwarded chunks verbatim and only emitted a +/// keepalive when the *last chunk* happened to end on a boundary. A TCP split +/// mid-event left that flag stuck `false` — and since only a new chunk could +/// clear it, an upstream that went quiet right then silenced the keepalive for +/// as long as it stayed quiet. The client saw zero bytes and aborted the +/// stream, which surfaced as `Response stalled mid-stream`. +fn keepalive_with_interval( + stream: S, + interval: Duration, + probe: &'static [u8], + counter: std::sync::Arc, ) -> impl futures_util::Stream> where S: futures_util::Stream> + Send + 'static, @@ -2358,42 +2913,50 @@ where async_stream::stream! { use futures_util::StreamExt; let mut inner = Box::pin(stream); - // True while the last bytes written ended an SSE event. - let mut at_boundary = true; + // Upstream bytes not yet forming a complete event. + let mut partial: Vec = Vec::new(); loop { tokio::select! { biased; item = inner.next() => { match item { - Some(chunk) => { - if let Ok(bytes) = &chunk { - at_boundary = ends_at_event_boundary(bytes); + Some(Ok(bytes)) => { + partial.extend_from_slice(&bytes); + if let Some(cut) = last_event_boundary(&partial) { + let whole: Vec = partial.drain(..cut).collect(); + yield Ok(Bytes::from(whole)); } - yield chunk; } + Some(Err(e)) => yield Err(e), None => break, } } - _ = tokio::time::sleep(SSE_KEEPALIVE_INTERVAL) => { - if at_boundary { - yield Ok(Bytes::from_static(b": keepalive\n\n")); - } + _ = tokio::time::sleep(interval) => { + counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + yield Ok(Bytes::from_static(probe)); } } } + // A final event that never received its blank line is still real data. + if !partial.is_empty() { + yield Ok(Bytes::from(partial)); + } } } -/// Whether a chunk ends on an SSE event boundary (a blank line). -fn ends_at_event_boundary(bytes: &[u8]) -> bool { - bytes.ends_with(b"\n\n") || bytes.ends_with(b"\r\n\r\n") -} - -fn build_sse_response(stream: S) -> Response +/// Wraps an SSE stream in a response, injecting `probe` during upstream +/// silences. The probe is protocol-specific — see the two `*_KEEPALIVE_PROBE` +/// constants; sending the wrong one either breaks the client's parser or fails +/// to reset its idle watchdog. +fn build_sse_response( + stream: S, + probe: &'static [u8], + counter: std::sync::Arc, +) -> Response where S: futures_util::Stream> + Send + 'static, { - let body = Body::from_stream(keepalive(stream)); + let body = Body::from_stream(keepalive(stream, probe, counter)); let mut resp = Response::new(body); *resp.headers_mut() = sse_headers(); resp @@ -2543,6 +3106,9 @@ async fn metrics_openmetrics(State(state): State) -> Response { let mut req_total_by_labels: HashMap<(String, u16, String), u64> = HashMap::new(); let mut in_tokens_by_model: HashMap = HashMap::new(); let mut out_tokens_by_model: HashMap = HashMap::new(); + let mut cache_read_by_model: HashMap = HashMap::new(); + let mut cache_creation_by_model: HashMap = HashMap::new(); + let mut premium_by_model: HashMap = HashMap::new(); let mut duration_sum_by_endpoint: HashMap = HashMap::new(); let mut duration_count_by_endpoint: HashMap = HashMap::new(); let mut cost_total = 0.0_f64; @@ -2560,6 +3126,12 @@ async fn metrics_openmetrics(State(state): State) -> Response { .entry((rec.endpoint.clone(), rec.status_code, model.clone())) .or_insert(0) += 1; *in_tokens_by_model.entry(model.clone()).or_insert(0) += rec.input_tokens; + *cache_read_by_model.entry(model.clone()).or_insert(0) += rec.cache_read_input_tokens; + *cache_creation_by_model.entry(model.clone()).or_insert(0) += + rec.cache_creation_input_tokens; + if let Some(multiplier) = rec.premium_multiplier { + *premium_by_model.entry(model.clone()).or_insert(0.0) += multiplier; + } *out_tokens_by_model.entry(model).or_insert(0) += rec.output_tokens; *duration_sum_by_endpoint .entry(rec.endpoint.clone()) @@ -2610,6 +3182,42 @@ async fn metrics_openmetrics(State(state): State) -> Response { )); } + out.push_str( + "# HELP ghc_proxy_cache_read_tokens_total Input tokens served from the prompt cache.\n", + ); + out.push_str("# TYPE ghc_proxy_cache_read_tokens_total counter\n"); + for (model, total) in cache_read_by_model { + out.push_str(&format!( + "ghc_proxy_cache_read_tokens_total{{model=\"{}\"}} {}\n", + metrics_label_escape(&model), + total + )); + } + + out.push_str( + "# HELP ghc_proxy_cache_creation_tokens_total Input tokens written into the prompt cache.\n", + ); + out.push_str("# TYPE ghc_proxy_cache_creation_tokens_total counter\n"); + for (model, total) in cache_creation_by_model { + out.push_str(&format!( + "ghc_proxy_cache_creation_tokens_total{{model=\"{}\"}} {}\n", + metrics_label_escape(&model), + total + )); + } + + out.push_str( + "# HELP ghc_proxy_premium_requests_total Copilot premium requests consumed, per the catalog's billing multiplier.\n", + ); + out.push_str("# TYPE ghc_proxy_premium_requests_total counter\n"); + for (model, total) in premium_by_model { + out.push_str(&format!( + "ghc_proxy_premium_requests_total{{model=\"{}\"}} {:.4}\n", + metrics_label_escape(&model), + total + )); + } + out.push_str( "# HELP ghc_proxy_request_duration_seconds_sum Sum of request durations by endpoint.\n", ); @@ -2726,7 +3334,28 @@ async fn api_requests( Query(params): Query>, ) -> Response { let (page, per_page, offset) = parse_pagination(¶ms); - let (items, total) = state.store.recent(per_page, offset); + // `failed=true` narrows to requests that never produced a usable answer — + // the ones worth looking at when something went wrong. + let failed_only = params.get("failed").map(|v| v == "true").unwrap_or(false); + // Several clients share one proxy, so isolating a single session is the + // difference between reading a log and guessing at one. + let session = params.get("session").map(|s| s.as_str()); + let (items, total) = if failed_only || session.is_some() { + state.store.filtered_page(per_page, offset, |r| { + if failed_only && r.status_code < 400 && r.failure_kind.is_none() { + return false; + } + match session { + Some(want) => r + .session_id + .as_deref() + .is_some_and(|id| id.starts_with(want)), + None => true, + } + }) + } else { + state.store.recent(per_page, offset) + }; let total_pages = total.div_ceil(per_page); Json(json!({ "items": items, @@ -2804,6 +3433,10 @@ async fn api_audit_summary(State(state): State) -> Response { let mut cache_hit_count = 0usize; let mut cache_write_count = 0usize; let mut record_count = 0usize; + let mut total_input_tokens = 0u64; + let mut total_cache_read_tokens = 0u64; + let mut total_cache_creation_tokens = 0u64; + let mut premium_requests = 0.0_f64; state.store.with_records(|records| { for rec in records { @@ -2836,6 +3469,16 @@ async fn api_audit_summary(State(state): State) -> Response { } else if rec.prompt_cache_hit == Some(false) { cache_write_count += 1; } + + // Token-level cache accounting. The per-request counters above + // only say how many turns touched the cache; these say how much of + // the prompt volume it actually absorbed. + total_input_tokens += rec.input_tokens; + total_cache_read_tokens += rec.cache_read_input_tokens; + total_cache_creation_tokens += rec.cache_creation_input_tokens; + if let Some(multiplier) = rec.premium_multiplier { + premium_requests += multiplier; + } } }); @@ -2857,17 +3500,80 @@ async fn api_audit_summary(State(state): State) -> Response { "stop_reasons": stop_reasons_sorted, "cache_hits": cache_hit_count, "cache_writes": cache_write_count, - "cache_hit_rate": if cache_hit_count + cache_write_count > 0 { - (cache_hit_count as f64 / (cache_hit_count + cache_write_count) as f64 * 100.0).round() / 100.0 - } else { - 0.0 - }, + // Share of requests that read from cache — a per-request count, not a + // token volume. Both rates below are 0–1 fractions. + "cache_hit_rate": ratio_2dp( + cache_hit_count as f64, + (cache_hit_count + cache_write_count) as f64, + ), + "total_input_tokens": total_input_tokens, + "total_cache_read_tokens": total_cache_read_tokens, + "total_cache_creation_tokens": total_cache_creation_tokens, + // Share of total prompt volume served from cache, which is what + // actually drives cost and latency — unlike the per-request rate above. + "token_cache_hit_rate": ratio_2dp( + total_cache_read_tokens as f64, + total_input_tokens as f64, + ), + "premium_requests": (premium_requests * 100.0).round() / 100.0, })) .into_response() } #[cfg(test)] mod tests { + #[test] + fn output_tokens_stay_provisional_until_message_delta() { + let mut st = super::DirectStreamState::default(); + // `message_start` opens with a placeholder output count. A stream cut + // short here (client stall abort) would otherwise publish "3" as the + // output of a turn that went on to emit 35899 — the exact numbers from + // the stalled Write call this guard exists for. + st.observe(&serde_json::json!({ + "type": "message_start", + "message": {"usage": { + "input_tokens": 2, + "cache_read_input_tokens": 342437, + "cache_creation_input_tokens": 6044, + "output_tokens": 3 + }} + })); + assert!(!st.usage_final); + assert_eq!(st.usage.output_tokens, 3); + assert_eq!(st.usage.input_tokens, 2 + 342437 + 6044); + + st.observe(&serde_json::json!({ + "type": "message_delta", + "usage": {"output_tokens": 35899}, + "delta": {"stop_reason": "tool_use"} + })); + assert!(st.usage_final); + assert_eq!(st.usage.output_tokens, 35899); + // `message_delta` restates output only; the input buckets captured at + // `message_start` must survive it. + assert_eq!(st.usage.input_tokens, 2 + 342437 + 6044); + assert_eq!(st.stop_reason.as_deref(), Some("tool_use")); + } + + #[test] + fn tool_names_are_collected_once_each() { + let mut st = super::DirectStreamState::default(); + for _ in 0..2 { + st.observe(&serde_json::json!({ + "type": "content_block_start", + "content_block": {"type": "tool_use", "name": "Write"} + })); + } + st.observe(&serde_json::json!({ + "type": "content_block_start", + "content_block": {"type": "text"} + })); + assert_eq!(st.tools_called, vec!["Write"]); + assert!(!st.saw_message_stop); + st.observe(&serde_json::json!({"type": "message_stop"})); + assert!(st.saw_message_stop); + } + use super::*; use axum::http::HeaderMap; @@ -2945,11 +3651,355 @@ mod tests { assert_eq!(model_rates("something-new"), (0.0005, 0.0015)); } + /// Minimal state for exercising the recording helpers. + fn test_state() -> SharedState { + std::sync::Arc::new(crate::state::AppState::new( + crate::config::Config::default(), + "test-token".into(), + )) + } + + #[test] + fn record_failure_captures_a_request_that_never_produced_a_response() { + let state = test_state(); + record_failure( + &state, + "/v1/messages", + "claude-opus-5", + Some("claude-opus-5"), + 429, + crate::store::failure::UPSTREAM_STATUS, + 1234, + None, + Some(r#"{"error":{"message":"rate limit"}}"#.into()), + Instant::now(), + Some("test-session".into()), + ); + let (items, total) = state.store.recent(10, 0); + assert_eq!(total, 1, "a failed request must still be recorded"); + let rec = &items[0]; + assert_eq!(rec.status_code, 429); + assert_eq!(rec.failure_kind.as_deref(), Some("upstream_status")); + assert_eq!(rec.request_size, 1234); + // The error body is the whole point of the record, so it is kept even + // with debug off. + assert!(rec.response_body.as_deref().unwrap().contains("rate limit")); + } + + #[test] + fn stream_recorder_records_a_client_disconnect_when_dropped_early() { + let state = test_state(); + { + let mut rec = StreamRecorder::new( + state.clone(), + "/v1/messages", + "claude-opus-5".into(), + "claude-opus-5".into(), + 500, + None, + None, + Instant::now(), + std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)), + Some("test-session".into()), + ); + rec.resp_size = 4096; + rec.usage.output_tokens = 77; + rec.debug_raw + .extend_from_slice(b"event: message_start\ndata: {}\n\n"); + // Deliberately no finalize() — this models axum dropping the + // response body when the client hangs up mid-stream. + } + let (items, total) = state.store.recent(10, 0); + assert_eq!(total, 1, "a client disconnect must leave a record behind"); + let rec = &items[0]; + assert_eq!(rec.status_code, 499); + assert_eq!(rec.failure_kind.as_deref(), Some("client_disconnected")); + // Whatever had streamed so far is preserved, not zeroed. + assert_eq!(rec.response_size, 4096); + assert_eq!(rec.output_tokens, 77); + // The partial upstream body is the main evidence for why the client + // gave up, so it must survive the disconnect too. + assert!( + rec.response_body + .as_deref() + .unwrap_or_default() + .contains("message_start"), + "partial response body lost on client disconnect" + ); + } + + #[test] + fn stream_recorder_does_not_double_record_after_finalize() { + let state = test_state(); + { + let mut rec = StreamRecorder::new( + state.clone(), + "/v1/messages", + "claude-opus-5".into(), + "claude-opus-5".into(), + 500, + None, + None, + Instant::now(), + std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)), + Some("test-session".into()), + ); + rec.finalize(200, None, None, None); + } // Drop runs here and must be a no-op. + let (_, total) = state.store.recent(10, 0); + assert_eq!( + total, 1, + "finalize followed by drop must record exactly once" + ); + } + + #[tokio::test] + async fn keepalive_probes_are_counted_for_diagnosis() { + use futures_util::StreamExt; + // A stalled stream is ambiguous after the fact: did the proxy stop + // signalling, or did the client ignore the signal? The counter is what + // separates the two, so it must track what actually went out. + let counter = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)); + let upstream = futures_util::stream::pending::>(); + let mut out = Box::pin(keepalive_with_interval( + upstream, + Duration::from_millis(25), + ANTHROPIC_KEEPALIVE_PROBE, + counter.clone(), + )); + for _ in 0..3 { + let _ = tokio::time::timeout(Duration::from_millis(300), out.next()).await; + } + assert!( + counter.load(std::sync::atomic::Ordering::Relaxed) >= 3, + "expected at least 3 probes, counted {}", + counter.load(std::sync::atomic::Ordering::Relaxed) + ); + } + + #[test] + fn session_id_is_read_from_the_nested_metadata_string() { + // `metadata.user_id` is a JSON *string* holding another JSON object, + // so it takes two parses. + let req = json!({"metadata": {"user_id": + "{\"device_id\":\"d3427\",\"account_uuid\":\"\",\"session_id\":\"7dea551a-c9f5-4ba1\"}"}}); + assert_eq!( + extract_session_id(&req).as_deref(), + Some("7dea551a-c9f5-4ba1") + ); + } + + #[test] + fn session_id_degrades_to_none_on_any_unexpected_shape() { + // This is diagnostic metadata — a malformed value must never take the + // request down with it. + for req in [ + json!({}), + json!({"metadata": {}}), + json!({"metadata": {"user_id": "not json at all"}}), + json!({"metadata": {"user_id": 12345}}), + json!({"metadata": {"user_id": "{\"device_id\":\"x\"}"}}), + json!({"metadata": {"user_id": "{\"session_id\":null}"}}), + json!({"metadata": {"user_id": "{\"session_id\":\"\"}"}}), + ] { + assert_eq!(extract_session_id(&req), None, "req was {req}"); + } + } + + #[test] + fn keepalive_probes_are_wellformed_for_their_protocol() { + // Both must terminate an SSE block, or they fuse with the next event. + assert!(ANTHROPIC_KEEPALIVE_PROBE.ends_with(b"\n\n")); + assert!(COMMENT_KEEPALIVE_PROBE.ends_with(b"\n\n")); + // Anthropic clients reset their idle watchdog on *events*. A comment + // is discarded by the SSE parser and never reaches them, which is how + // a live connection still gets aborted as "stalled". + assert!(ANTHROPIC_KEEPALIVE_PROBE.starts_with(b"event: ping")); + assert!(!ANTHROPIC_KEEPALIVE_PROBE.starts_with(b":")); + // OpenAI and Gemini have no ping event; sending one would surface as + // an unknown event to their parsers, so those paths keep the comment. + assert!(COMMENT_KEEPALIVE_PROBE.starts_with(b":")); + // The ping payload must be valid JSON — clients parse `data:`. + let payload = std::str::from_utf8(ANTHROPIC_KEEPALIVE_PROBE) + .unwrap() + .lines() + .find_map(|l| l.strip_prefix("data: ")) + .expect("ping carries a data line"); + let v: Value = serde_json::from_str(payload).expect("ping data is valid JSON"); + assert_eq!(v["type"], "ping"); + } + + #[tokio::test] + async fn anthropic_stall_emits_a_ping_event_rather_than_a_comment() { + use futures_util::StreamExt; + let upstream = futures_util::stream::pending::>(); + let mut out = Box::pin(keepalive_with_interval( + upstream, + Duration::from_millis(30), + ANTHROPIC_KEEPALIVE_PROBE, + std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)), + )); + let probe = tokio::time::timeout(Duration::from_millis(500), out.next()) + .await + .expect("a probe must be emitted while the upstream is silent") + .unwrap() + .unwrap(); + assert_eq!(&probe[..], ANTHROPIC_KEEPALIVE_PROBE); + } + + #[tokio::test] + async fn non_anthropic_stall_keeps_using_a_comment() { + use futures_util::StreamExt; + let upstream = futures_util::stream::pending::>(); + let mut out = Box::pin(keepalive_with_interval( + upstream, + Duration::from_millis(30), + COMMENT_KEEPALIVE_PROBE, + std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)), + )); + let probe = tokio::time::timeout(Duration::from_millis(500), out.next()) + .await + .expect("a probe must be emitted while the upstream is silent") + .unwrap() + .unwrap(); + assert_eq!(&probe[..], b": keepalive\n\n"); + } + + #[test] + fn last_event_boundary_finds_the_end_of_the_final_complete_event() { + assert_eq!(last_event_boundary(b"data: a\n\n"), Some(9)); + assert_eq!(last_event_boundary(b"data: a\r\n\r\n"), Some(11)); + // Two events: the cut is after the second one. + assert_eq!(last_event_boundary(b"data: a\n\ndata: b\n\n"), Some(18)); + // A trailing partial event is excluded from the cut. + assert_eq!(last_event_boundary(b"data: a\n\nevent: par"), Some(9)); + // Nothing complete yet. + assert_eq!(last_event_boundary(b"event: par"), None); + assert_eq!(last_event_boundary(b""), None); + } + + #[tokio::test] + async fn keepalive_still_fires_when_upstream_stalls_mid_event() { + use futures_util::StreamExt; + // One complete event followed by the first bytes of a second, then + // silence — exactly what a TCP split mid-event looks like when the + // model then goes quiet to think. + let upstream = futures_util::stream::iter(vec![Ok::<_, std::convert::Infallible>( + Bytes::from("event: content_block_delta\ndata: {}\n\nevent: content_bl"), + )]) + .chain(futures_util::stream::pending()); + + let mut out = Box::pin(keepalive_with_interval( + upstream, + Duration::from_millis(30), + ANTHROPIC_KEEPALIVE_PROBE, + std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)), + )); + + // Only the complete event is forwarded; the partial one is held back + // so the downstream byte stream stays on an event boundary. + let first = out.next().await.unwrap().unwrap(); + assert_eq!(&first[..], b"event: content_block_delta\ndata: {}\n\n"); + + // The upstream is now parked mid-event. A keepalive MUST still be + // emitted — otherwise the connection goes fully silent and the client + // aborts with "Response stalled mid-stream". + let next = tokio::time::timeout(Duration::from_millis(500), out.next()).await; + assert!( + next.is_ok(), + "no keepalive emitted while upstream was stalled mid-event" + ); + assert_eq!( + &next.unwrap().unwrap().unwrap()[..], + ANTHROPIC_KEEPALIVE_PROBE + ); + } + + #[test] + fn only_2xx_upstreams_are_streamed_as_sse() { + assert!(is_streamable_status(200)); + assert!(is_streamable_status(299)); + // These carry a JSON error body, not SSE. Streaming one yields a 200 + // with no events, which surfaces to the user as a stalled stream + // rather than the actual auth/quota failure. + for status in [401, 403, 429, 500, 502, 503] { + assert!(!is_streamable_status(status), "status {status}"); + } + } + + #[test] + fn ratios_are_two_decimal_fractions_never_percentages() { + // Half the prompt volume served from cache. + assert_eq!(ratio_2dp(7203.0, 14434.0), 0.5); + // A realistic Claude Code session: 98% cached, expressed as 0.98 — + // the same scale as the per-request `cache_hit_rate` beside it, so + // the two are not silently on different units. + assert_eq!(ratio_2dp(101_940.0, 103_673.0), 0.98); + assert_eq!(ratio_2dp(1.0, 1.0), 1.0); + // A zero denominator must yield 0.0 rather than NaN, which would + // serialize as `null` and break consumers. + assert_eq!(ratio_2dp(0.0, 0.0), 0.0); + assert!(ratio_2dp(5.0, 0.0).is_finite()); + } + #[test] fn cost_scales_with_tokens() { - let cost = calculate_cost("claude-opus-4.8", 1000, 1000); + // With no caching involved the whole prompt bills at the base rate, + // exactly as before cache-aware pricing was introduced. + let usage = TokenUsage { + input_tokens: 1000, + output_tokens: 1000, + ..TokenUsage::default() + }; + let cost = calculate_cost("claude-opus-4.8", &usage); assert!((cost - (0.015 + 0.075)).abs() < 1e-9); - assert_eq!(calculate_cost("gpt-4o", 0, 0), 0.0); + assert_eq!(calculate_cost("gpt-4o", &TokenUsage::default()), 0.0); + } + + #[test] + fn cost_prices_cache_reads_and_writes_at_their_own_rates() { + // 3k prompt: 1k uncached, 1k read from cache, 1k written to cache. + let usage = TokenUsage { + input_tokens: 3000, + cache_read_input_tokens: 1000, + cache_creation_input_tokens: 1000, + output_tokens: 1000, + }; + let (base_in, base_out) = model_rates("claude-opus-4.8"); + let expected = base_in // uncached remainder + + base_in * CACHE_WRITE_RATE_MULTIPLIER // cache write premium + + base_in * CACHE_READ_RATE_MULTIPLIER // cache read discount + + base_out; + let cost = calculate_cost("claude-opus-4.8", &usage); + assert!((cost - expected).abs() < 1e-9, "cost was {cost}"); + } + + #[test] + fn cached_prompts_cost_far_less_than_billing_them_at_full_rate() { + // The real shape of a Claude Code turn: ~98% of the prompt is a cache + // read. Charging the full input rate for all of it overstates the + // cost several times over. + let usage = TokenUsage { + input_tokens: 103_673, + cache_read_input_tokens: 101_940, + cache_creation_input_tokens: 1_731, + output_tokens: 561, + }; + let cache_aware = calculate_cost("claude-opus-5", &usage); + let (base_in, base_out) = model_rates("claude-opus-5"); + let naive = + (usage.input_tokens as f64 * base_in + usage.output_tokens as f64 * base_out) / 1000.0; + assert!( + cache_aware < naive, + "cache-aware {cache_aware} should be below naive {naive}" + ); + // ...but still far above counting only the 2 uncached tokens, which is + // what the dashboard used to report. + let broken = (2.0 * base_in + usage.output_tokens as f64 * base_out) / 1000.0; + assert!( + cache_aware > broken * 2.0, + "cache-aware {cache_aware} should dwarf the uncached-only {broken}" + ); } #[test] @@ -3009,13 +4059,14 @@ mod tests { } #[test] - fn keepalive_only_injects_at_event_boundaries() { - assert!(ends_at_event_boundary(b"data: {}\n\n")); - assert!(ends_at_event_boundary(b"data: {}\r\n\r\n")); - // Mid-event: injecting here would corrupt the passthrough stream. - assert!(!ends_at_event_boundary(b"data: {\"partial\":")); - assert!(!ends_at_event_boundary(b"data: {}\n")); - assert!(!ends_at_event_boundary(b"")); + fn keepalive_holds_back_a_partial_event_until_it_completes() { + // The old boundary check treated these as "unsafe to inject after". + // Now they are simply held back until the rest of the event arrives, + // so the downstream stream never sits mid-event. + assert_eq!(last_event_boundary(b"data: {\"partial\":"), None); + assert_eq!(last_event_boundary(b"data: {}\n"), None); + // Once the blank line lands, the whole event is released. + assert_eq!(last_event_boundary(b"data: {}\n\n"), Some(10)); } #[test] diff --git a/src/state.rs b/src/state.rs index 2319329..c5b0588 100644 --- a/src/state.rs +++ b/src/state.rs @@ -477,6 +477,16 @@ impl AppState { self.model_supports_endpoint(model, "/v1/messages").await } + /// Copilot premium-request cost for a model, from the catalog's + /// `billing.multiplier`. `None` when the catalog does not price the model. + pub async fn model_premium_multiplier(&self, model: &str) -> Option { + self.models + .read() + .await + .as_ref() + .and_then(|catalog| premium_multiplier_from_catalog(catalog, model)) + } + /// Maximum output tokens advertised for a model /// (`capabilities.limits.max_output_tokens`). Used to fill in a missing /// `max_tokens`, which some Copilot models reject the request without. @@ -669,3 +679,70 @@ pub fn summarize_usage(raw: &serde_json::Value) -> serde_json::Value { "raw": raw, }) } + +/// Extracts a model's Copilot premium-request cost from a `/models` catalog +/// payload. +/// +/// Returns `None` when the model is absent from the catalog or carries no +/// `billing.multiplier`, so a model the catalog never priced is not silently +/// counted as a full premium request. +fn premium_multiplier_from_catalog(catalog: &serde_json::Value, model: &str) -> Option { + catalog + .get("data")? + .as_array()? + .iter() + .find(|m| m.get("id").and_then(|i| i.as_str()) == Some(model))? + .get("billing")? + .get("multiplier")? + .as_f64() +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + #[test] + fn premium_multiplier_reads_billing_from_the_catalog() { + // Shape taken from a real `GET /models` response. + let catalog = json!({"data": [ + {"id": "claude-opus-5", "billing": {"is_premium": true, "multiplier": 1.0}}, + {"id": "gpt-4.1", "billing": {"is_premium": false, "multiplier": 0.0}}, + {"id": "discounted", "billing": {"is_premium": true, "multiplier": 0.33}}, + {"id": "no-billing-block"} + ]}); + assert_eq!( + premium_multiplier_from_catalog(&catalog, "claude-opus-5"), + Some(1.0) + ); + assert_eq!( + premium_multiplier_from_catalog(&catalog, "gpt-4.1"), + Some(0.0) + ); + assert_eq!( + premium_multiplier_from_catalog(&catalog, "discounted"), + Some(0.33) + ); + // A model without billing metadata, or one that is not in the catalog + // at all, must stay unknown rather than defaulting to a full request. + assert_eq!( + premium_multiplier_from_catalog(&catalog, "no-billing-block"), + None + ); + assert_eq!( + premium_multiplier_from_catalog(&catalog, "nonexistent"), + None + ); + } + + #[test] + fn premium_multiplier_survives_a_malformed_catalog() { + assert_eq!( + premium_multiplier_from_catalog(&json!({}), "claude-opus-5"), + None + ); + assert_eq!( + premium_multiplier_from_catalog(&json!({"data": "not-an-array"}), "x"), + None + ); + } +} diff --git a/src/store.rs b/src/store.rs index 3b15934..6647b04 100644 --- a/src/store.rs +++ b/src/store.rs @@ -17,6 +17,34 @@ pub struct RequestRecord { pub response_size: usize, pub input_tokens: u64, pub output_tokens: u64, + /// Whether `output_tokens` is the upstream's final figure. + /// + /// On an Anthropic stream the authoritative count arrives only in + /// `message_delta`, at the very end. A record finalized before that — a + /// client stall abort, an interrupted stream — carries the opening + /// placeholder from `message_start` instead, which is a single digit + /// against an eventual five-figure total. `Some(false)` marks the number + /// as unknown rather than letting the dashboard present it as fact. + /// `None` on paths that do not track this, so their counts render as + /// before. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub output_tokens_final: Option, + /// Portion of `input_tokens` that was served from a warm prompt cache. + /// Always serialized: the dashboard derives the cache hit ratio from it. + pub cache_read_input_tokens: u64, + /// Portion of `input_tokens` this request wrote into the cache. + pub cache_creation_input_tokens: u64, + /// Copilot premium-request cost, from the model catalog's + /// `billing.multiplier`. `None` when the model is unknown or the endpoint + /// carries no premium cost — never guessed as 1.0. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub premium_multiplier: Option, + /// Longest silence between two upstream chunks, in milliseconds. Only set + /// on streaming responses. When a stream stalls this is what assigns + /// blame: a large value means the upstream went quiet, a small one means + /// it kept sending and the stall was downstream of here. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub upstream_idle_max_ms: Option, pub duration: f64, /// Captured request body. Only populated when debug mode is enabled. #[serde(default, skip_serializing_if = "Option::is_none")] @@ -47,17 +75,55 @@ pub struct RequestRecord { /// Whether prompt caching was used (hit or write) #[serde(default, skip_serializing_if = "Option::is_none")] pub prompt_cache_hit: Option, + /// Client session this request belongs to, from `metadata.user_id`. + /// Several Claude Code instances can share one proxy; without this their + /// records interleave in the dashboard with no way to tell them apart. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub session_id: Option, + /// Keepalive probes actually written to the client during this response. + /// Pairs with `upstream_idle_max_ms` to place blame on a stalled stream: + /// a long idle with probes sent means the proxy kept signalling and the + /// client ignored it; a long idle with zero probes means the keepalive + /// itself failed. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub keepalive_probes: Option, + /// Why the request failed, when it did. `None` on a successful request. + /// Answers "which step broke" without needing a packet capture. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub failure_kind: Option, /// Estimated cost in USD based on token counts and model rates #[serde(default, skip_serializing_if = "Option::is_none")] pub estimated_cost_usd: Option, } +/// Values for [`RequestRecord::failure_kind`], ordered by how far the request +/// got before dying. +pub mod failure { + /// Never left the proxy — token refresh or the rate gate rejected it. + pub const PRECONDITION: &str = "precondition_failed"; + /// Sent, but the upstream never produced a response. + pub const CONNECT: &str = "connect_error"; + /// The upstream answered with a non-2xx status. + pub const UPSTREAM_STATUS: &str = "upstream_status"; + /// The upstream stream died partway through. + pub const STREAM_INTERRUPTED: &str = "stream_interrupted"; + /// The client hung up before the response finished. + pub const CLIENT_DISCONNECTED: &str = "client_disconnected"; +} + /// Aggregate statistics returned by `/api/stats`. #[derive(Debug, Default, Clone, Serialize)] pub struct Stats { pub request_count: u64, pub total_input_tokens: u64, pub total_output_tokens: u64, + /// Input tokens served from cache, across every recorded request. + pub total_cache_read_tokens: u64, + /// Input tokens written into the cache, across every recorded request. + pub total_cache_creation_tokens: u64, + /// Copilot premium requests consumed. Fractional because some models + /// bill at a discounted multiplier. + pub premium_requests: f64, pub bytes_sent: u64, pub bytes_received: u64, } @@ -90,6 +156,13 @@ impl RequestStore { inner.stats.request_count += 1; inner.stats.total_input_tokens += record.input_tokens; inner.stats.total_output_tokens += record.output_tokens; + inner.stats.total_cache_read_tokens += record.cache_read_input_tokens; + inner.stats.total_cache_creation_tokens += record.cache_creation_input_tokens; + // Only count what the catalog actually priced; an unknown multiplier + // is left out rather than assumed to be a full premium request. + if let Some(multiplier) = record.premium_multiplier { + inner.stats.premium_requests += multiplier; + } inner.stats.bytes_received += record.request_size as u64; inner.stats.bytes_sent += record.response_size as u64; inner.records.push_front(record); @@ -168,6 +241,13 @@ mod tests { response_size: 2, input_tokens: 3, output_tokens: 4, + output_tokens_final: None, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + premium_multiplier: None, + upstream_idle_max_ms: None, + session_id: None, + keepalive_probes: None, duration: 0.5, request_body: None, response_body: None, @@ -179,6 +259,7 @@ mod tests { is_agent_initiated: None, estimated_cost_usd: None, prompt_cache_hit: None, + failure_kind: None, } } @@ -216,6 +297,43 @@ mod tests { assert_eq!(items[0].status_code, 200); } + #[test] + fn stats_track_cache_buckets_and_premium_requests() { + let store = RequestStore::new(10); + + // A fully-cached Claude Code turn: almost the entire prompt is a + // cache read, so the uncached remainder is a rounding error. + let mut cached = record("/v1/messages", 200); + cached.input_tokens = 103_673; + cached.cache_read_input_tokens = 101_940; + cached.cache_creation_input_tokens = 1_731; + cached.premium_multiplier = Some(1.0); + store.add(cached); + + // A discounted model still counts, at its own rate. + let mut cheap = record("/v1/messages", 200); + cheap.input_tokens = 500; + cheap.premium_multiplier = Some(0.33); + store.add(cheap); + + // Embeddings burn no premium allowance; an unknown multiplier must + // not be invented as 1.0. + let mut free = record("/v1/embeddings", 200); + free.input_tokens = 20; + free.premium_multiplier = None; + store.add(free); + + let s = store.stats(); + assert_eq!(s.total_input_tokens, 103_673 + 500 + 20); + assert_eq!(s.total_cache_read_tokens, 101_940); + assert_eq!(s.total_cache_creation_tokens, 1_731); + assert!( + (s.premium_requests - 1.33).abs() < 1e-9, + "premium_requests was {}", + s.premium_requests + ); + } + #[test] fn with_records_visits_every_record() { let store = RequestStore::new(10); diff --git a/src/util.rs b/src/util.rs index c13f027..dad7ae8 100644 --- a/src/util.rs +++ b/src/util.rs @@ -291,6 +291,150 @@ pub async fn post_with_retry( } } +/// Token counts for one request, split by how the input was served. +/// +/// `input_tokens` is always the **true total** prompt size. Providers disagree +/// on how they slice that total, so the extractors below normalize it: +/// Anthropic reports three disjoint buckets that must be summed, while +/// OpenAI's `prompt_tokens` already contains its cached subset. +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub struct TokenUsage { + /// Total prompt tokens, cached and uncached alike. + pub input_tokens: u64, + /// Portion of `input_tokens` served from a warm cache. + pub cache_read_input_tokens: u64, + /// Portion of `input_tokens` written into the cache by this request. + pub cache_creation_input_tokens: u64, + pub output_tokens: u64, +} + +impl TokenUsage { + /// Fraction of the prompt served from cache, or `None` when the request + /// had no input at all (which would make the ratio meaningless rather + /// than zero). + pub fn cache_hit_ratio(&self) -> Option { + (self.input_tokens > 0) + .then(|| self.cache_read_input_tokens as f64 / self.input_tokens as f64) + } + + /// Adopts a usage object parsed from a streaming chunk, keeping the + /// previous totals when the new one is empty. + /// + /// Streams repeat `usage` across chunks — frequently as `null` or an empty + /// object until the final event — so overwriting unconditionally would + /// zero out counts a later chunk simply did not restate. + pub fn merge_stream_update(&mut self, next: TokenUsage) { + if next.input_tokens > 0 || next.output_tokens > 0 { + *self = next; + } + } +} + +/// Reads a `u64` field, treating absent/malformed values as zero. +fn usage_field(usage: &Value, key: &str) -> u64 { + usage.get(key).and_then(|t| t.as_u64()).unwrap_or(0) +} + +/// Normalizes an Anthropic `usage` object. +/// +/// `input_tokens`, `cache_read_input_tokens` and `cache_creation_input_tokens` +/// are **disjoint**, so the real prompt size is their sum. Reading +/// `input_tokens` on its own reports single digits for a fully-cached +/// hundred-thousand-token conversation, because Claude Code covers almost the +/// entire prompt — including the newest message — with `cache_control`. +pub fn anthropic_usage(usage: &Value) -> TokenUsage { + let cache_read = usage_field(usage, "cache_read_input_tokens"); + let cache_creation = usage_field(usage, "cache_creation_input_tokens"); + TokenUsage { + input_tokens: usage_field(usage, "input_tokens") + cache_read + cache_creation, + cache_read_input_tokens: cache_read, + cache_creation_input_tokens: cache_creation, + output_tokens: usage_field(usage, "output_tokens"), + } +} + +/// Normalizes an OpenAI `usage` object. +/// +/// Unlike Anthropic's disjoint buckets, `prompt_tokens` is already the total +/// and `prompt_tokens_details.cached_tokens` is a subset of it — adding them +/// would double-count. OpenAI has no cache-write concept, so that bucket stays +/// zero. +pub fn openai_usage(usage: &Value) -> TokenUsage { + TokenUsage { + input_tokens: usage_field(usage, "prompt_tokens"), + cache_read_input_tokens: usage + .get("prompt_tokens_details") + .map(|d| usage_field(d, "cached_tokens")) + .unwrap_or(0), + cache_creation_input_tokens: 0, + output_tokens: usage_field(usage, "completion_tokens"), + } +} + +/// Normalizes an OpenAI **Responses API** `usage` object. +/// +/// A trap worth naming: this API spells its totals `input_tokens` / +/// `output_tokens` — the same keys Anthropic uses — but with OpenAI's +/// semantics, where `input_tokens` is already the grand total and +/// `input_tokens_details.cached_tokens` is a subset of it. Feeding one of +/// these payloads to [`anthropic_usage`] would double-count the cached half. +pub fn responses_usage(usage: &Value) -> TokenUsage { + TokenUsage { + input_tokens: usage_field(usage, "input_tokens"), + cache_read_input_tokens: usage + .get("input_tokens_details") + .map(|d| usage_field(d, "cached_tokens")) + .unwrap_or(0), + cache_creation_input_tokens: 0, + output_tokens: usage_field(usage, "output_tokens"), + } +} + +/// Tracks the longest silence between upstream chunks on a streaming response. +/// +/// When a stream stalls, this is the number that assigns blame: a large value +/// means the upstream itself went quiet, while a small one means the upstream +/// kept sending and the stall happened on the proxy or client side. Without it +/// the two are indistinguishable after the fact. +/// +/// The wait before the first chunk counts as idle time too — a stream that +/// never produces anything is the worst stall of all. +pub struct IdleTracker { + last: std::time::Instant, + max_idle_ms: u64, +} + +impl IdleTracker { + pub fn new(start: std::time::Instant) -> Self { + IdleTracker { + last: start, + max_idle_ms: 0, + } + } + + /// Records that a chunk arrived at `now`. + pub fn mark(&mut self, now: std::time::Instant) { + let gap = now.saturating_duration_since(self.last).as_millis() as u64; + self.max_idle_ms = self.max_idle_ms.max(gap); + self.last = now; + } + + /// Records a chunk arriving right now. + pub fn mark_now(&mut self) { + self.mark(std::time::Instant::now()); + } + + /// Longest observed gap, including any trailing silence up to `now`. + pub fn max_idle_ms_including_now(&self) -> u64 { + let trailing = self.last.elapsed().as_millis() as u64; + self.max_idle_ms.max(trailing) + } + + pub fn max_idle_ms(&self) -> u64 { + self.max_idle_ms + } +} + #[cfg(test)] mod tests { use super::*; @@ -397,4 +541,139 @@ mod tests { assert!(!is_max_tokens_unsupported_error(200, body)); assert!(!is_max_tokens_unsupported_error(400, "some other error")); } + + #[test] + fn anthropic_usage_sums_all_three_input_buckets() { + // Verbatim from a real Copilot `/v1/messages` message_start event. + // Reading only `input_tokens` reports 2 for a 100k-token prompt. + let usage = serde_json::json!({ + "input_tokens": 2, + "cache_read_input_tokens": 101940, + "cache_creation_input_tokens": 1731, + "cache_creation": { + "ephemeral_5m_input_tokens": 1731, + "ephemeral_1h_input_tokens": 0 + }, + "output_tokens": 2 + }); + let u = anthropic_usage(&usage); + assert_eq!(u.input_tokens, 103_673); + assert_eq!(u.cache_read_input_tokens, 101_940); + assert_eq!(u.cache_creation_input_tokens, 1_731); + assert_eq!(u.output_tokens, 2); + } + + #[test] + fn anthropic_usage_tolerates_absent_cache_fields() { + // A cache-less turn: the total is just the uncached input. + let usage = serde_json::json!({"input_tokens": 4_096, "output_tokens": 128}); + let u = anthropic_usage(&usage); + assert_eq!(u.input_tokens, 4_096); + assert_eq!(u.cache_read_input_tokens, 0); + assert_eq!(u.cache_creation_input_tokens, 0); + assert_eq!(u.output_tokens, 128); + } + + #[test] + fn openai_usage_does_not_double_count_cached_tokens() { + // OpenAI's `prompt_tokens` ALREADY includes `cached_tokens`, unlike + // Anthropic's disjoint buckets. Summing them would inflate the total. + let usage = serde_json::json!({ + "prompt_tokens": 5_000, + "completion_tokens": 300, + "prompt_tokens_details": {"cached_tokens": 4_500} + }); + let u = openai_usage(&usage); + assert_eq!(u.input_tokens, 5_000); + assert_eq!(u.cache_read_input_tokens, 4_500); + // OpenAI has no cache-write concept. + assert_eq!(u.cache_creation_input_tokens, 0); + assert_eq!(u.output_tokens, 300); + } + + #[test] + fn idle_tracker_reports_the_longest_upstream_gap() { + let t0 = std::time::Instant::now(); + let mut t = IdleTracker::new(t0); + // Chunks arriving steadily. + t.mark(t0 + Duration::from_millis(100)); + t.mark(t0 + Duration::from_millis(150)); + assert_eq!(t.max_idle_ms(), 100); + // A long think between chunks — this is the number that distinguishes + // "upstream went quiet" from "proxy failed to forward". + t.mark(t0 + Duration::from_millis(2_150)); + assert_eq!(t.max_idle_ms(), 2_000); + // A later short gap must not lower the recorded maximum. + t.mark(t0 + Duration::from_millis(2_200)); + assert_eq!(t.max_idle_ms(), 2_000); + } + + #[test] + fn idle_tracker_counts_the_wait_before_the_first_chunk() { + // Time-to-first-byte is itself an upstream stall; a stream that never + // produces a chunk at all must not report zero idle time. + let t0 = std::time::Instant::now(); + let mut t = IdleTracker::new(t0); + t.mark(t0 + Duration::from_millis(4_000)); + assert_eq!(t.max_idle_ms(), 4_000); + } + + #[test] + fn responses_usage_reads_the_responses_api_field_names() { + // The Responses API reuses Anthropic's key names for a DIFFERENT + // meaning: its `input_tokens` is already the grand total and + // `input_tokens_details.cached_tokens` is a subset, so these must not + // be summed the way anthropic_usage sums its disjoint buckets. + let usage = serde_json::json!({ + "input_tokens": 8_000, + "output_tokens": 250, + "input_tokens_details": {"cached_tokens": 6_000} + }); + let u = responses_usage(&usage); + assert_eq!(u.input_tokens, 8_000); + assert_eq!(u.cache_read_input_tokens, 6_000); + assert_eq!(u.cache_creation_input_tokens, 0); + assert_eq!(u.output_tokens, 250); + } + + #[test] + fn stream_usage_updates_ignore_empty_repeats() { + let mut running = TokenUsage::default(); + running.merge_stream_update(TokenUsage { + input_tokens: 5_000, + cache_read_input_tokens: 4_000, + output_tokens: 10, + ..TokenUsage::default() + }); + assert_eq!(running.input_tokens, 5_000); + + // Streams repeat `usage` across chunks, often empty until the last + // one. Adopting an empty repeat would wipe totals already collected. + running.merge_stream_update(TokenUsage::default()); + assert_eq!(running.input_tokens, 5_000); + assert_eq!(running.cache_read_input_tokens, 4_000); + assert_eq!(running.output_tokens, 10); + + // A later chunk carrying real numbers does win. + running.merge_stream_update(TokenUsage { + input_tokens: 5_000, + cache_read_input_tokens: 4_000, + output_tokens: 250, + ..TokenUsage::default() + }); + assert_eq!(running.output_tokens, 250); + } + + #[test] + fn cache_hit_ratio_is_none_without_input() { + let none = TokenUsage::default(); + assert_eq!(none.cache_hit_ratio(), None); + let hit = TokenUsage { + input_tokens: 103_673, + cache_read_input_tokens: 101_940, + ..TokenUsage::default() + }; + let ratio = hit.cache_hit_ratio().unwrap(); + assert!((ratio - 0.983).abs() < 0.001, "ratio was {ratio}"); + } }
TimeEndpointModelStatusInOutDurationTimeSessionEndpointModelStatusInNewOutIdleProbesDuration
${open ? '▾' : '▸'} ${escapeHtml(r.timestamp)}${r.session_id ? `${escapeHtml(r.session_id.slice(0,8))}` : '—'} ${escapeHtml(r.endpoint)} ${escapeHtml(model)}${r.status_code}${r.input_tokens}${r.output_tokens}${r.status_code}${r.failure_kind ? `
${escapeHtml(r.failure_kind)}
` : ''}
${formatInput(r)}${compactTokens(r.cache_creation_input_tokens ?? 0)}${formatOutput(r)}${formatIdle(r)}${r.keepalive_probes ?? '—'} ${r.duration}s
${detail}
${detail}