proxy: native Anthropic/Gemini capture, usage auto-injection, retry classification

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

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 {

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
 		}

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
 }
 

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)
+	}
+}

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,
 		}
 	}