diff --git a/src/agents/plugins/claude/plugin/statusline.mjs b/src/agents/plugins/claude/plugin/statusline.mjs index 16dfd705e..645f71796 100644 --- a/src/agents/plugins/claude/plugin/statusline.mjs +++ b/src/agents/plugins/claude/plugin/statusline.mjs @@ -1,6 +1,9 @@ #!/usr/bin/env node // CodeMie statusline — shows model, project, branch, context, session cost/duration, // and (when a CodeMie profile is configured) the CLI budget for the authenticated user. +// When the request was routed to a different backend model (CodeMie Switchyard or the +// LiteLLM router), the actual model is read from the routing headers the proxy injects +// into the transcript and shown alongside the nominal one — see resolveActualModel(). // Deployed to ~/.claude/ by `codemie install statusline` (also triggered by the `--status` // CLI flag, which calls the same installer). Runs standalone — Node builtins only, no // project imports, since it executes via `node ` after the project process exits. @@ -91,6 +94,8 @@ export function extractBasicInfo(ctx) { return { projectName: cwd ? path.basename(cwd) : '', cwd, + transcriptPath: ctx?.transcript_path ?? '', + modelId: ctx?.model?.id ?? '', model: ctx?.model?.display_name ?? '', ctxPct: ctx?.context_window?.used_percentage ?? null, tokIn: ctx?.context_window?.total_input_tokens ?? null, @@ -100,6 +105,134 @@ export function extractBasicInfo(ctx) { }; } +// --- Actual (routed) model resolution --- +// +// Claude Code's own stdin JSON only ever reports the nominal model (`model.id`, the +// alias/tier the session was started with). When CodeMie Switchyard or the LiteLLM router +// dispatches a turn to a different backend model, that can surface two ways in the transcript's +// most recent assistant turn (transcript_path): +// 1. Routing headers — the proxy's routing-header-injector plugin copies the upstream +// router's response headers onto the response body (see +// src/providers/plugins/sso/proxy/plugins/routing-header-injector.plugin.ts), which +// Claude Code then persists verbatim. Authoritative when present: the proxy tags these +// explicitly, so they win over the body-model heuristic below. +// 2. The response body's own `model` field — every Anthropic-compatible response reports +// the model that actually generated it. A router that doesn't emit routing headers (or +// a deployment where this proxy isn't involved at all) still shows the truth here, so +// it's a fallback signal rather than depending on headers alone. +// +// Header field precedence mirrors src/cli/commands/analytics/cost/usage-readers.ts's +// `routedModel` so the statusline and the analytics report agree when both are present. + +const ROUTED_MODEL_TAIL_BYTES = 65_536; // last 64KB — comfortably covers the most recent turn(s) + +// Claude model family names, used to tell "genuinely routed to a different tier" apart from +// "alias resolved to its concrete dated/region-qualified snapshot", which happens on every +// request regardless of routing and must never be shown as if it were routing. +const MODEL_FAMILY_PATTERN = /(opus|sonnet|haiku|fable)/i; + +/** Pulls the actual dispatched model out of a transcript line's `message` object, if present. */ +export function extractRoutedModel(message) { + if (!message || typeof message !== 'object') return null; + return message.x_codemie_routing_capable_model ?? message['x-litellm-router-routed-model'] ?? null; +} + +function modelFamily(modelId) { + const match = modelId ? MODEL_FAMILY_PATTERN.exec(modelId) : null; + return match ? match[1].toLowerCase() : null; +} + +/** Strips Bedrock region/provider qualifiers (`converse/global.anthropic.` / `eu.anthropic.`) and its `-v1:0` suffix. */ +export function normalizeModelId(modelId) { + if (!modelId) return ''; + return modelId + .toLowerCase() + .replace(/^converse\//, '') + .replace(/^[a-z0-9-]+\.anthropic\./, '') + .replace(/-v\d+:\d+$/, ''); +} + +/** + * True when two model identifiers name the same tier — either a recognized family (opus/ + * sonnet/haiku/fable) matches on both sides, or, when neither side matches a known family, + * the Bedrock-normalized identifiers are identical. Used to avoid flagging an alias's normal + * resolution to its concrete provider snapshot as if it were routing. + */ +export function sameModelFamily(a, b) { + if (!a || !b) return false; + const famA = modelFamily(a); + const famB = modelFamily(b); + if (famA && famB) return famA === famB; + return normalizeModelId(a) === normalizeModelId(b); +} + +/** + * Scans transcript JSONL text backwards for the most recent assistant turn and returns the + * response body's own model plus any header-injected routed model. The first (partial) line + * of a tail read is expected to fail JSON.parse when the read didn't start at a line boundary + * — that's normal, not an error, so parse failures are skipped rather than treated as a reason + * to stop scanning. + */ +export function parseLastAssistantTurn(tailText) { + if (!tailText) return null; + const lines = tailText.split('\n'); + for (let i = lines.length - 1; i >= 0; i--) { + const line = lines[i].trim(); + if (!line) continue; + let parsed; + try { + parsed = JSON.parse(line); + } catch { + continue; + } + const message = parsed?.message; + if (parsed?.type === 'assistant' && message?.model) { + return { responseModel: message.model, headerRoutedModel: extractRoutedModel(message) }; + } + } + return null; +} + +async function defaultReadTail(filePath, maxBytes) { + const handle = await fs.open(filePath, 'r'); + try { + const { size } = await handle.stat(); + const start = Math.max(0, size - maxBytes); + const length = size - start; + if (length <= 0) return ''; + const { buffer, bytesRead } = await handle.read({ buffer: Buffer.alloc(length), position: start }); + return buffer.toString('utf8', 0, bytesRead); + } finally { + await handle.close(); + } +} + +/** + * Resolves the actual routed model for the current session, or null when there is nothing + * useful to show — no transcript, an unreadable transcript, or the best available signal + * (routing headers, falling back to the response body's own model) names the same tier as + * `nominalModelId`. Never throws: the statusline must keep rendering even if the transcript + * is mid-write or has already rotated away. + */ +export async function resolveActualModel(transcriptPath, nominalModelId, { readTail = defaultReadTail } = {}) { + if (!transcriptPath) return null; + let tail; + try { + tail = await readTail(transcriptPath, ROUTED_MODEL_TAIL_BYTES); + } catch { + return null; + } + const turn = parseLastAssistantTurn(tail); + if (!turn) return null; + const candidate = turn.headerRoutedModel ?? turn.responseModel; + if (!candidate) return null; + const nominal = nominalModelId || turn.responseModel; + // Display the Bedrock-stripped form — the raw candidate may be a fully qualified backend + // id (e.g. `converse/global.anthropic.claude-haiku-4-5-20251001-v1:0`), which is accurate + // but not what a human wants to read in a one-line statusline. + return sameModelFamily(nominal, candidate) ? null : normalizeModelId(candidate); +} + export function formatDuration(ms) { if (typeof ms !== 'number' || Number.isNaN(ms) || ms < 0) return null; const mins = Math.floor(ms / 60000); @@ -138,14 +271,14 @@ export function ctxBar(pct) { return `${c(color, bar)} ${pct}%`; } -export function buildStatusLine({ projectName, branch, model, ctxPct, tokIn, tokOut, cost, durationMs, budget, budgetError }) { +export function buildStatusLine({ projectName, branch, model, actualModel, ctxPct, tokIn, tokOut, cost, durationMs, budget, budgetError }) { const parts = []; if (projectName) parts.push(c(C.purple, `[${projectName}]`)); if (budget) parts.push(c(budgetColor(budget.pct), budget.text)); else if (budgetError) parts.push(c(C.yellow, `⚠ ${budgetError}`)); if (branch) parts.push(c(C.blue, `(${branch})`)); - if (model) parts.push(c(C.cyan, `[${model}]`)); + if (model) parts.push(c(C.cyan, `[${actualModel ? `${model} → ${actualModel}` : model}]`)); const bar = ctxBar(ctxPct); if (bar) parts.push(bar); @@ -256,9 +389,13 @@ export async function main() { } const branchPromise = basic.cwd ? gitBranch(basic.cwd) : Promise.resolve(''); - const [budgetResult, branch] = await Promise.all([resolveBudget(), branchPromise]); + const [budgetResult, branch, actualModel] = await Promise.all([ + resolveBudget(), + branchPromise, + resolveActualModel(basic.transcriptPath, basic.modelId), + ]); - process.stdout.write(buildStatusLine({ ...basic, branch, ...budgetResult })); + process.stdout.write(buildStatusLine({ ...basic, branch, actualModel, ...budgetResult })); } // Compares decoded paths (not raw strings) so this correctly matches even when the diff --git a/src/cli/commands/analytics/cost/cost-enricher.ts b/src/cli/commands/analytics/cost/cost-enricher.ts index ea8bbdd12..61e110539 100644 --- a/src/cli/commands/analytics/cost/cost-enricher.ts +++ b/src/cli/commands/analytics/cost/cost-enricher.ts @@ -198,7 +198,13 @@ export function buildModelTimeline(records: UsageRecord[]): ModelTimelinePoint[] }; if (r.requestedModel != null) point.requestedModel = r.requestedModel; if (r.capableModel != null) point.capableModel = r.capableModel; + if (r.routingFamily != null) point.routingFamily = r.routingFamily; if (r.routingTier != null) point.routingTier = r.routingTier; + if (r.routingTierRaw != null) point.routingTierRaw = r.routingTierRaw; + if (r.routedModel != null) point.routedModel = r.routedModel; + if (r.classifierModel != null) point.classifierModel = r.classifierModel; + if (r.routerType != null) point.routerType = r.routerType; + if (r.routerScore != null) point.routerScore = r.routerScore; if (r.routingConfidence != null) point.routingConfidence = r.routingConfidence; if (r.routingSource != null) point.routingSource = r.routingSource; if (r.signalScore != null) point.signalScore = r.signalScore; @@ -402,18 +408,29 @@ export async function enrichCosts( if (records.length) { const timeline = buildModelTimeline(records); if (timeline.length) cost.modelTimeline = timeline; - // Accumulate classifier (Switchyard routing LLM) cost and tokens from per-turn metadata. + // Accumulate classifier (routing LLM) cost and tokens from per-turn metadata. let judgeCostUSD = 0; let judgeInputTokens = 0; let judgeOutputTokens = 0; let judgeCachedTokens = 0; let judgeCacheCreationTokens = 0; + // A session is only "cost known" if every routed turn came from a family that reports + // cost. One LiteLLM turn makes the session total an understatement, not a measurement. + let routedTurns = 0; + let costKnownTurns = 0; for (const r of records) { judgeCostUSD += r.judgeCostUSD ?? 0; judgeInputTokens += r.judgeInputTokens ?? 0; judgeOutputTokens += r.judgeOutputTokens ?? 0; judgeCachedTokens += r.judgeCachedTokens ?? 0; judgeCacheCreationTokens += r.judgeCacheCreationTokens ?? 0; + if (r.routingFamily != null) { + routedTurns++; + if (r.routingCostKnown) costKnownTurns++; + } + } + if (routedTurns > 0) { + cost.routingCostKnown = costKnownTurns === routedTurns; } if (judgeInputTokens > 0 || judgeOutputTokens > 0 || judgeCostUSD > 0) { cost.judgeCostUSD = Math.round(judgeCostUSD * 1e8) / 1e8; diff --git a/src/cli/commands/analytics/cost/types.ts b/src/cli/commands/analytics/cost/types.ts index 309fbc795..9f09c9cec 100644 --- a/src/cli/commands/analytics/cost/types.ts +++ b/src/cli/commands/analytics/cost/types.ts @@ -38,7 +38,13 @@ export interface ModelTimelinePoint { tokens: number; // per-turn total tokens for this turn requestedModel?: string; // capable model that was originally requested capableModel?: string; // alias for requestedModel + routingFamily?: 'switchyard' | 'litellm'; // which header family carried the decision routingTier?: 'efficient' | 'capable' | string; + routingTierRaw?: string; // tier as emitted, before the two vocabularies are folded + routedModel?: string; // model the router actually dispatched to + classifierModel?: string; // LLM that made the routing decision + routerType?: string; // router strategy, e.g. 'complexity' + routerScore?: number; // numeric score from LiteLLM's heuristic scorer routingConfidence?: number; routingSource?: 'stage_router' | 'judge' | 'classifier' | string; decisionSource?: string; @@ -89,12 +95,19 @@ export interface SessionCost { perModel: ModelCost[]; priced: boolean; // true if the native log was found & parsed hadLog: boolean; // true if a native log path was located (priced): UsageReco return out; } +function parseOptInt(s: string | undefined): number | undefined { + if (s == null) return undefined; + const v = parseInt(s, 10); + return Number.isNaN(v) ? undefined : v; +} + +function parseOptFloat(s: string | undefined): number | undefined { + if (s == null) return undefined; + const v = parseFloat(s); + return Number.isNaN(v) ? undefined : v; +} + function normalizeRoutingTier(raw: string | undefined): string | undefined { if (raw == null) return undefined; const key = String(raw).trim().toLowerCase(); @@ -308,11 +352,24 @@ export function extractClaudeUsageRecords(parsed: ParsedSession): UsageRecord[] const codemieDecisionSource = normalizeDecisionSource(msg?.x_codemie_routing_decision_source); const litellmCause = normalizeDecisionSource(msg?.['x-litellm-router-cause']); - const requestedModel = msg?.x_codemie_requested_model; + // Discriminate between the two routing header families so consumers know if routing cost is known. + const routingFamily: 'switchyard' | 'litellm' | undefined = + codemieTier != null || msg?.x_codemie_routing_source != null + ? 'switchyard' + : litellmTier != null || litellmCause != null + ? 'litellm' + : undefined; + + const requestedModel = msg?.x_codemie_requested_model ?? msg?.['x-litellm-router-model-name']; const capableModel = msg?.x_codemie_routing_capable_model; const routingTier = codemieTier ?? litellmTier ?? signals.routingTier; + const routingTierRaw = msg?.x_codemie_routing_tier ?? msg?.['x-litellm-router-tier']; const decisionSource = codemieDecisionSource ?? litellmCause ?? signals.decisionSource; const routingSource = msg?.x_codemie_routing_source ?? (decisionSource === 'llm-classifier' ? 'judge' : (decisionSource != null ? 'stage_router' : undefined)); + const routedModel = msg?.x_codemie_routing_capable_model ?? msg?.['x-litellm-router-routed-model']; + const classifierModel = msg?.['x-litellm-router-classifier-model']; + const routerType = msg?.['x-litellm-router-type']; + const routerScore = parseOptFloat(msg?.['x-litellm-router-score']); const rawConfidence = msg?.x_codemie_routing_confidence; const routingConfidence: number | undefined = (() => { if (rawConfidence != null) { @@ -330,16 +387,6 @@ export function extractClaudeUsageRecords(parsed: ParsedSession): UsageRecord[] } return undefined; })(); - const parseOptInt = (s: string | undefined): number | undefined => { - if (s == null) return undefined; - const v = parseInt(s, 10); - return Number.isNaN(v) ? undefined : v; - }; - const parseOptFloat = (s: string | undefined): number | undefined => { - if (s == null) return undefined; - const v = parseFloat(s); - return Number.isNaN(v) ? undefined : v; - }; const judgeInputTokens = parseOptInt(raw.message?.x_codemie_routing_judge_input_tokens ?? raw.message?.x_codemie_routing_classifier_input_tokens); const judgeOutputTokens = parseOptInt(raw.message?.x_codemie_routing_judge_output_tokens ?? raw.message?.x_codemie_routing_classifier_output_tokens); const judgeCachedTokens = parseOptInt(raw.message?.x_codemie_routing_judge_cached_tokens ?? raw.message?.x_codemie_routing_classifier_cached_tokens); @@ -356,17 +403,27 @@ export function extractClaudeUsageRecords(parsed: ParsedSession): UsageRecord[] const judgePrimaryRule = raw.message?.x_codemie_routing_judge_primary_rule ?? raw.message?.x_codemie_routing_classifier_primary_rule; const judgeCapabilityBoundary = raw.message?.x_codemie_routing_judge_capability_boundary ?? raw.message?.x_codemie_routing_classifier_capability_boundary; + appendDedupedRecord(records, keyedRecords, { key, ts, model, usage: { input, output, cacheRead, cacheCreation, cacheCreation1h, total: input + output + cacheRead + cacheCreation }, ...(requestedModel != null && { requestedModel }), + ...(routingFamily != null && { routingFamily }), ...(routingTier != null && { routingTier }), + ...(routingTierRaw != null && { routingTierRaw }), ...(capableModel != null && { capableModel }), + ...(routedModel != null && { routedModel }), + ...(classifierModel != null && { classifierModel }), + ...(routerType != null && { routerType }), + ...(routerScore != null && { routerScore }), ...(routingConfidence != null && { routingConfidence }), ...(routingSource != null && { routingSource }), ...(decisionSource != null && { decisionSource }), + // Switchyard reports classifier cost explicitly (possibly $0); LiteLLM never reports it at + // all. Record which is which so a session total doesn't read a LiteLLM "unmeasured" as "free". + ...(routingFamily != null && { routingCostKnown: routingFamily === 'switchyard' }), ...(judgeInputTokens != null && { judgeInputTokens }), ...(judgeOutputTokens != null && { judgeOutputTokens }), ...(judgeCachedTokens != null && { judgeCachedTokens }), diff --git a/src/cli/commands/analytics/report/client/app.js b/src/cli/commands/analytics/report/client/app.js index e521014e1..5f0bd53e9 100644 --- a/src/cli/commands/analytics/report/client/app.js +++ b/src/cli/commands/analytics/report/client/app.js @@ -732,17 +732,18 @@ }); host.appendChild(grid); - // Routing KPI section — only shown when at least one session has judge routing cost. - var routedSessions = fs.filter(function (s) { return s.judgeCostUSD != null && s.judgeCostUSD > 0; }); + // Routing KPI section — only shown when at least one session has routing. + var routedSessions = fs.filter(function (s) { return s.judgeCostUSD != null || s.routingCostKnown === false; }); if (routedSessions.length > 0) { host.appendChild(el('h3', 'section-title', 'Routing')); var rsGrid = el('div', 'kpi-grid'); rsGrid.style.gridTemplateColumns = 'repeat(3,1fr)'; - var totalJudgeCost = sum(routedSessions, function (s) { return s.judgeCostUSD || 0; }); - var totalJudgeIn = routedSessions.reduce(function (acc, s) { return acc + (s.judgeInputTokens || 0); }, 0); - var totalJudgeOut = routedSessions.reduce(function (acc, s) { return acc + (s.judgeOutputTokens || 0); }, 0); + var measuredSessions = routedSessions.filter(function (s) { return s.judgeCostUSD != null && s.judgeCostUSD > 0; }); + var totalJudgeCost = sum(measuredSessions, function (s) { return s.judgeCostUSD || 0; }); + var totalJudgeIn = measuredSessions.reduce(function (acc, s) { return acc + (s.judgeInputTokens || 0); }, 0); + var totalJudgeOut = measuredSessions.reduce(function (acc, s) { return acc + (s.judgeOutputTokens || 0); }, 0); [ - ['Sessions with routing', fmtNum(routedSessions.length) + ' / ' + fmtNum(fs.length)], - ['Judge routing cost', fmtUSD(totalJudgeCost)], + ['Sessions with routing', fmtNum(routedSessions.length) + ' / ' + fmtNum(fs.length) + (routedSessions.some(function(s) { return s.routingCostKnown === false; }) ? ' (cost unmeasured on ' + fmtNum(routedSessions.filter(function(s) { return s.routingCostKnown === false; }).length) + ')' : '')], + ['Judge routing cost', fmtUSD(totalJudgeCost) + (measuredSessions.length < routedSessions.length ? ' (LiteLLM unmeasured)' : '')], ['Judge tokens (in/out)', fmtTokens(totalJudgeIn) + ' / ' + fmtTokens(totalJudgeOut)] ].forEach(function (k) { var c = el('div', 'kpi'); c.innerHTML = '
' + k[0] + '
' + k[1] + '
'; rsGrid.appendChild(c); @@ -1156,7 +1157,10 @@ ['Started', '' + esc(fmtWhen(s.startTime)) + '', ''] ]; if (s.judgeCostUSD != null) { - costRows.push(['Routing (judge)', fmtUSD(s.judgeCostUSD), 'included in cost']); + var detail = s.routingCostKnown === false ? 'unmeasured (LiteLLM)' : 'included in cost'; + costRows.push(['Routing (judge)', s.judgeCostUSD > 0 ? fmtUSD(s.judgeCostUSD) : '—', detail]); + } else if (s.routingCostKnown === false) { + costRows.push(['Routing (judge)', '—', 'unmeasured (LiteLLM, no cost headers)']); } if (s.premiumRequests !== undefined) { costRows.push(['Premium requests', fmtNum(s.premiumRequests), 'provider billing unit']); diff --git a/src/cli/commands/analytics/report/payload-builder.ts b/src/cli/commands/analytics/report/payload-builder.ts index 6d0b436a7..ab623965a 100644 --- a/src/cli/commands/analytics/report/payload-builder.ts +++ b/src/cli/commands/analytics/report/payload-builder.ts @@ -115,6 +115,7 @@ export function buildPayload( ...(cost?.judgeOutputTokens != null ? { judgeOutputTokens: cost.judgeOutputTokens } : {}), ...(cost?.judgeCachedTokens != null ? { judgeCachedTokens: cost.judgeCachedTokens } : {}), ...(cost?.judgeCacheCreationTokens != null ? { judgeCacheCreationTokens: cost.judgeCacheCreationTokens } : {}), + ...(cost?.routingCostKnown != null ? { routingCostKnown: cost.routingCostKnown } : {}), ...(cost?.dispatches && cost.dispatches.length ? { dispatches: cost.dispatches } : {}), skillInvocations, agentInvocations, diff --git a/src/cli/commands/analytics/report/types.ts b/src/cli/commands/analytics/report/types.ts index db290a12a..86ef35861 100644 --- a/src/cli/commands/analytics/report/types.ts +++ b/src/cli/commands/analytics/report/types.ts @@ -45,12 +45,19 @@ export interface ReportSessionRecord { modelTimeline?: ModelTimelinePoint[]; // per-turn model + routing metadata; absent when no routing data dispatches?: DispatchEvent[]; // timed top-level agent/skill/command invocations; absent when none - // === Switchyard routing classifier cost (included in costUSD) === + // === Routing classifier cost (included in costUSD; Switchyard only — see routingCostKnown) === judgeCostUSD?: number; // USD spent on the routing classifier LLM judgeInputTokens?: number; judgeOutputTokens?: number; judgeCachedTokens?: number; judgeCacheCreationTokens?: number; + /** + * True when every routed turn in this session came from a family that reports classifier + * cost (Switchyard). False when any turn used LiteLLM, which never reports cost — so + * `judgeCostUSD` is absent even though classifiers ran. Absent when the session had no + * routed turns at all. + */ + routingCostKnown?: boolean; // === Usage provenance (optional; absent for agents that always record full usage) === /** diff --git a/src/providers/plugins/sso/proxy/plugins/__tests__/routing-header-injector.plugin.test.ts b/src/providers/plugins/sso/proxy/plugins/__tests__/routing-header-injector.plugin.test.ts new file mode 100644 index 000000000..e14238177 --- /dev/null +++ b/src/providers/plugins/sso/proxy/plugins/__tests__/routing-header-injector.plugin.test.ts @@ -0,0 +1,243 @@ +import { describe, it, expect } from 'vitest'; +import { + extractRoutingHeaders, + injectIntoJsonBody, + injectIntoSseChunk, +} from '../routing-header-injector.plugin.js'; + +describe('RoutingHeaderInjectorPlugin', () => { + describe('extractRoutingHeaders', () => { + it('captures x-litellm-* headers with hyphens intact', () => { + const headers = { + 'x-litellm-router-tier': 'COMPLEX', + 'x-litellm-router-cause': 'llm_classifier', + 'x-litellm-model-name': 'claude-sonnet-5', + }; + const result = extractRoutingHeaders(headers); + expect(result).toEqual({ + 'x-litellm-router-tier': 'COMPLEX', + 'x-litellm-router-cause': 'llm_classifier', + 'x-litellm-model-name': 'claude-sonnet-5', + }); + }); + + it('captures x-codemie-routing-* and x-codemie-requested-* headers, converting hyphens to underscores', () => { + const headers = { + 'x-codemie-routing-tier': 'efficient', + 'x-codemie-routing-source': 'judge', + 'x-codemie-requested-model': 'claude-sonnet-4-6', + 'x-codemie-routing-judge-cost-usd': '0.002475', + }; + const result = extractRoutingHeaders(headers); + expect(result).toEqual({ + 'x_codemie_routing_tier': 'efficient', + 'x_codemie_routing_source': 'judge', + 'x_codemie_requested_model': 'claude-sonnet-4-6', + 'x_codemie_routing_judge_cost_usd': '0.002475', + }); + }); + + it('ignores non-routing headers', () => { + const headers = { + 'authorization': 'Bearer secret-token', + 'content-type': 'application/json', + 'x-custom-header': 'value', + 'x-litellm-router-tier': 'SIMPLE', + }; + const result = extractRoutingHeaders(headers); + expect(result).toEqual({ + 'x-litellm-router-tier': 'SIMPLE', + }); + }); + + it('handles array-valued headers by taking the first element', () => { + const headers = { + 'x-litellm-router-signals': ['["llm-classifier:COMPLEX"]', 'ignored'], + }; + const result = extractRoutingHeaders(headers); + expect(result).toEqual({ + 'x-litellm-router-signals': '["llm-classifier:COMPLEX"]', + }); + }); + + it('ignores null and undefined values', () => { + const headers = { + 'x-litellm-router-tier': 'COMPLEX', + 'x-litellm-missing': null as unknown as string, + 'x-codemie-routing-tier': undefined as unknown as string, + }; + const result = extractRoutingHeaders(headers); + expect(result).toEqual({ + 'x-litellm-router-tier': 'COMPLEX', + }); + }); + + it('returns empty object when no routing headers are present', () => { + const headers = { + 'authorization': 'Bearer token', + 'content-type': 'application/json', + }; + const result = extractRoutingHeaders(headers); + expect(result).toEqual({}); + }); + + it('case-insensitive matching for header names', () => { + const headers = { + 'X-LiteLLM-Router-Tier': 'COMPLEX', + 'X-CodeMie-Routing-Tier': 'efficient', + }; + const result = extractRoutingHeaders(headers); + expect(result).toEqual({ + 'x-litellm-router-tier': 'COMPLEX', + 'x_codemie_routing_tier': 'efficient', + }); + }); + }); + + describe('injectIntoJsonBody', () => { + it('merges injections into a JSON object', () => { + const body = Buffer.from(JSON.stringify({ id: 'msg_123', role: 'assistant', content: [] })); + const injections = { + 'x-litellm-router-tier': 'COMPLEX', + 'x-litellm-model-name': 'claude-sonnet-5', + }; + const result = injectIntoJsonBody(body, injections); + const parsed = JSON.parse(result.toString('utf-8')); + expect(parsed).toEqual({ + id: 'msg_123', + role: 'assistant', + content: [], + 'x-litellm-router-tier': 'COMPLEX', + 'x-litellm-model-name': 'claude-sonnet-5', + }); + }); + + it('returns original buffer when injections are empty', () => { + const body = Buffer.from(JSON.stringify({ id: 'msg_123' })); + const result = injectIntoJsonBody(body, {}); + expect(result).toBe(body); + }); + + it('returns original buffer when body is not a JSON object', () => { + const testCases = [ + Buffer.from('not json'), + Buffer.from('[]'), + Buffer.from('null'), + Buffer.from('123'), + Buffer.from('true'), + ]; + const injections = { 'x-litellm-router-tier': 'COMPLEX' }; + for (const body of testCases) { + const result = injectIntoJsonBody(body, injections); + expect(result).toBe(body); + } + }); + + it('returns original buffer on parse error', () => { + const body = Buffer.from('{broken json}'); + const injections = { 'x-litellm-router-tier': 'COMPLEX' }; + const result = injectIntoJsonBody(body, injections); + expect(result).toBe(body); + }); + + it('overwrites existing keys if injections have the same key', () => { + const body = Buffer.from(JSON.stringify({ id: 'msg_123', 'x-litellm-router-tier': 'old' })); + const injections = { 'x-litellm-router-tier': 'COMPLEX' }; + const result = injectIntoJsonBody(body, injections); + const parsed = JSON.parse(result.toString('utf-8')); + expect(parsed['x-litellm-router-tier']).toBe('COMPLEX'); + }); + }); + + describe('injectIntoSseChunk', () => { + it('injects fields into the message object of a message_start event', () => { + const chunk = Buffer.from( + 'data: ' + + JSON.stringify({ + type: 'message_start', + message: { id: 'msg_123', role: 'assistant', model: 'claude-sonnet-5' }, + }) + ); + const injections = { + 'x-litellm-router-tier': 'COMPLEX', + 'x-litellm-model-name': 'claude-sonnet-5', + }; + const result = injectIntoSseChunk(chunk, injections); + const line = result.toString('utf-8'); + const parsed = JSON.parse(line.slice(6)); + expect(parsed.message).toEqual({ + id: 'msg_123', + role: 'assistant', + model: 'claude-sonnet-5', + 'x-litellm-router-tier': 'COMPLEX', + 'x-litellm-model-name': 'claude-sonnet-5', + }); + }); + + it('preserves non-message_start events unchanged', () => { + const chunk = Buffer.from( + 'data: ' + JSON.stringify({ type: 'content_block_delta', delta: { type: 'text_delta' } }) + ); + const injections = { 'x-litellm-router-tier': 'COMPLEX' }; + const result = injectIntoSseChunk(chunk, injections); + expect(result).toBe(chunk); + }); + + it('preserves non-JSON lines unchanged', () => { + const chunk = Buffer.from('data: [DONE]'); + const injections = { 'x-litellm-router-tier': 'COMPLEX' }; + const result = injectIntoSseChunk(chunk, injections); + expect(result).toBe(chunk); + }); + + it('preserves other event types in a multi-line chunk', () => { + const multiLine = `data: ${JSON.stringify({ type: 'message_start', message: { id: 'msg_1' } })} +data: ${JSON.stringify({ type: 'content_block_delta', delta: { type: 'text_delta', text: 'hi' } })} +data: [DONE]`; + const chunk = Buffer.from(multiLine); + const injections = { 'x-litellm-router-tier': 'COMPLEX' }; + const result = injectIntoSseChunk(chunk, injections); + const lines = result.toString('utf-8').split('\n'); + // First line should be modified, others unchanged + expect(JSON.parse(lines[0].slice(6)).message['x-litellm-router-tier']).toBe('COMPLEX'); + expect(lines[1]).toBe( + `data: ${JSON.stringify({ type: 'content_block_delta', delta: { type: 'text_delta', text: 'hi' } })}` + ); + expect(lines[2]).toBe('data: [DONE]'); + }); + + it('returns original buffer when injections are empty', () => { + const chunk = Buffer.from( + 'data: ' + JSON.stringify({ type: 'message_start', message: { id: 'msg_123' } }) + ); + const result = injectIntoSseChunk(chunk, {}); + expect(result).toBe(chunk); + }); + + it('returns original buffer when no message_start is present', () => { + const chunk = Buffer.from('data: ' + JSON.stringify({ type: 'content_block_delta' })); + const injections = { 'x-litellm-router-tier': 'COMPLEX' }; + const result = injectIntoSseChunk(chunk, injections); + expect(result).toBe(chunk); + }); + + it('handles partial lines gracefully (does not parse or modify)', () => { + const chunk = Buffer.from('data: {"type":"message_st'); + const injections = { 'x-litellm-router-tier': 'COMPLEX' }; + const result = injectIntoSseChunk(chunk, injections); + expect(result).toBe(chunk); + }); + + it('injects into all message_start events in a multi-line chunk', () => { + const multiStart = `data: ${JSON.stringify({ type: 'message_start', message: { id: 'msg_1' } })} +data: ${JSON.stringify({ type: 'message_start', message: { id: 'msg_2' } })}`; + const chunk = Buffer.from(multiStart); + const injections = { 'x-litellm-router-tier': 'COMPLEX' }; + const result = injectIntoSseChunk(chunk, injections); + const lines = result.toString('utf-8').split('\n'); + // Both get injected (both are message_start events) + expect(JSON.parse(lines[0].slice(6)).message['x-litellm-router-tier']).toBe('COMPLEX'); + expect(JSON.parse(lines[1].slice(6)).message['x-litellm-router-tier']).toBe('COMPLEX'); + }); + }); +}); diff --git a/src/providers/plugins/sso/proxy/plugins/index.ts b/src/providers/plugins/sso/proxy/plugins/index.ts index f1ab3a4a3..f043ecbd6 100644 --- a/src/providers/plugins/sso/proxy/plugins/index.ts +++ b/src/providers/plugins/sso/proxy/plugins/index.ts @@ -20,6 +20,7 @@ import { CodexEncryptedContentSanitizerPlugin } from './codex-encrypted-content- import { CopilotEncryptedContentSanitizerPlugin } from './copilot-encrypted-content-sanitizer.plugin.js'; import { VsCodeRequestNormalizerPlugin } from './vscode-request-normalizer.plugin.js'; import { LoggingPlugin } from './logging.plugin.js'; +import { RoutingHeaderInjectorPlugin } from './routing-header-injector.plugin.js'; import { SSOSessionSyncPlugin } from './sso.session-sync.plugin.js'; /** @@ -44,6 +45,7 @@ export function registerCorePlugins(): void { registry.register(new VsCodeRequestNormalizerPlugin()); // Priority 17 - constrains VS Code user identifiers registry.register(new HeaderInjectionPlugin()); registry.register(new LoggingPlugin()); // Always enabled - logs to log files at INFO level + registry.register(new RoutingHeaderInjectorPlugin()); // Priority 55 - copies router decision headers onto the response body so agents persist them registry.register(new SSOSessionSyncPlugin()); // Priority 100 - syncs sessions via multiple processors } @@ -66,6 +68,7 @@ export { CopilotEncryptedContentSanitizerPlugin, VsCodeRequestNormalizerPlugin, LoggingPlugin, + RoutingHeaderInjectorPlugin, }; export { SSOSessionSyncPlugin } from './sso.session-sync.plugin.js'; export { getPluginRegistry, resetPluginRegistry } from './registry.js'; diff --git a/src/providers/plugins/sso/proxy/plugins/routing-header-injector.plugin.ts b/src/providers/plugins/sso/proxy/plugins/routing-header-injector.plugin.ts new file mode 100644 index 000000000..3087140ac --- /dev/null +++ b/src/providers/plugins/sso/proxy/plugins/routing-header-injector.plugin.ts @@ -0,0 +1,193 @@ +/** + * Routing Header Injector Plugin + * Priority: 55 (after logging at 50, before session-sync at 100) + * + * The upstream router (CodeMie Switchyard or the LiteLLM router) reports which model + * tier it picked, and why, in HTTP *response headers*. Headers are forwarded downstream + * but agents do not persist them, so the decision is lost the moment the turn ends. + * + * This plugin copies those headers onto the response *body*, where the agent stores them + * verbatim in its own transcript alongside `usage`. The analytics pipeline + * (cost/usage-readers.ts) then reads routing metadata per turn with no join and no + * sidecar file. + * + * Two response paths: + * JSON (non-streaming) — buffer the body, merge fields at the top level (which is the + * message object for the Messages API), return a new response. + * SSE (streaming) — stash headers in `context.metadata` during onUpstreamResponse, + * then merge them into the first `message_start` event's nested + * `message` object as it streams past. + * + * Header → body key mapping (matches what the analytics reader expects): + * x-litellm-* → key keeps its hyphens (`x-litellm-router-tier`) + * x-codemie-* → key converts hyphens to `_` (`x_codemie_routing_tier`) + * + * Only one family is ever present: a deployment routes through Switchyard or the LiteLLM + * router, not both. Capturing by prefix means a new field in either family is recorded + * without a code change here. + */ + +import { IncomingHttpHeaders, IncomingMessage } from 'http'; +import { ProxyPlugin, PluginContext, ProxyInterceptor, UpstreamResponseTools } from './types.js'; +import { ProxyContext } from '../proxy-types.js'; +import { logger } from '../../../../../utils/logger.js'; + +const LITELLM_PREFIX = 'x-litellm-'; +const CODEMIE_ROUTING_PREFIXES = ['x-codemie-routing-', 'x-codemie-requested-']; +const METADATA_HEADERS_KEY = '_routingInjectionHeaders'; +const METADATA_INJECTED_KEY = '_routingInjected'; + +/** Fields extracted from routing headers, keyed as they will appear in the body. */ +export type RoutingInjections = Record; + +/** + * Extract routing-relevant upstream response headers into a flat body-key → value map. + * Returns an empty object when the response carries no routing metadata, which is the + * common case for non-routed deployments and for requests that name a literal model ID. + */ +export function extractRoutingHeaders(headers: IncomingHttpHeaders): RoutingInjections { + const out: RoutingInjections = {}; + for (const [key, value] of Object.entries(headers)) { + if (value == null) continue; + const raw = Array.isArray(value) ? value[0] : value; + if (raw == null) continue; + const lower = key.toLowerCase(); + if (lower.startsWith(LITELLM_PREFIX)) { + out[lower] = raw; + } else if (CODEMIE_ROUTING_PREFIXES.some((p) => lower.startsWith(p))) { + out[lower.replace(/-/g, '_')] = raw; + } + } + return out; +} + +/** + * Merge fields into the top-level JSON object of a body buffer. Returns the buffer + * unchanged when there is nothing to inject, the payload is not a JSON object, or + * parsing fails — a routing annotation must never corrupt a response. + */ +export function injectIntoJsonBody(body: Buffer, injections: RoutingInjections): Buffer { + if (Object.keys(injections).length === 0) return body; + try { + const parsed: unknown = JSON.parse(body.toString('utf-8')); + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) { + return body; + } + return Buffer.from(JSON.stringify({ ...parsed, ...injections }), 'utf-8'); + } catch { + return body; + } +} + +/** + * Merge fields into the `message` object of a `message_start` SSE event. + * + * Rewrites only whole `data:` lines that parse as a `message_start` event; every other + * line — including partial trailing lines at a chunk boundary — is passed through byte + * for byte. Returns the original buffer when no event was modified. + */ +export function injectIntoSseChunk(chunk: Buffer, injections: RoutingInjections): Buffer { + if (Object.keys(injections).length === 0) return chunk; + const lines = chunk.toString('utf-8').split('\n'); + let modified = false; + const outLines: string[] = []; + + for (const line of lines) { + if (!line.startsWith('data: ')) { + outLines.push(line); + continue; + } + let parsed: unknown; + try { + parsed = JSON.parse(line.slice(6)); + } catch { + // Not JSON (e.g. `[DONE]`) or a line split across chunks — pass through untouched. + outLines.push(line); + continue; + } + const event = parsed as { type?: unknown; message?: unknown }; + if ( + typeof parsed === 'object' && + parsed !== null && + !Array.isArray(parsed) && + event.type === 'message_start' && + typeof event.message === 'object' && + event.message !== null + ) { + const message = { ...(event.message as Record), ...injections }; + outLines.push('data: ' + JSON.stringify({ ...event, message })); + modified = true; + } else { + outLines.push(line); + } + } + + return modified ? Buffer.from(outLines.join('\n'), 'utf-8') : chunk; +} + +export class RoutingHeaderInjectorPlugin implements ProxyPlugin { + id = '@codemie/proxy-routing-header-injector'; + name = 'Routing Header Injector'; + version = '1.0.0'; + priority = 55; // After logging (50), before session-sync (100) + + async createInterceptor(_context: PluginContext): Promise { + return new RoutingHeaderInjectorInterceptor(); + } +} + +class RoutingHeaderInjectorInterceptor implements ProxyInterceptor { + name = 'routing-header-injector'; + + async onUpstreamResponse( + context: ProxyContext, + response: IncomingMessage, + tools: UpstreamResponseTools + ): Promise { + const injections = extractRoutingHeaders(response.headers); + if (Object.keys(injections).length === 0) { + return response; + } + + const contentType = String(response.headers['content-type'] ?? '').toLowerCase(); + if (contentType.includes('text/event-stream')) { + // Streaming: defer to onResponseChunk so the stream is never buffered. + context.metadata[METADATA_HEADERS_KEY] = injections; + logger.debug( + `[${this.name}] Captured ${Object.keys(injections).length} routing header(s) for SSE injection` + ); + return response; + } + + try { + const body = await tools.readBody(response); + const modified = injectIntoJsonBody(body, injections); + if (modified !== body) { + logger.debug( + `[${this.name}] Injected ${Object.keys(injections).length} routing header(s) into JSON response` + ); + } + return tools.fromBuffer(response, modified); + } catch (error) { + logger.debug(`[${this.name}] JSON injection failed, forwarding original response:`, error); + return response; + } + } + + async onResponseChunk(context: ProxyContext, chunk: Buffer): Promise { + if (context.metadata[METADATA_INJECTED_KEY]) { + return chunk; + } + const injections = context.metadata[METADATA_HEADERS_KEY] as RoutingInjections | undefined; + if (!injections || Object.keys(injections).length === 0) { + return chunk; + } + + const modified = injectIntoSseChunk(chunk, injections); + if (modified !== chunk) { + context.metadata[METADATA_INJECTED_KEY] = true; + logger.debug(`[${this.name}] Injected routing headers into SSE message_start`); + } + return modified; + } +}