diff --git a/README.md b/README.md index 4ec5eb4..df7fee8 100644 --- a/README.md +++ b/README.md @@ -81,7 +81,10 @@ budget.yaml ───┘ │ (distribution The proxy captures calls at the **HTTP layer** (you point your agent's `base_url` at it), so Augur is framework- and language-agnostic and records the *real* call graph — retries and fan-out included — without you rewriting the -agent. +agent. It speaks three wire dialects natively — **OpenAI**, **Anthropic** +(`/v1/messages`, including the cache-read/cache-write token split), and **Gemini** +(`usageMetadata`) — so you point it at whichever provider your agent already +uses. --- @@ -271,14 +274,21 @@ example. ## Honest limitations -- **Streaming usage** is captured exactly only when the client sets - `stream_options.include_usage` (the provider then emits a usage block). Without - it, Augur records a zero-token row rather than tokenizing on the proxy side — - proxy-side tokenization is fragile (per-model tokenizers) and `include_usage` - is the correct, exact path. -- **Multipliers.** Augur reports *calls per run* — the multiplier it can observe - truthfully. Classifying those calls into retries vs sub-agent fan-out needs - labeling the trace does not yet carry. +- **Streaming usage.** For OpenAI streaming, the provider only emits a usage + block when the request sets `stream_options.include_usage`. Augur injects that + flag automatically (disable with `--inject-usage=false`), so streamed calls are + captured exactly rather than as a zero-token row — the trade-off is one extra + usage-only SSE chunk the agent receives (benign for standard clients). + Anthropic and Gemini report usage on every streamed response with no opt-in. +- **Multipliers.** Augur reports *calls per run* and splits it into + *retries per run* — a call whose request body is byte-identical to an earlier + one in the same run, the observable signature of a client-library retry — vs + the rest. It still does **not** isolate sub-agent fan-out from sequential + tool-loop steps: both are distinct-body calls, and call concurrency isn't + visible at the HTTP layer. +- **Cache-write pricing** is modelled from `cache_write` in `pricing.yaml` + (Anthropic's ~1.25x input premium on `cache_creation_input_tokens`). It is a + dated snapshot like every other price — verify it against the current page. - **Running the agent in CI spends real tokens.** Keep the scenario set small and `runs` modest, or record once and replay (`--record`/`--replay`) so CI pushes spend nothing. @@ -317,6 +327,11 @@ LangChain **callback shim** that writes Augur's trace schema from the usage ever call reports — *without editing CloudOracle*. Full walkthrough and harness: [`examples/cloudoracle/`](examples/cloudoracle/). +(The proxy now speaks Anthropic's `/v1/messages` natively, so an Anthropic agent +that *can* set a `base_url` — e.g. LangChain's `ChatAnthropic(base_url=…)` — can +take the HTTP path directly and skip the shim; the shim remains the fallback for +frameworks that don't expose a base URL at all.) + **What it found (Claude Haiku 4.5, 20 runs): gate PASS** at `$/request p95 $0.0198` (budget $0.02). The headline is the `find-savings` scenario — its **p95 cost is 2.3× its median**, driven by a call-count tail (5 → 13 calls/run when the diff --git a/aggregate/aggregate.go b/aggregate/aggregate.go index 3056863..5232b38 100644 --- a/aggregate/aggregate.go +++ b/aggregate/aggregate.go @@ -12,22 +12,26 @@ import ( // within a scenario — the "decompose by model" view that shows which model // drives a scenario's bill. type ModelUsage struct { - Model string `json:"model"` - Calls int `json:"calls"` - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` - CachedTokens int `json:"cached_tokens"` - CostUSD float64 `json:"cost_usd"` + Model string `json:"model"` + Calls int `json:"calls"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CachedTokens int `json:"cached_tokens"` + CacheWriteTokens int `json:"cache_write_tokens,omitempty"` + CostUSD float64 `json:"cost_usd"` } // Run is the cost of a single (scenario, run): the sum over every LLM call the // agent made during that one execution. Kept individually so the checkpoint can // reconcile per-run totals against the raw trace by hand. type Run struct { - ScenarioID string `json:"scenario_id"` - RunID string `json:"run_id"` - Calls int `json:"calls"` - CostUSD float64 `json:"cost_usd"` + ScenarioID string `json:"scenario_id"` + RunID string `json:"run_id"` + Calls int `json:"calls"` + // Retries is how many of this run's calls repeated an earlier call's request + // body (trace.KindRetry) — the observed retry count for the run. + Retries int `json:"retries"` + CostUSD float64 `json:"cost_usd"` } // Scenario aggregates all runs of one scenario into the cost distribution and @@ -39,10 +43,14 @@ type Scenario struct { // projection engine and the gate ultimately care about. CostPerRun Distribution `json:"cost_per_run_usd"` // CallsPerRun is the distribution of how many LLM calls one run made — the - // primary OBSERVED agentic multiplier. Retry/fan-out classification needs - // labeling the trace does not yet carry; calls-per-run is what we can state - // truthfully from the data. + // primary OBSERVED agentic multiplier. CallsPerRun Distribution `json:"calls_per_run"` + // RetriesPerRun is the distribution of how many of a run's calls were + // retries (trace.KindRetry: a request body identical to an earlier call in + // the same run). It splits the calls-per-run multiplier into the part driven + // by client-library retries versus genuinely new calls. Fan-out is not + // isolated — the proxy cannot observe call concurrency. + RetriesPerRun Distribution `json:"retries_per_run"` // ByModel is the per-model breakdown, sorted by cost descending. ByModel []ModelUsage `json:"by_model"` // TotalCost is the summed cost of every run of this scenario. @@ -127,9 +135,10 @@ func AggregateWithKnobs(records []trace.Record, pricing cost.Pricing, knobs Knob for _, rec := range records { u := cost.Usage{ - InputTokens: rec.InputTokens, - OutputTokens: rec.OutputTokens, - CachedTokens: rec.CachedTokens, + InputTokens: rec.InputTokens, + OutputTokens: rec.OutputTokens, + CachedTokens: rec.CachedTokens, + CacheWriteTokens: rec.CacheWriteTokens, } b, err := pricing.Breakdown(rec.Model, u) if err != nil { @@ -146,6 +155,9 @@ func AggregateWithKnobs(records []trace.Record, pricing cost.Pricing, knobs Knob runOrder = append(runOrder, k) } r.Calls++ + if rec.Kind == trace.KindRetry { + r.Retries++ + } r.CostUSD += c models := scenarioModels[rec.ScenarioID] @@ -162,6 +174,7 @@ func AggregateWithKnobs(records []trace.Record, pricing cost.Pricing, knobs Knob mu.InputTokens += rec.InputTokens mu.OutputTokens += rec.OutputTokens mu.CachedTokens += rec.CachedTokens + mu.CacheWriteTokens += rec.CacheWriteTokens mu.CostUSD += c } @@ -176,19 +189,22 @@ func AggregateWithKnobs(records []trace.Record, pricing cost.Pricing, knobs Knob for id, rs := range scenarioRuns { costs := make([]float64, len(rs)) calls := make([]float64, len(rs)) + retries := make([]float64, len(rs)) var total float64 for i, r := range rs { costs[i] = r.CostUSD calls[i] = float64(r.Calls) + retries[i] = float64(r.Retries) total += r.CostUSD } scenarios = append(scenarios, Scenario{ - ScenarioID: id, - Runs: len(rs), - CostPerRun: Summarize(costs), - CallsPerRun: Summarize(calls), - ByModel: sortedModels(scenarioModels[id]), - TotalCost: total, + ScenarioID: id, + Runs: len(rs), + CostPerRun: Summarize(costs), + CallsPerRun: Summarize(calls), + RetriesPerRun: Summarize(retries), + ByModel: sortedModels(scenarioModels[id]), + TotalCost: total, }) } sort.Slice(scenarios, func(i, j int) bool { diff --git a/aggregate/table.go b/aggregate/table.go index 0265436..1d9f65a 100644 --- a/aggregate/table.go +++ b/aggregate/table.go @@ -29,6 +29,12 @@ func (r Result) WriteTable(w io.Writer) error { s.CostPerRun.Mean, s.CostPerRun.P50, s.CostPerRun.P95, s.CostPerRun.Stdev, s.CostPerRun.Min, s.CostPerRun.Max) fmt.Fprintf(tw, " calls/run\t%.2f\t%.2f\t%.2f\t%.2f\t%.0f\t%.0f\n", s.CallsPerRun.Mean, s.CallsPerRun.P50, s.CallsPerRun.P95, s.CallsPerRun.Stdev, s.CallsPerRun.Min, s.CallsPerRun.Max) + // Only surface retries when any were observed — keeps the common + // no-retry trace's table uncluttered. + if s.RetriesPerRun.Max > 0 { + fmt.Fprintf(tw, " retries/run\t%.2f\t%.2f\t%.2f\t%.2f\t%.0f\t%.0f\n", + s.RetriesPerRun.Mean, s.RetriesPerRun.P50, s.RetriesPerRun.P95, s.RetriesPerRun.Stdev, s.RetriesPerRun.Min, s.RetriesPerRun.Max) + } if err := tw.Flush(); err != nil { return err } diff --git a/cost/cost.go b/cost/cost.go index d4a3e93..be37e1b 100644 --- a/cost/cost.go +++ b/cost/cost.go @@ -34,18 +34,25 @@ type ModelPrice struct { // cache. When a model has no cache discount this should equal Input; // LoadPricing fills it in that way when the field is omitted. CachedInput float64 + // CacheWrite is USD per Mtok for prompt tokens WRITTEN to the provider cache + // (Anthropic's cache_creation_input_tokens), typically a premium over Input. + // When a provider does not bill cache writes separately this equals Input; + // LoadPricing fills it in that way when the field is omitted. + CacheWrite float64 } // Usage is the token accounting for a single LLM call, mirroring how providers -// report it. CachedTokens is a SUBSET of InputTokens (the cached portion of the -// prompt), not an additional bucket — this matches OpenAI's +// report it. CachedTokens and CacheWriteTokens are both SUBSETS of InputTokens +// (portions of the prompt), not additional buckets — this matches OpenAI's // prompt_tokens / prompt_tokens_details.cached_tokens and Anthropic's -// cache_read_input_tokens. Billing therefore splits InputTokens into a cached -// part and a full-price part. +// cache_read_input_tokens / cache_creation_input_tokens. Billing therefore +// splits InputTokens into a cached-read part, a cache-write part, and a +// full-price remainder. type Usage struct { - InputTokens int // total prompt tokens, INCLUDING the cached portion - OutputTokens int // completion tokens - CachedTokens int // cached prompt tokens, billed at the cached rate + InputTokens int // total prompt tokens, INCLUDING the cached and cache-write portions + OutputTokens int // completion tokens + CachedTokens int // cached-read prompt tokens, billed at the cached rate + CacheWriteTokens int // cache-write prompt tokens, billed at the cache-write rate } // Validate reports whether the usage is internally consistent. Negative counts @@ -60,9 +67,11 @@ func (u Usage) Validate() error { return fmt.Errorf("cost: negative output tokens (%d)", u.OutputTokens) case u.CachedTokens < 0: return fmt.Errorf("cost: negative cached tokens (%d)", u.CachedTokens) - case u.CachedTokens > u.InputTokens: - return fmt.Errorf("cost: cached tokens (%d) exceed input tokens (%d)", - u.CachedTokens, u.InputTokens) + case u.CacheWriteTokens < 0: + return fmt.Errorf("cost: negative cache-write tokens (%d)", u.CacheWriteTokens) + case u.CachedTokens+u.CacheWriteTokens > u.InputTokens: + return fmt.Errorf("cost: cached (%d) + cache-write (%d) tokens exceed input tokens (%d)", + u.CachedTokens, u.CacheWriteTokens, u.InputTokens) } return nil } @@ -73,31 +82,37 @@ func (u Usage) Validate() error { type Breakdown struct { // InputUSD is the cost of the non-cached prompt tokens. InputUSD float64 - // CachedUSD is the cost of the cached prompt tokens. + // CachedUSD is the cost of the cached-read prompt tokens. CachedUSD float64 + // CacheWriteUSD is the cost of the cache-write prompt tokens. + CacheWriteUSD float64 // OutputUSD is the cost of the completion tokens. OutputUSD float64 } -// Total is the full call cost: the three components summed. -func (b Breakdown) Total() float64 { return b.InputUSD + b.CachedUSD + b.OutputUSD } +// Total is the full call cost: the four components summed. +func (b Breakdown) Total() float64 { + return b.InputUSD + b.CachedUSD + b.CacheWriteUSD + b.OutputUSD +} -// PromptUSD is the cost attributable to the prompt (input + cached) — the part -// that scales with context growth. -func (b Breakdown) PromptUSD() float64 { return b.InputUSD + b.CachedUSD } +// PromptUSD is the cost attributable to the prompt (input + cached + cache +// write) — the part that scales with context growth. +func (b Breakdown) PromptUSD() float64 { return b.InputUSD + b.CachedUSD + b.CacheWriteUSD } // Breakdown returns the per-component cost of a single call priced at p. The -// cached portion of the prompt is billed at CachedInput, the remainder at -// Input, and completion tokens at Output. It errors if the usage is invalid. +// cached-read portion of the prompt is billed at CachedInput, the cache-write +// portion at CacheWrite, the remainder at Input, and completion tokens at +// Output. It errors if the usage is invalid. func (p ModelPrice) Breakdown(u Usage) (Breakdown, error) { if err := u.Validate(); err != nil { return Breakdown{}, err } - fullInput := u.InputTokens - u.CachedTokens + fullInput := u.InputTokens - u.CachedTokens - u.CacheWriteTokens return Breakdown{ - InputUSD: float64(fullInput) / tokensPerMtok * p.Input, - CachedUSD: float64(u.CachedTokens) / tokensPerMtok * p.CachedInput, - OutputUSD: float64(u.OutputTokens) / tokensPerMtok * p.Output, + InputUSD: float64(fullInput) / tokensPerMtok * p.Input, + CachedUSD: float64(u.CachedTokens) / tokensPerMtok * p.CachedInput, + CacheWriteUSD: float64(u.CacheWriteTokens) / tokensPerMtok * p.CacheWrite, + OutputUSD: float64(u.OutputTokens) / tokensPerMtok * p.Output, }, nil } diff --git a/cost/cost_test.go b/cost/cost_test.go index 449256c..30f628f 100644 --- a/cost/cost_test.go +++ b/cost/cost_test.go @@ -162,3 +162,58 @@ func TestPricePresence(t *testing.T) { t.Error("Price(nope) ok = true, want false") } } + +// haiku mirrors the Anthropic snapshot: $1.00 input / $5.00 output / +// $0.10 cached-read / $1.25 cache-write per Mtok. +var haiku = ModelPrice{Input: 1.00, Output: 5.00, CachedInput: 0.10, CacheWrite: 1.25} + +func TestCacheWriteBilling(t *testing.T) { + // Hand calc for an Anthropic call with all three input buckets: + // full input = 1M - 300k read - 200k write = 500k @ $1.00 = 0.50 + // cache read = 300k @ $0.10 = 0.03 + // cache write = 200k @ $1.25 = 0.25 + // output = 100k @ $5.00 = 0.50 + // total = 1.28 + u := Usage{InputTokens: 1_000_000, OutputTokens: 100_000, CachedTokens: 300_000, CacheWriteTokens: 200_000} + b, err := haiku.Breakdown(u) + if err != nil { + t.Fatalf("Breakdown: %v", err) + } + if !approxEqual(b.InputUSD, 0.50) || !approxEqual(b.CachedUSD, 0.03) || + !approxEqual(b.CacheWriteUSD, 0.25) || !approxEqual(b.OutputUSD, 0.50) { + t.Errorf("breakdown = %+v, want in=0.50 cached=0.03 write=0.25 out=0.50", b) + } + if !approxEqual(b.Total(), 1.28) { + t.Errorf("Total = %v, want 1.28", b.Total()) + } + // Cache-write tokens are part of the prompt, so context growth scales them. + if !approxEqual(b.PromptUSD(), 0.78) { + t.Errorf("PromptUSD = %v, want 0.78 (0.50+0.03+0.25)", b.PromptUSD()) + } +} + +func TestCacheWriteValidation(t *testing.T) { + // cached + write may not exceed the total input. + u := Usage{InputTokens: 100, CachedTokens: 60, CacheWriteTokens: 60} + if err := u.Validate(); err == nil { + t.Error("expected error when cached+write exceed input, got nil") + } + // A negative write bucket is nonsensical. + if err := (Usage{InputTokens: 100, CacheWriteTokens: -1}).Validate(); err == nil { + t.Error("expected error for negative cache-write tokens, got nil") + } +} + +func TestCacheWriteDefaultsToInputRate(t *testing.T) { + // When cache_write is omitted, LoadPricing sets it to Input, so a model with + // no write premium bills cache-write tokens at the plain input rate. + yml := []byte("version: 1\nunit: per_mtok\nmodels:\n m:\n input: 2.0\n output: 4.0\n") + p, err := ParsePricing(yml) + if err != nil { + t.Fatalf("ParsePricing: %v", err) + } + mp, _ := p.Price("m") + if mp.CacheWrite != 2.0 { + t.Errorf("CacheWrite defaulted to %v, want 2.0 (the input rate)", mp.CacheWrite) + } +} diff --git a/cost/pricing.go b/cost/pricing.go index a7d795c..d3358f6 100644 --- a/cost/pricing.go +++ b/cost/pricing.go @@ -22,6 +22,10 @@ type yamlPricing struct { // discount → bill cached tokens at the full input rate) from an // explicit 0.0 (a genuinely free cache). CachedInput *float64 `yaml:"cached_input"` + // CacheWrite is a pointer for the same reason: omitted → cache-write + // tokens are billed at the full input rate (the provider does not charge + // a separate write premium). Anthropic models set this ~1.25× input. + CacheWrite *float64 `yaml:"cache_write"` } `yaml:"models"` } @@ -53,14 +57,19 @@ func ParsePricing(data []byte) (Pricing, error) { models := make(map[string]ModelPrice, len(yp.Models)) for name, m := range yp.Models { - cached := m.Input // default: no cache discount + cached := m.Input // default: no cache-read discount if m.CachedInput != nil { cached = *m.CachedInput } + write := m.Input // default: no separate cache-write premium + if m.CacheWrite != nil { + write = *m.CacheWrite + } models[name] = ModelPrice{ Input: m.Input, Output: m.Output, CachedInput: cached, + CacheWrite: write, } } diff --git a/pricing.yaml b/pricing.yaml index 3061179..3f9168b 100644 --- a/pricing.yaml +++ b/pricing.yaml @@ -10,6 +10,10 @@ # cached_input USD/Mtok for prompt tokens served from the provider's cache # (a discounted subset of input). Omit if the model has no # cache discount — Augur then bills cached tokens at `input`. +# cache_write USD/Mtok for prompt tokens WRITTEN to the provider's cache +# (Anthropic's cache_creation_input_tokens, ~1.25x input). Omit +# for providers that don't bill writes separately (OpenAI, +# Gemini) — Augur then bills them at `input`. version: 1 snapshot_date: "2026-06-21" @@ -31,16 +35,19 @@ models: output: 4.40 cached_input: 0.55 - # --- Anthropic --- + # --- Anthropic --- (cache_write ≈ 1.25× input; cached_input is 5-min read rate) claude-opus-4-8: input: 15.00 output: 75.00 cached_input: 1.50 + cache_write: 18.75 claude-sonnet-4-6: input: 3.00 output: 15.00 cached_input: 0.30 + cache_write: 3.75 claude-haiku-4-5: input: 1.00 output: 5.00 cached_input: 0.10 + cache_write: 1.25 diff --git a/proxy/multiprovider_test.go b/proxy/multiprovider_test.go new file mode 100644 index 0000000..1b4a88a --- /dev/null +++ b/proxy/multiprovider_test.go @@ -0,0 +1,174 @@ +package proxy + +import ( + "bytes" + "net/http/httptest" + "strings" + "testing" + + "augur/trace" +) + +// Anthropic's non-streaming Messages response: input_tokens EXCLUDES the cache +// buckets, so the recorded InputTokens must be their sum. +func TestParseUsageAnthropic(t *testing.T) { + body := `{"type":"message","model":"claude-haiku-4-5","usage":{"input_tokens":100,"output_tokens":50,"cache_read_input_tokens":20,"cache_creation_input_tokens":10}}` + u, has, model := parseUsageJSON([]byte(body)) + if !has { + t.Fatal("hasUsage = false, want true") + } + if model != "claude-haiku-4-5" { + t.Errorf("model = %q, want claude-haiku-4-5", model) + } + // total prompt = 100 + 20 read + 10 write = 130 + want := tokenUsage{InputTokens: 130, OutputTokens: 50, CachedTokens: 20, CacheWriteTokens: 10} + if u != want { + t.Errorf("usage = %+v, want %+v", u, want) + } +} + +// Anthropic streams usage across two events: message_start carries the input and +// cache buckets (with a placeholder output_tokens), message_delta carries the +// real output_tokens. The SSE merge must combine them, not overwrite. +func TestAnthropicStreamingMerge(t *testing.T) { + sse := "event: message_start\n" + + `data: {"type":"message_start","message":{"model":"claude-haiku-4-5","usage":{"input_tokens":100,"cache_read_input_tokens":20,"cache_creation_input_tokens":10,"output_tokens":1}}}` + "\n\n" + + "event: message_delta\n" + + `data: {"type":"message_delta","usage":{"output_tokens":50}}` + "\n\n" + + u, model := usageFromSSE([]byte(sse)) + want := tokenUsage{InputTokens: 130, OutputTokens: 50, CachedTokens: 20, CacheWriteTokens: 10} + if u != want { + t.Errorf("merged usage = %+v, want %+v", u, want) + } + if model != "claude-haiku-4-5" { + t.Errorf("model = %q, want claude-haiku-4-5", model) + } +} + +// Gemini reports usage under usageMetadata with camelCase counts; cached content +// is a subset of the prompt count. +func TestParseUsageGemini(t *testing.T) { + body := `{"usageMetadata":{"promptTokenCount":200,"candidatesTokenCount":80,"cachedContentTokenCount":30,"totalTokenCount":280}}` + u, has, _ := parseUsageJSON([]byte(body)) + if !has { + t.Fatal("hasUsage = false, want true") + } + want := tokenUsage{InputTokens: 200, OutputTokens: 80, CachedTokens: 30} + if u != want { + t.Errorf("usage = %+v, want %+v", u, want) + } +} + +func TestModelFromPath(t *testing.T) { + cases := []struct{ path, want string }{ + {"/v1beta/models/gemini-2.0-flash:generateContent", "gemini-2.0-flash"}, + {"/v1beta/models/gemini-1.5-pro:streamGenerateContent", "gemini-1.5-pro"}, + {"/v1/chat/completions", ""}, // OpenAI names the model in the body + {"/v1/messages", ""}, // Anthropic too + } + for _, c := range cases { + if got := modelFromPath(c.path); got != c.want { + t.Errorf("modelFromPath(%q) = %q, want %q", c.path, got, c.want) + } + } +} + +func TestMaybeInjectIncludeUsage(t *testing.T) { + // A streaming OpenAI request without the option gets it injected. + in := `{"model":"gpt-4o","stream":true,"messages":[]}` + out := maybeInjectIncludeUsage("/v1/chat/completions", []byte(in)) + if !strings.Contains(string(out), `"include_usage":true`) { + t.Errorf("expected include_usage injected, got %s", out) + } + if !strings.Contains(string(out), `"model":"gpt-4o"`) { + t.Errorf("injection dropped other fields: %s", out) + } + + // Non-streaming request is left byte-for-byte unchanged. + ns := `{"model":"gpt-4o","messages":[]}` + if got := maybeInjectIncludeUsage("/v1/chat/completions", []byte(ns)); string(got) != ns { + t.Errorf("non-streaming body mutated: %s", got) + } + + // Already opted in → unchanged. + opted := `{"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true}}` + if got := maybeInjectIncludeUsage("/v1/chat/completions", []byte(opted)); string(got) != opted { + t.Errorf("opted-in body mutated: %s", got) + } + + // Non-OpenAI path (Anthropic) → never injected, even when streaming. + anthropic := `{"model":"claude-haiku-4-5","stream":true}` + if got := maybeInjectIncludeUsage("/v1/messages", []byte(anthropic)); string(got) != anthropic { + t.Errorf("Anthropic body mutated: %s", got) + } + + // Malformed body → returned unchanged (never break the request). + bad := `not json` + if got := maybeInjectIncludeUsage("/v1/chat/completions", []byte(bad)); string(got) != bad { + t.Errorf("malformed body mutated: %s", got) + } +} + +// The proxy classifies a byte-identical repeat within a run as a retry, distinct +// bodies as initial calls, and resets per run. +func TestRetryClassification(t *testing.T) { + up := newFakeUpstream() + defer up.close() + up.respBody = chatResponse + + var buf bytes.Buffer + s := newTestProxy(t, up, &buf) + srv := httptest.NewServer(s) + defer srv.Close() + + same := `{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}` + other := `{"model":"gpt-4o","messages":[{"role":"user","content":"different"}]}` + + doTagged(t, srv.URL, "s", "run-1", same) // initial + doTagged(t, srv.URL, "s", "run-1", same) // retry (identical body) + doTagged(t, srv.URL, "s", "run-1", other) // initial (new body) + doTagged(t, srv.URL, "s", "run-2", same) // initial (new run resets) + + recs, err := trace.ReadAll(&buf) + if err != nil { + t.Fatalf("ReadAll: %v", err) + } + want := []string{trace.KindInitial, trace.KindRetry, trace.KindInitial, trace.KindInitial} + if len(recs) != len(want) { + t.Fatalf("got %d rows, want %d", len(recs), len(want)) + } + for i, w := range want { + if recs[i].Kind != w { + t.Errorf("row %d Kind = %q, want %q", i, recs[i].Kind, w) + } + } +} + +// End-to-end: an Anthropic-shaped response flows through the proxy and lands in +// the trace with the cache buckets split correctly. +func TestProxyRecordsAnthropicUsage(t *testing.T) { + up := newFakeUpstream() + defer up.close() + up.respBody = `{"type":"message","model":"claude-haiku-4-5","usage":{"input_tokens":100,"output_tokens":50,"cache_read_input_tokens":20,"cache_creation_input_tokens":10}}` + + var buf bytes.Buffer + s := newTestProxy(t, up, &buf) + srv := httptest.NewServer(s) + defer srv.Close() + + doTagged(t, srv.URL, "s", "run-1", `{"model":"claude-haiku-4-5","messages":[]}`) + + recs, err := trace.ReadAll(&buf) + if err != nil || len(recs) != 1 { + t.Fatalf("trace: err=%v rows=%d", err, len(recs)) + } + r := recs[0] + if r.Model != "claude-haiku-4-5" { + t.Errorf("model = %q", r.Model) + } + if r.InputTokens != 130 || r.OutputTokens != 50 || r.CachedTokens != 20 || r.CacheWriteTokens != 10 { + t.Errorf("tokens = in %d out %d cached %d write %d, want 130/50/20/10", + r.InputTokens, r.OutputTokens, r.CachedTokens, r.CacheWriteTokens) + } +} diff --git a/proxy/proxy.go b/proxy/proxy.go index 5425b3f..03546b1 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -17,6 +17,7 @@ import ( "bytes" "encoding/json" "fmt" + "hash/fnv" "io" "net/http" "net/url" @@ -60,26 +61,36 @@ type Server struct { client *http.Client now nowFunc + // InjectUsage, when true (the default from New), rewrites streaming OpenAI + // chat-completion requests to set stream_options.include_usage=true so the + // provider emits an exact usage block even when the agent didn't ask for it. + // Set false to forward the request byte-for-byte. + InjectUsage bool + mode Mode cassette *cassette.Cassette - mu sync.Mutex - seq map[string]int // per (scenario|run) next call ordinal + mu sync.Mutex + seq map[string]int // per (scenario|run) next call ordinal + seen map[string]map[uint64]bool // per (scenario|run) request-body hashes seen (retry detection) } // New returns a Server that forwards to upstream (e.g. https://api.openai.com) // and appends trace rows via tracer. A nil client uses a sensible default. The -// server starts in ModeLive; call Record or Replay to change mode. +// server starts in ModeLive with usage injection on; call Record or Replay to +// change mode. func New(upstream *url.URL, tracer *trace.Writer, client *http.Client) *Server { if client == nil { client = &http.Client{Timeout: 10 * time.Minute} } return &Server{ - upstream: upstream, - tracer: tracer, - client: client, - now: time.Now, - seq: make(map[string]int), + upstream: upstream, + tracer: tracer, + client: client, + now: time.Now, + InjectUsage: true, + seq: make(map[string]int), + seen: make(map[string]map[uint64]bool), } } @@ -108,6 +119,32 @@ func (s *Server) nextSeq(scenario, run string) int { return n } +// classify labels a call by whether its request body has been seen before in +// the same (scenario, run): a repeat is the observable signature of a client- +// library retry, a first sighting is an initial call. This is the honest +// classification the proxy can make — it deliberately does not try to separate +// fan-out from sequential tool-loop steps (both are initial calls with distinct +// bodies), because call concurrency is not visible at the HTTP layer. +func (s *Server) classify(scenario, run string, body []byte) string { + h := fnv.New64a() + _, _ = h.Write(body) + sum := h.Sum64() + + key := scenario + "|" + run + s.mu.Lock() + defer s.mu.Unlock() + seen := s.seen[key] + if seen == nil { + seen = make(map[uint64]bool) + s.seen[key] = seen + } + if seen[sum] { + return trace.KindRetry + } + seen[sum] = true + return trace.KindInitial +} + func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { scenario := r.Header.Get(HeaderScenarioID) run := r.Header.Get(HeaderRunID) @@ -122,17 +159,31 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { } _ = r.Body.Close() reqModel := modelFromRequest(reqBody) + if reqModel == "" { + // Gemini names the model in the URL path, not the body. + reqModel = modelFromPath(r.URL.Path) + } // Assign the call's ordinal once, up front: it is both the trace's seq and // the cassette key, so record and replay must compute it identically. seq := s.nextSeq(scenario, run) + // Classify from the ORIGINAL request body (before any usage injection), so a + // retry — a byte-identical repeat — is recognised regardless of injection, + // and record/replay classify identically. + kind := s.classify(scenario, run, reqBody) if s.mode == ModeReplay { - s.replay(w, scenario, run, seq, reqModel, r.URL.Path) + s.replay(w, scenario, run, seq, reqModel, r.URL.Path, kind) return } - outReq, err := s.buildUpstreamRequest(r, reqBody) + // Forward the (possibly usage-injected) body upstream; classification and + // the cassette are unaffected because both key off the original request. + fwdBody := reqBody + if s.InjectUsage { + fwdBody = maybeInjectIncludeUsage(r.URL.Path, reqBody) + } + outReq, err := s.buildUpstreamRequest(r, fwdBody) if err != nil { http.Error(w, "augur proxy: building upstream request: "+err.Error(), http.StatusBadGateway) return @@ -152,7 +203,7 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { if s.mode == ModeLive && isEventStream(resp.Header) { usage, respModel := s.streamResponse(w, resp) latency := s.now().Sub(start) - s.recordTrace(start, scenario, run, seq, pickModel(reqModel, respModel), usage, latency.Milliseconds(), r.URL.Path, resp.StatusCode) + s.recordTrace(start, scenario, run, seq, pickModel(reqModel, respModel), usage, latency.Milliseconds(), r.URL.Path, resp.StatusCode, kind) return } @@ -179,14 +230,14 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { } usage, respModel := extractUsage(contentType, body) - s.recordTrace(start, scenario, run, seq, pickModel(reqModel, respModel), usage, latency.Milliseconds(), r.URL.Path, resp.StatusCode) + s.recordTrace(start, scenario, run, seq, pickModel(reqModel, respModel), usage, latency.Milliseconds(), r.URL.Path, resp.StatusCode, kind) } // replay serves a previously recorded response from the cassette without // contacting the provider, and regenerates the call's trace row from it. A miss // means the agent made a call that was not recorded — surfaced as a 502 so the // divergence is loud rather than silently mis-costed. -func (s *Server) replay(w http.ResponseWriter, scenario, run string, seq int, reqModel, path string) { +func (s *Server) replay(w http.ResponseWriter, scenario, run string, seq int, reqModel, path, kind string) { e, ok := s.cassette.Lookup(scenario, run, seq) if !ok { http.Error(w, fmt.Sprintf("augur proxy: replay miss for scenario %q run %q seq %d (agent diverged from the recording?)", @@ -201,24 +252,26 @@ func (s *Server) replay(w http.ResponseWriter, scenario, run string, seq int, re _, _ = w.Write(body) usage, respModel := extractUsage(e.ContentType, body) - s.recordTrace(s.now(), scenario, run, seq, pickModel(reqModel, respModel), usage, e.LatencyMs, path, e.Status) + s.recordTrace(s.now(), scenario, run, seq, pickModel(reqModel, respModel), usage, e.LatencyMs, path, e.Status, kind) } // recordTrace writes one trace row. A write failure must not corrupt the // agent's response but must be loud: a dropped row means an under-counted bill. -func (s *Server) recordTrace(ts time.Time, scenario, run string, seq int, model string, u oaiUsage, latencyMs int64, path string, status int) { +func (s *Server) recordTrace(ts time.Time, scenario, run string, seq int, model string, u tokenUsage, latencyMs int64, path string, status int, kind string) { rec := trace.Record{ - Timestamp: ts.UTC().Format(time.RFC3339Nano), - ScenarioID: scenario, - RunID: run, - Seq: seq, - Model: model, - InputTokens: u.InputTokens, - OutputTokens: u.OutputTokens, - CachedTokens: u.CachedTokens, - LatencyMs: latencyMs, - Endpoint: path, - Status: status, + Timestamp: ts.UTC().Format(time.RFC3339Nano), + ScenarioID: scenario, + RunID: run, + Seq: seq, + Model: model, + InputTokens: u.InputTokens, + OutputTokens: u.OutputTokens, + CachedTokens: u.CachedTokens, + CacheWriteTokens: u.CacheWriteTokens, + LatencyMs: latencyMs, + Endpoint: path, + Status: status, + Kind: kind, } if err := s.tracer.Write(rec); err != nil { fmt.Printf("augur proxy: WARNING trace write failed: %v\n", err) @@ -259,52 +312,141 @@ func (s *Server) buildUpstreamRequest(r *http.Request, body []byte) (*http.Reque return req, nil } -// oaiUsage is the token accounting block of an OpenAI-compatible response. -type oaiUsage struct { - InputTokens int - OutputTokens int - CachedTokens int +// tokenUsage is the provider-neutral token accounting Augur records for one LLM +// call. CachedTokens and CacheWriteTokens are subsets of InputTokens (see +// package cost). parseUsageJSON maps the OpenAI, Anthropic, and Gemini wire +// shapes onto it. +type tokenUsage struct { + InputTokens int + OutputTokens int + CachedTokens int + CacheWriteTokens int } -// parseUsageJSON extracts token usage and the resolved model from one -// OpenAI-compatible JSON object — a full non-streaming response body or a single -// streamed chunk's payload. usage is a pointer in the wire shape so we can tell -// "no usage block" (every streamed chunk before the last) from "usage with zero -// tokens": hasUsage is false in the former. Both the buffered and streaming -// paths funnel through here so their accounting cannot drift apart. -func parseUsageJSON(data []byte) (u oaiUsage, hasUsage bool, model string) { +// merge folds a newly parsed chunk's usage into u, taking each field's latest +// non-zero value. Streaming responses spread usage across chunks differently by +// provider — OpenAI emits one final block, Anthropic splits input (message_start) +// from output (message_delta), Gemini repeats a cumulative block — and +// last-non-zero-wins reconciles all three without provider-specific stream state. +func (u *tokenUsage) merge(n tokenUsage) { + if n.InputTokens != 0 { + u.InputTokens = n.InputTokens + } + if n.OutputTokens != 0 { + u.OutputTokens = n.OutputTokens + } + if n.CachedTokens != 0 { + u.CachedTokens = n.CachedTokens + } + if n.CacheWriteTokens != 0 { + u.CacheWriteTokens = n.CacheWriteTokens + } +} + +// wireUsage is the union of the OpenAI and Anthropic per-call usage blocks (both +// carried under a "usage" key). Pointer fields distinguish "key absent" from a +// genuine zero, so parseUsageJSON can tell which provider's shape it holds. +type wireUsage struct { + // OpenAI + PromptTokens *int `json:"prompt_tokens"` + CompletionTokens *int `json:"completion_tokens"` + PromptTokensDetails *struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + // Anthropic + InputTokens *int `json:"input_tokens"` + OutputTokens *int `json:"output_tokens"` + CacheReadInputTokens int `json:"cache_read_input_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens"` +} + +// parseUsageJSON extracts token usage and the resolved model from one OpenAI-, +// Anthropic-, or Gemini-compatible JSON object: a full non-streaming response +// body or a single streamed chunk's payload. hasUsage is false when the object +// carries no usage block at all (e.g. an OpenAI content chunk), letting callers +// tell "no usage here" from "usage that is genuinely zero". Every path funnels +// through here so the providers cannot drift apart in accounting. +func parseUsageJSON(data []byte) (u tokenUsage, hasUsage bool, model string) { var parsed struct { - Model string `json:"model"` - Usage *struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - PromptTokensDetails struct { - CachedTokens int `json:"cached_tokens"` - } `json:"prompt_tokens_details"` - } `json:"usage"` + Model string `json:"model"` + Usage *wireUsage `json:"usage"` + Message *struct { // Anthropic streaming message_start nests usage + model + Model string `json:"model"` + Usage *wireUsage `json:"usage"` + } `json:"message"` + UsageMetadata *struct { // Gemini + PromptTokenCount int `json:"promptTokenCount"` + CandidatesTokenCount int `json:"candidatesTokenCount"` + CachedContentTokenCount int `json:"cachedContentTokenCount"` + } `json:"usageMetadata"` } if err := json.Unmarshal(data, &parsed); err != nil { - return oaiUsage{}, false, "" + return tokenUsage{}, false, "" } - if parsed.Usage == nil { - return oaiUsage{}, false, parsed.Model + + model = parsed.Model + if parsed.Message != nil && parsed.Message.Model != "" { + model = parsed.Message.Model + } + + // Gemini: usageMetadata with camelCase counts (cached is a subset of prompt). + if g := parsed.UsageMetadata; g != nil { + return tokenUsage{ + InputTokens: g.PromptTokenCount, + OutputTokens: g.CandidatesTokenCount, + CachedTokens: g.CachedContentTokenCount, + }, true, model + } + + // OpenAI / Anthropic: a "usage" block, possibly nested under "message". + w := parsed.Usage + if w == nil && parsed.Message != nil { + w = parsed.Message.Usage + } + if w == nil { + return tokenUsage{}, false, model + } + + switch { + case w.PromptTokens != nil || w.CompletionTokens != nil || w.PromptTokensDetails != nil: + // OpenAI: cached_tokens is a subset of prompt_tokens (already the total). + u = tokenUsage{InputTokens: deref(w.PromptTokens), OutputTokens: deref(w.CompletionTokens)} + if w.PromptTokensDetails != nil { + u.CachedTokens = w.PromptTokensDetails.CachedTokens + } + return u, true, model + case w.InputTokens != nil || w.OutputTokens != nil || w.CacheReadInputTokens != 0 || w.CacheCreationInputTokens != 0: + // Anthropic: input_tokens EXCLUDES the cache buckets, so the total prompt + // is the sum. A message_delta carries only output_tokens (input nil) — + // the stream merge folds it into the message_start usage. + return tokenUsage{ + InputTokens: deref(w.InputTokens) + w.CacheReadInputTokens + w.CacheCreationInputTokens, + OutputTokens: deref(w.OutputTokens), + CachedTokens: w.CacheReadInputTokens, + CacheWriteTokens: w.CacheCreationInputTokens, + }, true, model + default: + return tokenUsage{}, false, model } - return oaiUsage{ - InputTokens: parsed.Usage.PromptTokens, - OutputTokens: parsed.Usage.CompletionTokens, - CachedTokens: parsed.Usage.PromptTokensDetails.CachedTokens, - }, true, parsed.Model +} + +func deref(p *int) int { + if p == nil { + return 0 + } + return *p } // usageFromResponse is a thin wrapper over parseUsageJSON for the non-streaming // path and tests. -func usageFromResponse(body []byte) (oaiUsage, string) { +func usageFromResponse(body []byte) (tokenUsage, string) { u, _, model := parseUsageJSON(body) return u, model } -// modelFromRequest reads the "model" field from an OpenAI-compatible request -// body. Returns "" if absent or unparseable. +// modelFromRequest reads the "model" field from an OpenAI- or Anthropic- +// compatible request body. Returns "" if absent or unparseable (Gemini names the +// model in the URL path instead — see modelFromPath). func modelFromRequest(body []byte) string { var parsed struct { Model string `json:"model"` @@ -315,6 +457,78 @@ func modelFromRequest(body []byte) string { return parsed.Model } +// modelFromPath extracts the model from a Gemini-style request path such as +// /v1beta/models/gemini-2.0-flash:generateContent. Returns "" for paths that do +// not carry a "/models/" segment (OpenAI/Anthropic name it in the body). +func modelFromPath(path string) string { + const marker = "/models/" + i := strings.Index(path, marker) + if i < 0 { + return "" + } + rest := path[i+len(marker):] + // Strip the ":method" suffix (generateContent, streamGenerateContent, …). + if c := strings.IndexByte(rest, ':'); c >= 0 { + rest = rest[:c] + } + // And any trailing path segment. + if sl := strings.IndexByte(rest, '/'); sl >= 0 { + rest = rest[:sl] + } + return rest +} + +// maybeInjectIncludeUsage returns body with stream_options.include_usage set to +// true when it is an OpenAI chat-completions streaming request that hasn't opted +// in — so the provider emits an exact usage block Augur can record instead of a +// zero-token row. It is a no-op (returns body unchanged) for non-matching paths, +// non-streaming requests, requests that already set the option, and bodies that +// don't parse — the request must never break because of injection. Restricting +// to the OpenAI chat path keeps the extra field away from Anthropic/Gemini +// (which report usage without opting in and would reject an unknown parameter). +func maybeInjectIncludeUsage(path string, body []byte) []byte { + if !strings.Contains(path, "/chat/completions") { + return body + } + var probe struct { + Stream *bool `json:"stream"` + StreamOptions *struct { + IncludeUsage *bool `json:"include_usage"` + } `json:"stream_options"` + } + if err := json.Unmarshal(body, &probe); err != nil { + return body + } + if probe.Stream == nil || !*probe.Stream { + return body // not a streaming request + } + if probe.StreamOptions != nil && probe.StreamOptions.IncludeUsage != nil && *probe.StreamOptions.IncludeUsage { + return body // already opted in + } + + // Merge into a generic map so every other field the agent set is preserved + // (key order is not, which is immaterial to the provider). + var m map[string]json.RawMessage + if err := json.Unmarshal(body, &m); err != nil { + return body + } + opts := map[string]json.RawMessage{} + if raw, ok := m["stream_options"]; ok { + _ = json.Unmarshal(raw, &opts) // best-effort; overwrite include_usage below + } + opts["include_usage"] = json.RawMessage("true") + optsRaw, err := json.Marshal(opts) + if err != nil { + return body + } + m["stream_options"] = optsRaw + out, err := json.Marshal(m) + if err != nil { + return body + } + return out +} + // hopByHopHeaders are connection-specific headers that must not be forwarded by // a proxy (RFC 7230 §6.1). var hopByHopHeaders = []string{ diff --git a/proxy/proxy_stream.go b/proxy/proxy_stream.go index d986492..c77f6ad 100644 --- a/proxy/proxy_stream.go +++ b/proxy/proxy_stream.go @@ -24,18 +24,19 @@ func isEventStreamCT(ct string) bool { // response body, handling both a single JSON object and a full SSE stream. The // buffered and replay paths use it so a recorded streaming response is costed // the same as it was live. -func extractUsage(contentType string, body []byte) (oaiUsage, string) { +func extractUsage(contentType string, body []byte) (tokenUsage, string) { if isEventStreamCT(contentType) { return usageFromSSE(body) } return usageFromResponse(body) } -// usageFromSSE scans a buffered SSE body line-by-line for the usage block and -// the resolved model, mirroring what streamResponse captures while relaying. -func usageFromSSE(body []byte) (oaiUsage, string) { +// usageFromSSE scans a buffered SSE body line-by-line, merging every chunk's +// usage (see tokenUsage.merge) and tracking the resolved model, mirroring what +// streamResponse captures while relaying. +func usageFromSSE(body []byte) (tokenUsage, string) { var ( - usage oaiUsage + usage tokenUsage model string ) for line := range bytes.SplitSeq(body, []byte("\n")) { @@ -44,7 +45,7 @@ func usageFromSSE(body []byte) (oaiUsage, string) { model = m } if has { - usage = u + usage.merge(u) } } } @@ -59,7 +60,7 @@ func usageFromSSE(body []byte) (oaiUsage, string) { // // Relaying is byte-exact: each line is read with its delimiter intact and // written straight back, so the client sees precisely the provider's stream. -func (s *Server) streamResponse(w http.ResponseWriter, resp *http.Response) (oaiUsage, string) { +func (s *Server) streamResponse(w http.ResponseWriter, resp *http.Response) (tokenUsage, string) { copyHeader(w.Header(), resp.Header) w.WriteHeader(resp.StatusCode) @@ -67,7 +68,7 @@ func (s *Server) streamResponse(w http.ResponseWriter, resp *http.Response) (oai _ = rc.Flush() // flush headers so the client opens the stream immediately var ( - usage oaiUsage + usage tokenUsage model string ) br := bufio.NewReader(resp.Body) @@ -83,7 +84,7 @@ func (s *Server) streamResponse(w http.ResponseWriter, resp *http.Response) (oai model = m } if has { - usage = u + usage.merge(u) } } } @@ -107,14 +108,14 @@ var doneMarker = []byte("[DONE]") // parseSSEChunk pulls usage/model out of a single SSE line. Non-data lines, // comments, the [DONE] marker, and chunks without a usage block return // hasUsage=false (model may still be set). -func parseSSEChunk(line []byte) (u oaiUsage, hasUsage bool, model string) { +func parseSSEChunk(line []byte) (u tokenUsage, hasUsage bool, model string) { t := bytes.TrimSpace(line) if !bytes.HasPrefix(t, dataPrefix) { - return oaiUsage{}, false, "" + return tokenUsage{}, false, "" } payload := bytes.TrimSpace(t[len(dataPrefix):]) if len(payload) == 0 || bytes.Equal(payload, doneMarker) { - return oaiUsage{}, false, "" + return tokenUsage{}, false, "" } return parseUsageJSON(payload) } diff --git a/proxy/proxy_test.go b/proxy/proxy_test.go index 8240b00..620df88 100644 --- a/proxy/proxy_test.go +++ b/proxy/proxy_test.go @@ -132,6 +132,7 @@ func TestProxyRecordsUsage(t *testing.T) { LatencyMs: 0, // fixed clock → zero elapsed Endpoint: "/v1/chat/completions", Status: 200, + Kind: trace.KindInitial, } if got != want { t.Errorf("trace row mismatch:\n got %+v\nwant %+v", got, want) @@ -267,7 +268,7 @@ func TestUsageFromResponse(t *testing.T) { // Garbage / no usage → zero, no panic. z, m := usageFromResponse([]byte(`not json`)) - if z != (oaiUsage{}) || m != "" { + if z != (tokenUsage{}) || m != "" { t.Errorf("garbage body = (%+v, %q), want zero", z, m) } } diff --git a/proxy_cmd.go b/proxy_cmd.go index 8d33d09..5d29c67 100644 --- a/proxy_cmd.go +++ b/proxy_cmd.go @@ -16,8 +16,9 @@ import ( func runProxy(args []string) error { fs := flag.NewFlagSet("proxy", flag.ContinueOnError) listen := fs.String("listen", ":8080", "address to listen on") - upstream := fs.String("upstream", "https://api.openai.com", "base URL of the real OpenAI-compatible provider") + upstream := fs.String("upstream", "https://api.openai.com", "base URL of the real provider (OpenAI-, Anthropic-, or Gemini-compatible)") tracePath := fs.String("trace", "trace.jsonl", "path to append the cost trace to (JSONL)") + injectUsage := fs.Bool("inject-usage", true, "auto-set stream_options.include_usage on OpenAI streaming requests so usage is captured exactly") if err := fs.Parse(args); err != nil { return err } @@ -37,6 +38,7 @@ func runProxy(args []string) error { defer tracer.Close() srv := proxy.New(up, tracer, nil) + srv.InjectUsage = *injectUsage fmt.Printf("augur proxy: listening on %s → forwarding to %s, tracing to %s\n", *listen, up.String(), *tracePath) diff --git a/run_cmd.go b/run_cmd.go index 100fbe8..849609e 100644 --- a/run_cmd.go +++ b/run_cmd.go @@ -36,6 +36,7 @@ func runRun(args []string) error { continueOnError := fs.Bool("continue-on-error", false, "keep going after an agent invocation fails") record := fs.String("record", "", "record every response to this cassette file (real provider calls)") replay := fs.String("replay", "", "replay responses from this cassette file (no provider calls, no tokens)") + injectUsage := fs.Bool("inject-usage", true, "auto-set stream_options.include_usage on OpenAI streaming requests so usage is captured exactly") if err := fs.Parse(args); err != nil { return err } @@ -66,6 +67,7 @@ func runRun(args []string) error { defer tracer.Close() pxy := proxy.New(up, tracer, nil) + pxy.InjectUsage = *injectUsage cass, err := configureCassette(pxy, *record, *replay) if err != nil { return err diff --git a/trace/trace.go b/trace/trace.go index 7222270..dbdd0fa 100644 --- a/trace/trace.go +++ b/trace/trace.go @@ -44,11 +44,15 @@ type Record struct { // model named in the request (it matches pricing.yaml keys), falling back // to the model echoed in the response. Model string `json:"model"` - // Token accounting, mirroring the provider's report. CachedTokens is the - // cached SUBSET of InputTokens, not an additional bucket (see package cost). - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` - CachedTokens int `json:"cached_tokens"` + // Token accounting, mirroring the provider's report. CachedTokens and + // CacheWriteTokens are cached-read and cache-write SUBSETS of InputTokens, + // not additional buckets (see package cost). CacheWriteTokens is Anthropic's + // cache_creation_input_tokens, billed at a premium over the base input rate; + // OpenAI/Gemini leave it zero (their cache writes are not separately billed). + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CachedTokens int `json:"cached_tokens"` + CacheWriteTokens int `json:"cache_write_tokens,omitempty"` // LatencyMs is the wall-clock time the upstream provider took to respond. LatencyMs int64 `json:"latency_ms"` // Endpoint is the request path (e.g. /v1/chat/completions), kept so a trace @@ -59,8 +63,26 @@ type Record struct { // hunts for, and a row with zero tokens still records that an attempt was // made. Status int `json:"status,omitempty"` + // Kind classifies the call within its run using only what the proxy can + // observe truthfully: KindRetry marks a call whose request body is byte- + // identical to an earlier call in the same (scenario, run) — the shape a + // client library's retry takes — versus KindInitial for the first time a + // given request is seen. It is empty on traces written before this field + // existed. Fan-out vs sequential tool-loop steps are NOT distinguished here + // (both are KindInitial): the proxy cannot see call concurrency. + Kind string `json:"kind,omitempty"` } +// Call-kind values for Record.Kind. See the Kind field for the classification +// rule and its deliberate limits. +const ( + // KindInitial is the first time a given request body is seen in a run. + KindInitial = "initial" + // KindRetry is a call whose request body repeats an earlier one in the same + // run — the observable signature of a client-library retry. + KindRetry = "retry" +) + // Writer appends Records to an io.Writer as JSON Lines. It is safe for // concurrent use: the proxy serves requests in parallel, so Write is guarded by // a mutex to keep lines from interleaving.