From 1d23c363521815aa1a4680fd7849114dd28e53cc Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Tue, 1 Sep 2026 16:27:46 -0400 Subject: [PATCH 01/28] docs: Propose directional body capabilities and tool-prune plugin Specifies two changes. First, a directional split of the body-write capability: PluginCapabilities.WritesBody becomes WritesRequestBody and gains a WritesResponseBody sibling, so response streaming is gated on the response-side flag alone. A request-only mutator currently disables incremental SSE relay for a body it never touches, because Pipeline.WritesBody() is an undirected OR that both proxy listeners consult. Second, tool-prune: an outbound plugin that deletes named entries from the tools array of an inference request. One registration, a static remove list, no persistence. Measure-only mode comes free from the framework's per-plugin on_error: observe policy. Includes a compatibility audit of all three plugins that declare the capability today (context-guru writes requests only and regains streaming; sparc and cpex write both and are unchanged), plus two adjacent fixes: cloneCatalog silently drops new capability fields, and SetBody's godoc describes an enforcement the code does not implement. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- docs/proposals/tool-prune.md | 405 +++++++++++++++++++++++++++++++++++ 1 file changed, 405 insertions(+) create mode 100644 docs/proposals/tool-prune.md diff --git a/docs/proposals/tool-prune.md b/docs/proposals/tool-prune.md new file mode 100644 index 000000000..78ad96957 --- /dev/null +++ b/docs/proposals/tool-prune.md @@ -0,0 +1,405 @@ +# Directional Body Capabilities and the `tool-prune` Plugin + +**Status**: Draft +**Date**: September 2026 + +This document specifies two changes that together let AuthBridge cut an agent's +token bill by removing tool definitions the agent never calls: + +1. **A directional split of the body-write capability.** `PluginCapabilities.WritesBody` + is renamed to `WritesRequestBody` and joined by `WritesResponseBody`. Response + streaming is then gated on the response-side flag alone, so a plugin that only + rewrites requests no longer forfeits incremental server-sent events (SSE). +2. **`tool-prune`**, an outbound plugin that deletes named entries from the `tools` + array of an inference request. The list is static, produced at setup time by a + new `abctl tools scan` subcommand that analyses local Claude Code transcripts. + +Part 1 is a prerequisite for part 2 but stands on its own merits: it is a +framework correctness fix that any future request-only mutator benefits from. + +## Motivation + +Claude Code sends the full tool manifest on every request. On the author's machine +that manifest is roughly 20,500 tokens, and it sits at the front of the cached +prompt prefix, so it is billed on every turn of every session. Measurements across +seven developers show most of those definitions are never called even once in a +30-day window. + +Removing the dead entries is a pure win: fewer prompt tokens, no behaviour change +for tools the agent actually uses. Doing it in the proxy rather than in each +client's configuration means it works for every agent behind AuthBridge without +per-client setup, and it is measurable centrally. + +### Why the list is static + +The `tools` array is at the front of the cached prompt prefix. Any change to it +invalidates every cache breakpoint after it, at 1.25x write cost. A plugin that +learned at runtime and revised its verdict would repeatedly bust the prompt cache +and could plausibly destroy more value than it saves. A list fixed at setup time +busts the cache exactly once, then stabilises. + +This is the deciding argument for setup-time analysis over runtime learning, and +it removes the need for any persistence layer inside the plugin. + +## Part 1: Directional body capabilities + +### The defect + +`PluginCapabilities.WritesBody` is a single boolean covering both directions. +`Pipeline.WritesBody()` (`authlib/pipeline/pipeline.go:383-390`) is a plain OR +across the chain, with no notion of direction: + +```go +func (p *Pipeline) WritesBody() bool { + for _, plugin := range p.plugins { + if plugin.Capabilities().Normalize().WritesBody { + return true + } + } + return false +} +``` + +Both proxy listeners consult that predicate to decide whether an SSE response may +be relayed incrementally (`forwardproxy/server.go:383-387`, +`reverseproxy/server.go:440-454`). A plugin that rewrites only the **request** body +therefore disables **response** streaming, for a body it never touches: + +```go +if isEventStream(resp.Header.Get("Content-Type")) && resp.Body != nil { + if s.OutboundPipeline.WritesBody() { + slog.Warn("forward-proxy: text/event-stream response with WritesBody plugin — falling back to buffered path", ...) +``` + +The cost is latency and feel, not correctness: the buffered path restores the body +verbatim (`reverseproxy/server.go:466`) and the `event:` line is preserved on the +re-framing path (`reverseproxy/server.go:749-754`). But the +user-visible effect is that a long completion arrives in one lump after a silent +wait instead of appearing incrementally, which is the first thing anyone notices. + +The fallback is also independent of `ErrorPolicy`: `WritesBody()` asks whether any +plugin *declares* the capability, not whether it is currently *permitted* to +mutate. So a plugin running in `on_error: observe` (measure-only) loses response +streaming before it has gained anything. + +Request buffering is not part of this cost. Requests are never streamed — Claude +Code sends one complete `POST /v1/messages` with a `Content-Length` — and the +request body is already read end to end before dispatch whenever any plugin +declares `ReadsBody` (`forwardproxy/server.go:256-266`, `pipeline.go:368-376`). +`inference-parser` declares it, so every request through the demo chain is already +fully buffered. A request-only mutator adds no buffering at all. + +### The change + +```go +type PluginCapabilities struct { + ReadsBody bool + + // WritesRequestBody declares the plugin may call pctx.SetBody. + WritesRequestBody bool + + // WritesResponseBody declares the plugin may call pctx.SetResponseBody. + // Listeners fall back from incremental SSE relay to the buffered path + // only when some plugin declares this. + WritesResponseBody bool + + Requires []string + RequiresAny []string + Description string +} +``` + +`Normalize()` promotes `ReadsBody` from either write flag. +`Pipeline.WritesResponseBody()` becomes the streaming predicate; +`Pipeline.WritesRequestBody()` keeps gating request propagation. + +### Why rename rather than add + +Adding `WritesResponseBody` alongside an unchanged `WritesBody` would default the new +field to `false`, so an out-of-tree plugin that rewrites responses would silently +start streaming and then call `SetResponseBody` after bytes had already been sent. +That converts a latency annoyance into a correctness bug, in code the change does +not touch. + +Plugins are compiled into the binary — registration is via +`plugins.RegisterPlugin`, with no dynamic loading — so renaming the field gives +every such author a **compile error** instead. That is the safest available +failure: impossible to miss, trivially fixed, and it forces the author to answer +"which body?" rather than inherit a default they never considered. It also retires +the ambiguity permanently, since after the change there is no undirected option to +pick. + +The rename is confined to Go source and documentation. `PluginCapabilities` has no +struct tags and is never marshalled directly; the session API defines its own +tagged wire types (`sessionapi.CatalogEntry` at `sessionapi/server.go:55-63`, and +the pipeline view whose `readsBody` field is at `:167`) and neither exposes +`writesBody`. So **no wire key +and no configuration key changes.** Capabilities are not configurable, so +`authlib/config` is untouched. + +### Compatibility audit + +Every in-tree plugin that declares the capability today, and what it actually does: + +| Plugin | Rewrites request | Rewrites response | Evidence | After the change | +|---|---|---|---|---| +| `context-guru` | yes | **no** | `contextguru/plugin.go:160`; no `SetResponseBody` call anywhere | `WritesRequestBody` — **gains** response streaming | +| `sparc` | yes | yes | `sparc/plugin.go:214`; `sparc/respond.go:111,122` | both flags — unchanged | +| `cpex` | yes | yes | `cpex/plugin.go:122`; `cmf_body.go:609`, `cmf_a2a.go:216`, `cmf_inference.go:218` | both flags — unchanged | +| `tool-prune` | yes | no | new | `WritesRequestBody` — streams | + +Two of the three genuinely need the buffered path and keep it. Only `context-guru` +changes behaviour, and only by regaining streaming it never needed to lose. + +`validateCapabilities` (`pipeline.go:549-573`) becomes direction-aware, but the +**outcome is identical for every configuration that exists today**: all three +current plugins write requests, so they remain mutually exclusive exactly as +before. The reader-ordering rule stays triggered by either write flag, so no +configuration that validates today starts failing and none that fails starts +passing. + +### One adjacent fix + +`cloneCatalog` (`plugins/registry.go:202-222`) copies capability fields one at a +time, so any field added to `PluginCapabilities` is silently dropped from +`/v1/plugins`: + +```go +Capabilities: pipeline.PluginCapabilities{ + ReadsBody: caps.ReadsBody, + WritesBody: caps.WritesBody, + Description: caps.Description, + Requires: append([]string(nil), caps.Requires...), + RequiresAny: append([]string(nil), caps.RequiresAny...), +}, +``` + +Replace the field-by-field construction with a struct copy plus explicit slice +reallocation, which preserves the deep-copy guarantee and picks up future fields +automatically: + +```go +c := caps +c.Requires = append([]string(nil), caps.Requires...) +c.RequiresAny = append([]string(nil), caps.RequiresAny...) +``` + +### Documented contract fix + +`SetBody`'s godoc (`pipeline/context.go:390-396`) states that a plugin without the +write capability which calls `SetBody` mutates only the in-memory context and +leaves the wire unchanged. The code does not do this: `SetBody` sets +`c.bodyMutated = true` unconditionally outside observe mode (`context.go:424`), and +the listeners gate purely on `pctx.BodyMutated()` (`forwardproxy/server.go:335`, +`reverseproxy/server.go:358`). An undeclared mutation therefore does reach the +wire. + +This proposal does not add the missing enforcement — doing so silently would break +any plugin currently relying on the actual behaviour. It corrects the comment to +describe what the code does, and notes the divergence so a future change can close +it deliberately. Left as documented, it is a live trap: it makes "just don't +declare the capability" look like a legitimate way to keep streaming. + +## Part 2: The `tool-prune` plugin + +### Behaviour + +One registration, `plugins.RegisterPlugin("tool-prune", ...)`: + +```go +pipeline.PluginCapabilities{ + WritesRequestBody: true, + RequiresAny: []string{"inference-parser"}, + Description: "Removes unused tool definitions from inference requests", +} +``` + +On each outbound request: + +1. Skip unless the path matches `paths` (default `/v1/chat/completions`, + `/v1/completions`, `/v1/messages`), matched by suffix as `context-guru` does. +2. Read the parsed manifest from `pctx.Extensions.Inference.Tools`. +3. For each configured name present in the manifest, delete its element from the + `tools` array of the **original** request bytes with `sjson.DeleteBytes`, + iterating indices in descending order so earlier deletions do not shift later + ones. +4. Call `pctx.SetBody` once with the result. + +Every byte outside the deleted array elements is unchanged. `gjson`/`sjson` are +already in `authlib/go.mod` (currently indirect), so no new dependency. + +Any error or panic fails open: the original body is forwarded unmodified. A +cost optimisation must never be able to break a request. + +### Configuration + +```yaml +pipeline: + outbound: + plugins: + - inference-parser + - mcp-parser + - a2a-parser + - name: tool-prune + on_error: observe # measure only; switch to enforce when trusted + config: + remove: [NotebookEdit, ScheduleWakeup, TaskOutput] +``` + +`remove` is the complete verdict. There is no learning, no state, and no storage +dependency. + +### Measure-only mode comes from the framework + +`on_error` is a per-plugin policy already parsed by `authlib/config` +(`config.go:257`, values `enforce | observe | off`). Under `observe`, `SetBody` is +a no-op on bytes but still records a modify `Invocation` with `Shadow=true` +(`context.go:397-421`), so "would have removed" is countable without changing a +single request. `off` skips dispatch entirely. + +This is why one registration suffices: the same plugin code serves measure and +enforce, selected by one word of configuration. `context.go` states the intent +directly — "Plugin code therefore looks identical under enforce and observe." + +Off-by-default is satisfied structurally: the plugin is absent from the shipped +pipeline list until a user adds it, and the documented first step adds it with +`on_error: observe`. + +### Where the list comes from: `abctl tools scan` + +A new subcommand ports the discovery core of `claude-tool-audit.py` (about 40 of +its 814 lines) into Go: + +- Read `~/.claude/projects/**/*.jsonl`. +- Hot-path line filter on the literal `"tool_use"` before any JSON parsing. +- Deduplicate tool calls by the unique `tool_use` block id. +- Window to the last `--days` (default 30). + +`abctl` currently has no subcommand dispatch — `main.go` parses two flags and +launches the terminal UI. The change checks for a non-flag first argument before +`flag.Parse()` and dispatches, falling through to the UI otherwise. + +``` +abctl tools scan [--days 30] [--keep Name,Name] [--write ] +``` + +Without `--write` it prints the YAML block. With `--write` it patches the +`remove:` list of the `tool-prune` entry in place, idempotently. + +### The offered-set problem, and how the scan stays safe + +Transcripts record tools that were **called**, never tools that were **offered**. +This is structural, not a defect: a configured-but-never-invoked tool leaves no +trace. Two consequences: + +- Tools never called in the window but bundled in the known Claude Code tool set + are the removal candidates, and they are where most of the 20,500 tokens sit. +- A tool name the scan has never heard of is **kept**. Removing a tool the model + needs is the harmful direction of failure; carrying a few extra definitions is + not. + +The bundled set is version-sensitive: developers on newer Claude Code releases +produced tool calls the current table does not recognise. Two mitigations: + +1. Unknown names are always kept, so drift costs savings, never correctness. +2. At startup the plugin compares its configured `remove` list against the names + it observes in `ext.Tools` and logs any configured name that never appears, so + a stale list surfaces as a warning rather than a silent no-op. + +A `--keep` flag and a small "implies" table cover tools whose use is indirect — +for example `Agent` implying `SendMessage`, which a transcript may not show being +called directly. + +### Installation flow + +`install-demo.sh` already downloads both binaries with checksum verification and +prints next steps. `authbridge-proxy` writes `cortex-ca/demo.yaml` on first run +(`cmd/authbridge-proxy/demo.go`), and that file is hot-reloaded — its own header +says so — so the list can be filled in without a restart. + +- `demoConfigYAML()` gains the `tool-prune` entry with `on_error: observe` and an + empty `remove: []`. +- `install-demo.sh` runs `abctl tools scan --write` when the config already + exists, and otherwise prints the block in its next-steps output alongside the + existing "Watch traffic" hint. + +### What the user sees in-session + +Claude Code's `/context` breakdown has `System tools`, `Tool schemas`, `MCP tools`, +`Custom agents`, `Memory files` and `Free space` line items — confirmed by string +inspection of the installed 2.1.257 binary. **It will not show this saving.** It is +a client-side pre-flight breakdown of what the CLI assembled; it necessarily +computes `Free space` itself, and the pruning happens downstream. This is the +first place a user would look, and it must be documented as unaffected. + +What does move is `/cost` and any figure derived from the API response `usage` +block: the server bills the request it received, so `input_tokens` and +`cache_read_input_tokens` genuinely drop. + +The honest limit: proxy-side pruning saves money but does **not** return context +window to the user. The client still believes it sent the full manifest, so +auto-compact triggers at the same point. Recovering headroom requires client-side +configuration (`--allowedTools`, disabling unused MCP servers). AuthBridge's +advantage is the complement — it applies to every agent behind it with no +per-client change, and it measures. + +## Delivery + +Three commits, sequenced so the regression argument survives review. + +1. **Mechanical rename.** `WritesBody` to `WritesRequestBody` across 107 + references in 28 files (Go and documentation), with no semantic change. + Reviewable as a single token substitution, and every existing test passing + still carries meaning because nothing but the name moved. +2. **The split.** Add `WritesResponseBody`; declare it on `sparc` and `cpex`; + point both listener branches at `Pipeline.WritesResponseBody()`; convert + `cloneCatalog` to a struct copy; correct the `SetBody` godoc; add tests. +3. **`tool-prune`.** Plugin, `abctl tools scan`, `demoConfigYAML()` entry, + `install-demo.sh` wiring, and documentation. + +### Testing + +For part 1, the primary regression argument is that the existing body-capability +tests pass with only the identifier renamed. On top of that: + +- A truth table for `Pipeline.WritesResponseBody()` across the four plugin shapes + (request-only, response-only, both, neither). +- Listener tests: an SSE response with a request-only writer in the chain relays + incrementally; with a `sparc`-shaped or `cpex`-shaped chain it still buffers. +- A reflection-based `cloneCatalog` round-trip that fails if any future + capability field is dropped. +- `validateCapabilities` table assertions covering the current plugin + combinations, to show acceptance and rejection are unchanged. + +For part 2: + +- Byte-level assertions that pruning a manifest leaves every other byte of the + request identical, including key order and whitespace. +- Descending-index deletion verified against a manifest where a naive ascending + loop would delete the wrong elements. +- Names absent from the manifest are ignored without error. +- Malformed and truncated bodies fail open, forwarding the original bytes. +- Under `on_error: observe`, the body is unchanged and a `Shadow=true` invocation + is recorded. +- Scanner tests over fixture transcripts: window boundaries, `tool_use` block + deduplication, unknown names retained, `--keep` honoured. + +### Risks + +| Risk | Mitigation | +|---|---| +| Removing a tool the agent needs | Unknown names always kept; `--keep` override; ship with `on_error: observe`; fail open on any error | +| Stale bundled tool set as Claude Code evolves | Drift reduces savings only; plugin warns on configured names never observed in `ext.Tools` | +| One-off prompt-cache invalidation when the list changes | Inherent and bounded: static list means it happens once, then the prefix is stable | +| Commit 1 conflicts with in-flight branches declaring `WritesBody` | One-line fix per branch; the compile error makes it self-evident | +| `context-guru` regaining response streaming exposes a latent bug in that path | Covered by the listener tests above; the path is already exercised by chains with no body writer | + +## Open questions + +None blocking. Two items deliberately deferred: + +- Adding the missing enforcement so `SetBody` matches its documented contract. + Needs its own compatibility review. +- Surfacing per-plugin savings in `abctl`. The in-session `/cost` readout covers + the first cut; a `Metrics:` section in the existing plugin detail pane is a + natural follow-up if operators want it centrally. From 39326de88c149c444eab5022d919dbcbfc600f2a Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Tue, 1 Sep 2026 16:38:06 -0400 Subject: [PATCH 02/28] docs: Add plugin metrics channel to the tool-prune proposal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Claude Code's /cost reports a session total with no baseline, so it cannot attribute a saving to the plugin. Adds a third part to the proposal: a generic Metric / MetricsProvider optional interface in authlib/pipeline, surfaced through describePipeline and the existing abctl plugin detail pane. Both extension points already have the pattern needed — the wire side mirrors the RawConfigProvider assertion, the pane mirrors the Config section. Counters are in-memory and per-process, so no storage dependency. The plugin distinguishes enforce from observe by checking pctx.BodyMutated() after SetBody, which makes observe mode a projection: it reports what it would save before any request changes. Bytes removed are exact. Tokens are estimated by calibrating a bytes-per-token ratio on the user's own traffic via response usage, rather than bundling a tokenizer or hardcoding a constant. This re-adds OnFinish for two counter reads only; the removal list stays entirely configuration-driven. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- docs/proposals/tool-prune.md | 169 +++++++++++++++++++++++++++++++++-- 1 file changed, 161 insertions(+), 8 deletions(-) diff --git a/docs/proposals/tool-prune.md b/docs/proposals/tool-prune.md index 78ad96957..f5f80d284 100644 --- a/docs/proposals/tool-prune.md +++ b/docs/proposals/tool-prune.md @@ -3,8 +3,9 @@ **Status**: Draft **Date**: September 2026 -This document specifies two changes that together let AuthBridge cut an agent's -token bill by removing tool definitions the agent never calls: +This document specifies three changes that together let AuthBridge cut an agent's +token bill by removing tool definitions the agent never calls, and show the +operator what that saved: 1. **A directional split of the body-write capability.** `PluginCapabilities.WritesBody` is renamed to `WritesRequestBody` and joined by `WritesResponseBody`. Response @@ -13,9 +14,13 @@ token bill by removing tool definitions the agent never calls: 2. **`tool-prune`**, an outbound plugin that deletes named entries from the `tools` array of an inference request. The list is static, produced at setup time by a new `abctl tools scan` subcommand that analyses local Claude Code transcripts. +3. **A plugin metrics channel**, surfaced in the existing `abctl` plugin detail + pane. Claude Code's `/cost` reports a session total, which is too coarse to + attribute a saving to the plugin, so the plugin reports its own counters. Part 1 is a prerequisite for part 2 but stands on its own merits: it is a framework correctness fix that any future request-only mutator benefits from. +Part 3 is likewise generic — any plugin gains a display channel. ## Motivation @@ -336,6 +341,10 @@ What does move is `/cost` and any figure derived from the API response `usage` block: the server bills the request it received, so `input_tokens` and `cache_read_input_tokens` genuinely drop. +But `/cost` reports a **session total**, with no baseline to compare against. A +user cannot tell from one aggregate number how much of it the plugin saved, or +whether enabling the plugin was worth it. That is what part 3 is for. + The honest limit: proxy-side pruning saves money but does **not** return context window to the user. The client still believes it sent the full manifest, so auto-compact triggers at the same point. Recovering headroom requires client-side @@ -343,9 +352,135 @@ configuration (`--allowedTools`, disabling unused MCP servers). AuthBridge's advantage is the complement — it applies to every agent behind it with no per-client change, and it measures. +## Part 3: Plugin metrics in `abctl` + +### What the plugin counts + +Counters are in-memory and per-process, guarded by a mutex. No persistence, no +storage backend, no new dependency. + +```go +type metrics struct { + mu sync.Mutex + + requestsSeen uint64 // matched the path gate + requestsPruned uint64 // body actually rewritten (enforce) + requestsProjected uint64 // would have been rewritten (observe) + + toolsRemoved uint64 + perTool map[string]uint64 + + bytesRemoved uint64 // sum of deleted tools array elements + + promptTokens uint64 // from response usage, via OnFinish + requestBytes uint64 // body size of the same requests + requestsWithUsage uint64 +} +``` + +The plugin distinguishes enforce from observe without inspecting policy: under +`ErrorPolicyObserve`, `SetBody` leaves `bodyMutated` false +(`pipeline/context.go:418-421`), so checking `pctx.BodyMutated()` after the call +tells the plugin which counter to increment. This makes **observe mode a +projection**: the plugin computes exactly what it would remove and reports the +saving before a single request changes. Read the projection, then flip to +`enforce`. + +### Turning bytes into tokens without guessing + +The plugin knows removed bytes exactly, but billing is denominated in tokens. +Rather than bundle a tokenizer or hardcode a bytes-per-token constant, the ratio +is calibrated on the user's own traffic: `OnFinish` reads +`pctx.Extensions.Inference.PromptTokens` — populated by `inference-parser` from +the response `usage` block — alongside the request body size for that same +request. Estimated tokens saved is then `bytesRemoved x (promptTokens / +requestBytes)`, reported with its sample size and labelled an estimate. + +`OnFinish` is the correct hook for response-derived data: `inference-parser` is a +`StreamingResponder`, and `RunResponse` skips `OnResponse` for such plugins. + +This re-adds `OnFinish`, which the static-list decision had removed. The scope is +deliberately narrow — two counter reads, no persistence and no influence on the +removal list, which stays entirely configuration-driven. + +One approximation to state plainly: under `enforce`, `PromptTokens` is already the +post-pruning count, so the ratio is measured on pruned requests. That is +acceptable for a bytes-to-tokens conversion factor, which is a property of the +tokenizer and content mix rather than of the pruning, but it is why the figure is +labelled an estimate rather than a measurement. + +### A generic metrics interface + +Added to `authlib/pipeline/plugin.go` beside the existing optional interfaces: + +```go +// Metric is one operator-facing counter reported by a plugin. +type Metric struct { + Name string `json:"name"` + Value float64 `json:"value"` + Unit string `json:"unit,omitempty"` // count | bytes | tokens | ratio + Note string `json:"note,omitempty"` // e.g. "estimate, n=1284" +} + +// MetricsProvider is implemented by plugins that expose counters for +// operator display. Called on demand from the session API; must be safe +// for concurrent use and must not block. +type MetricsProvider interface { + Metrics() []Metric +} +``` + +`plugins.StatsSource` and `auth.Stats` already exist but are the wrong vehicle: +`auth.Stats` is auth-specific, with typed approval and denial enums and a custom +`MarshalJSON` (`auth/auth.go:61,183`). Carrying "bytes removed" through it would +distort its meaning. A separate plugin-defined interface keeps auth statistics +auth-shaped and gives every future plugin a display channel. + +### Wire and UI + +Both extension points already have the exact pattern needed. + +- `sessionapi`: add `Metrics []pipeline.Metric` to the pipeline plugin view and + populate it in `describePipeline` with a three-line type assertion mirroring the + `RawConfigProvider` case at `sessionapi/server.go:234`. Matching field on + `apiclient.PipelinePlugin`. +- `cmd/abctl/tui/plugin_detail_pane.go`: a `Metrics:` section after the dependency + sections and before `Config:`, following the always-newline convention that the + comment at `:67-70` records as deliberate — it exists to stop layout jitter when + navigating between plugins that do and do not have the section. `Note` renders + in `styleHint`, as `Description` already does at `:27`. + +Roughly 20 lines in the pane, 3 in `describePipeline`, 2 struct fields, and about +35 lines for the interface plus the plugin's counters. + +The operator reads something like: + +``` +Metrics: + requests seen 1284 count + requests pruned 1284 count + tools removed 11556 count + bytes removed 9389184 bytes + bytes removed / request 7312 bytes + tokens saved / request ~1830 tokens estimate, n=1284 +``` + +In observe mode the same rows appear with `requests projected` in place of +`requests pruned`, so the distinction between a projection and a realised saving +is visible in the readout rather than inferred from configuration. + +### Limits + +Counters are per-process and in-memory: they reset when the proxy restarts and are +not aggregated across a fleet. That is the right trade for the laptop scenario this +targets, and it is what keeps the plugin free of a storage dependency. Fleet-wide +aggregation belongs on the existing stats server +(`runtimeutil.StartStatServer`, port 47602 in the demo config), which is a +natural later addition and does not change the plugin. + ## Delivery -Three commits, sequenced so the regression argument survives review. +Four commits, sequenced so the regression argument survives review. 1. **Mechanical rename.** `WritesBody` to `WritesRequestBody` across 107 references in 28 files (Go and documentation), with no semantic change. @@ -354,8 +489,12 @@ Three commits, sequenced so the regression argument survives review. 2. **The split.** Add `WritesResponseBody`; declare it on `sparc` and `cpex`; point both listener branches at `Pipeline.WritesResponseBody()`; convert `cloneCatalog` to a struct copy; correct the `SetBody` godoc; add tests. -3. **`tool-prune`.** Plugin, `abctl tools scan`, `demoConfigYAML()` entry, - `install-demo.sh` wiring, and documentation. +3. **The metrics channel.** `Metric` and `MetricsProvider` in `authlib/pipeline`; + the `describePipeline` type assertion and wire field; the `abctl` pane section. + Lands before the plugin so the plugin arrives already visible, and so this + generic addition is reviewed on its own merits rather than as plugin scaffolding. +4. **`tool-prune`.** Plugin and its counters, `abctl tools scan`, + `demoConfigYAML()` entry, `install-demo.sh` wiring, and documentation. ### Testing @@ -384,6 +523,18 @@ For part 2: - Scanner tests over fixture transcripts: window boundaries, `tool_use` block deduplication, unknown names retained, `--keep` honoured. +For part 3: + +- A plugin that does not implement `MetricsProvider` produces no `metrics` key on + the wire and renders `(none)` without disturbing pane layout. +- Concurrent `Metrics()` calls against a live counter update, under the race + detector. +- Enforce mode increments `requestsPruned`; observe mode increments + `requestsProjected` and leaves `bytesRemoved` accumulating, so the projection is + non-zero while the body is untouched. +- The bytes-to-tokens ratio is reported as zero-valued rather than dividing by + zero when no response usage has been seen yet. + ### Risks | Risk | Mitigation | @@ -393,6 +544,8 @@ For part 2: | One-off prompt-cache invalidation when the list changes | Inherent and bounded: static list means it happens once, then the prefix is stable | | Commit 1 conflicts with in-flight branches declaring `WritesBody` | One-line fix per branch; the compile error makes it self-evident | | `context-guru` regaining response streaming exposes a latent bug in that path | Covered by the listener tests above; the path is already exercised by chains with no body writer | +| The estimated token saving is read as a measurement | Unit and sample size shown on the row; the underlying byte counts are exact and reported separately, so the estimate is never the only number | +| Counters reset on restart and mislead someone comparing across restarts | Documented; `requests seen` is displayed alongside every derived figure so the sample behind it is always visible | ## Open questions @@ -400,6 +553,6 @@ None blocking. Two items deliberately deferred: - Adding the missing enforcement so `SetBody` matches its documented contract. Needs its own compatibility review. -- Surfacing per-plugin savings in `abctl`. The in-session `/cost` readout covers - the first cut; a `Metrics:` section in the existing plugin detail pane is a - natural follow-up if operators want it centrally. +- Fleet-wide metric aggregation on the stats server, and a Prometheus exposition + of `MetricsProvider`. The per-process counters in part 3 cover the laptop case + this targets; neither addition changes the plugin. From 6ca3fbbbfa0462825770af7fd9a9346a8819e744 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 11:19:24 -0400 Subject: [PATCH 03/28] refactor(authbridge): Rename WritesBody to WritesRequestBody Mechanical rename of PluginCapabilities.WritesBody and Pipeline.WritesBody() to WritesRequestBody, across 107 references in 28 files (Go source, documentation, and one YAML comment). No semantic change: every call site keeps the behaviour it had, so the existing body-capability tests still pass with only the identifier moved. This prepares the directional split in the next commit, where WritesResponseBody becomes the SSE streaming predicate and a request-only mutator stops forfeiting incremental relay. Renaming rather than adding gives out-of-tree plugin authors a compile error instead of a silently-defaulted field. Also gofmt's a pre-existing mis-sorted import in contextguru/plugin.go, which the rename's struct-alignment reflow pulled in. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- authbridge/authlib/listener/extproc/server.go | 2 +- .../extproc/server_contentlength_test.go | 2 +- .../authlib/listener/extproc/server_test.go | 6 ++-- .../forwardproxy/mcp_sse_stream_test.go | 2 +- .../authlib/listener/forwardproxy/server.go | 14 ++++---- .../forwardproxy/server_headerdiff_test.go | 2 +- .../listener/forwardproxy/server_test.go | 6 ++-- .../listener/forwardproxy/streaming_test.go | 14 ++++---- .../authlib/listener/reverseproxy/server.go | 10 +++--- .../listener/reverseproxy/server_test.go | 8 ++--- .../listener/reverseproxy/streaming_test.go | 4 +-- .../authlib/pipeline/bodymutation_test.go | 34 +++++++++---------- authbridge/authlib/pipeline/context.go | 4 +-- authbridge/authlib/pipeline/holder.go | 4 +-- authbridge/authlib/pipeline/pipeline.go | 22 ++++++------ authbridge/authlib/pipeline/pipeline_test.go | 2 +- authbridge/authlib/pipeline/plugin.go | 12 +++---- .../authlib/plugins/contextguru/build_test.go | 2 +- .../authlib/plugins/contextguru/plugin.go | 12 +++---- authbridge/authlib/plugins/cpex/plugin.go | 12 +++---- .../authlib/plugins/cpex/plugin_test.go | 4 +-- authbridge/authlib/plugins/registry.go | 10 +++--- authbridge/authlib/plugins/sparc/plugin.go | 8 ++--- .../authlib/plugins/sparc/plugin_test.go | 4 +-- authbridge/demos/context-guru/README.md | 4 +-- .../context-guru/k8s/authbridge-config.yaml | 2 +- authbridge/docs/cpex-plugin.md | 4 +-- authbridge/docs/framework-architecture.md | 20 +++++------ authbridge/docs/plugin-reference.md | 14 ++++---- authbridge/docs/plugin-tutorial.md | 8 ++--- 30 files changed, 126 insertions(+), 126 deletions(-) diff --git a/authbridge/authlib/listener/extproc/server.go b/authbridge/authlib/listener/extproc/server.go index 8de88d384..dcb8c1348 100644 --- a/authbridge/authlib/listener/extproc/server.go +++ b/authbridge/authlib/listener/extproc/server.go @@ -657,7 +657,7 @@ func (s *Server) handleResponseBody(ctx context.Context, body []byte, pctx *pipe s.recordOutboundResponseSession(pctx) } - // A plugin that declared WritesBody: true and called pctx.SetResponseBody + // A plugin that declared WritesRequestBody: true and called pctx.SetResponseBody // flips the ResponseBodyMutated flag. Emit the replacement bytes via // BodyMutation so Envoy rewrites the downstream response; otherwise // pass through with no mutation. The flag avoids the O(n) string diff --git a/authbridge/authlib/listener/extproc/server_contentlength_test.go b/authbridge/authlib/listener/extproc/server_contentlength_test.go index 5100c4b13..17401ab0f 100644 --- a/authbridge/authlib/listener/extproc/server_contentlength_test.go +++ b/authbridge/authlib/listener/extproc/server_contentlength_test.go @@ -60,7 +60,7 @@ type responseMutator struct{ newBody []byte } func (*responseMutator) Name() string { return "response-mutator" } func (*responseMutator) Capabilities() pipeline.PluginCapabilities { - return pipeline.PluginCapabilities{WritesBody: true} + return pipeline.PluginCapabilities{WritesRequestBody: true} } func (*responseMutator) OnRequest(context.Context, *pipeline.Context) pipeline.Action { return pipeline.Action{Type: pipeline.Continue} diff --git a/authbridge/authlib/listener/extproc/server_test.go b/authbridge/authlib/listener/extproc/server_test.go index 1554d22ab..dca5a02ea 100644 --- a/authbridge/authlib/listener/extproc/server_test.go +++ b/authbridge/authlib/listener/extproc/server_test.go @@ -398,7 +398,7 @@ func (p *bodyRecorderPlugin) OnResponse(_ context.Context, _ *pipeline.Context) return pipeline.Action{Type: pipeline.Continue} } -// bodyMutatorPlugin declares WritesBody and rewrites pctx.Body via +// bodyMutatorPlugin declares WritesRequestBody and rewrites pctx.Body via // SetBody. Used to assert extproc emits a BodyMutation on the wire // when a plugin rewrites the request body. type bodyMutatorPlugin struct { @@ -407,7 +407,7 @@ type bodyMutatorPlugin struct { func (p *bodyMutatorPlugin) Name() string { return "body-mutator" } func (p *bodyMutatorPlugin) Capabilities() pipeline.PluginCapabilities { - return pipeline.PluginCapabilities{WritesBody: true} + return pipeline.PluginCapabilities{WritesRequestBody: true} } func (p *bodyMutatorPlugin) OnRequest(_ context.Context, pctx *pipeline.Context) pipeline.Action { pctx.SetBody(p.newBody) @@ -417,7 +417,7 @@ func (p *bodyMutatorPlugin) OnResponse(_ context.Context, _ *pipeline.Context) p return pipeline.Action{Type: pipeline.Continue} } -// TestExtProc_RequestBodyMutation_Inbound: a WritesBody plugin must +// TestExtProc_RequestBodyMutation_Inbound: a WritesRequestBody plugin must // produce a RequestBody ProcessingResponse carrying BodyMutation with // the new bytes, and the header mutation must request content-encoding // be removed. diff --git a/authbridge/authlib/listener/forwardproxy/mcp_sse_stream_test.go b/authbridge/authlib/listener/forwardproxy/mcp_sse_stream_test.go index 6b4f286b9..e8aa9d9d3 100644 --- a/authbridge/authlib/listener/forwardproxy/mcp_sse_stream_test.go +++ b/authbridge/authlib/listener/forwardproxy/mcp_sse_stream_test.go @@ -59,7 +59,7 @@ func TestForwardProxy_SSE_StreamsWithoutResponder(t *testing.T) { store := session.New(5*time.Minute, 100, 0) defer store.Close() - // Empty pipeline: HasStreamingResponders()==false and WritesBody()==false, + // Empty pipeline: HasStreamingResponders()==false and WritesRequestBody()==false, // so serveOutbound routes to streamPassthrough — the reporter's plain-proxy // shape (their only outbound plugin, token-exchange, is likewise not a // StreamingResponder). diff --git a/authbridge/authlib/listener/forwardproxy/server.go b/authbridge/authlib/listener/forwardproxy/server.go index c46421ea7..b4fce2d4b 100644 --- a/authbridge/authlib/listener/forwardproxy/server.go +++ b/authbridge/authlib/listener/forwardproxy/server.go @@ -352,7 +352,7 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, isBridge r.Header[k] = append([]string(nil), vv...) // set / overwrite } - // If a WritesBody plugin rewrote pctx.Body, ship the new bytes + // If a WritesRequestBody plugin rewrote pctx.Body, ship the new bytes // upstream and clear Content-Encoding (see forwardproxy response // path for the rationale). if pctx.BodyMutated() { @@ -400,14 +400,14 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, isBridge // response (the client Accepts both), so the same tool may return // JSON on one call and SSE on the next. Decide here rather than // negotiating, and don't take the streaming path when a plugin - // declares WritesBody (mutating a body we've already started + // declares WritesRequestBody (mutating a body we've already started // forwarding is incompatible with streaming) — fall back to // buffered with a warning log instead. if isEventStream(resp.Header.Get("Content-Type")) && resp.Body != nil { - if s.OutboundPipeline.WritesBody() { + if s.OutboundPipeline.WritesRequestBody() { // A body mutator needs the whole body to rewrite it, so it // can't stream — fall back to the buffered path with a warning. - slog.Warn("forward-proxy: text/event-stream response with WritesBody plugin — falling back to buffered path", "host", r.Host) + slog.Warn("forward-proxy: text/event-stream response with WritesRequestBody plugin — falling back to buffered path", "host", r.Host) } else if s.OutboundPipeline.HasStreamingResponders() { // Streaming-aware plugins (inference-parser, a2a-parser) parse // each SSE frame; handleStreamingResponse re-frames via sseframe. @@ -419,16 +419,16 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, isBridge // would drop the event:/id:/retry: lines that generic SSE // clients (e.g. an MCP Streamable HTTP client) depend on. Fixes #642. // - // A plugin that declares ReadsBody (but not WritesBody, and is + // A plugin that declares ReadsBody (but not WritesRequestBody, and is // not a StreamingResponder) also lands here, and its OnResponse // runs against an empty pctx.ResponseBody: streamPassthrough // forwards the stream without buffering it. We deliberately don't // buffer to satisfy such a plugin — that would reintroduce the // #642 timeout on a live stream. A plugin that must inspect a // streamed body should implement StreamingResponder. Warn - // (mirroring the WritesBody fallback above) so the + // (mirroring the WritesRequestBody fallback above) so the // misconfiguration surfaces instead of the plugin silently seeing - // no body. WritesBody is already false in this branch, so + // no body. WritesRequestBody is already false in this branch, so // NeedsBody() here implies ReadsBody. if s.OutboundPipeline.NeedsBody() { slog.Warn("forward-proxy: text/event-stream response with a ReadsBody plugin that is not a StreamingResponder — streaming byte-for-byte; its OnResponse will see an empty body (implement StreamingResponder to inspect a streamed body)", "host", r.Host) diff --git a/authbridge/authlib/listener/forwardproxy/server_headerdiff_test.go b/authbridge/authlib/listener/forwardproxy/server_headerdiff_test.go index d9b61924d..23b9ab473 100644 --- a/authbridge/authlib/listener/forwardproxy/server_headerdiff_test.go +++ b/authbridge/authlib/listener/forwardproxy/server_headerdiff_test.go @@ -17,7 +17,7 @@ import ( // point of PR #760 is that EVERY plugin header mutation — not just the // old Authorization special case — must reach the upstream request. The // plugin declares no capabilities: a header write does not need -// ReadsBody/WritesBody, mirroring how staticinject/cpex mutate headers. +// ReadsBody/WritesRequestBody, mirroring how staticinject/cpex mutate headers. type headerMutatorPlugin struct { set map[string]string // header -> value to Set (set or overwrite) del []string // headers to Del diff --git a/authbridge/authlib/listener/forwardproxy/server_test.go b/authbridge/authlib/listener/forwardproxy/server_test.go index 1ff66a55a..0bab01d72 100644 --- a/authbridge/authlib/listener/forwardproxy/server_test.go +++ b/authbridge/authlib/listener/forwardproxy/server_test.go @@ -323,7 +323,7 @@ func (p *bodyRecorderPlugin) OnResponse(_ context.Context, _ *pipeline.Context) return pipeline.Action{Type: pipeline.Continue} } -// bodyMutatorPlugin declares WritesBody and rewrites pctx.Body via +// bodyMutatorPlugin declares WritesRequestBody and rewrites pctx.Body via // SetBody. Used below to confirm the forwardproxy propagates the // mutation to the upstream request. type bodyMutatorPlugin struct { @@ -332,7 +332,7 @@ type bodyMutatorPlugin struct { func (p *bodyMutatorPlugin) Name() string { return "body-mutator" } func (p *bodyMutatorPlugin) Capabilities() pipeline.PluginCapabilities { - return pipeline.PluginCapabilities{WritesBody: true} + return pipeline.PluginCapabilities{WritesRequestBody: true} } func (p *bodyMutatorPlugin) OnRequest(_ context.Context, pctx *pipeline.Context) pipeline.Action { pctx.SetBody(p.newBody) @@ -342,7 +342,7 @@ func (p *bodyMutatorPlugin) OnResponse(_ context.Context, _ *pipeline.Context) p return pipeline.Action{Type: pipeline.Continue} } -// TestForwardProxy_RequestBodyMutation: a WritesBody plugin rewriting +// TestForwardProxy_RequestBodyMutation: a WritesRequestBody plugin rewriting // pctx.Body must cause the upstream backend to receive the new bytes // with a correct Content-Length and no Content-Encoding. func TestForwardProxy_RequestBodyMutation(t *testing.T) { diff --git a/authbridge/authlib/listener/forwardproxy/streaming_test.go b/authbridge/authlib/listener/forwardproxy/streaming_test.go index 47d1c276f..c14eaa822 100644 --- a/authbridge/authlib/listener/forwardproxy/streaming_test.go +++ b/authbridge/authlib/listener/forwardproxy/streaming_test.go @@ -32,11 +32,11 @@ type streamingProbe struct { caps pipeline.PluginCapabilities } -func newStreamingProbe(writesBody bool) *streamingProbe { +func newStreamingProbe(writesRequestBody bool) *streamingProbe { return &streamingProbe{ caps: pipeline.PluginCapabilities{ - ReadsBody: true, - WritesBody: writesBody, + ReadsBody: true, + WritesRequestBody: writesRequestBody, }, } } @@ -158,13 +158,13 @@ func TestForwardProxy_Streaming_FramesFlowThrough(t *testing.T) { } } -// TestForwardProxy_Streaming_WritesBodyFallsBackToBuffered asserts the -// safety guard: a pipeline with a WritesBody plugin can't take the +// TestForwardProxy_Streaming_WritesRequestBodyFallsBackToBuffered asserts the +// safety guard: a pipeline with a WritesRequestBody plugin can't take the // streaming path (the plugin can't rewrite a body we've already // started forwarding). The proxy logs a warning and falls back to // buffered, so the response is delivered correctly even though it // loses the streaming property. -func TestForwardProxy_Streaming_WritesBodyFallsBackToBuffered(t *testing.T) { +func TestForwardProxy_Streaming_WritesRequestBodyFallsBackToBuffered(t *testing.T) { upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) @@ -174,7 +174,7 @@ func TestForwardProxy_Streaming_WritesBodyFallsBackToBuffered(t *testing.T) { })) defer upstream.Close() - probe := newStreamingProbe(true) // WritesBody=true → buffered fallback + probe := newStreamingProbe(true) // WritesRequestBody=true → buffered fallback pipe, err := pipeline.New([]pipeline.Plugin{probe}) if err != nil { t.Fatalf("New pipeline: %v", err) diff --git a/authbridge/authlib/listener/reverseproxy/server.go b/authbridge/authlib/listener/reverseproxy/server.go index 15e108b5f..88516a37c 100644 --- a/authbridge/authlib/listener/reverseproxy/server.go +++ b/authbridge/authlib/listener/reverseproxy/server.go @@ -142,7 +142,7 @@ func NewServer(inbound *pipeline.Holder, sessions *session.Store, backendURL str } // Strip the client's Accept-Encoding, but only when a plugin will // actually inspect the response body: a StreamingResponder (SSE - // re-framing) or any ReadsBody/WritesBody plugin (buffered read into + // re-framing) or any ReadsBody/WritesRequestBody plugin (buffered read into // pctx.ResponseBody). Those paths must see plaintext — with no explicit // Accept-Encoding, Go's transport negotiates gzip itself and // transparently decompresses the response (dropping Content-Encoding / @@ -352,7 +352,7 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) { return } - // If a WritesBody plugin rewrote pctx.Body, send the new bytes to + // If a WritesRequestBody plugin rewrote pctx.Body, send the new bytes to // the backend and clear Content-Encoding (same rationale as the // response path — plugin may have decompressed). if pctx.BodyMutated() { @@ -440,14 +440,14 @@ func (s *Server) modifyResponse(resp *http.Response) error { // called on this path — streaming-aware plugins finalize via // OnResponseFrame(last=true). // - // WritesBody is incompatible with streaming (we can't rewrite a + // WritesRequestBody is incompatible with streaming (we can't rewrite a // body we've already started forwarding) — fall back to buffered // with a warning. if isEventStream(resp.Header.Get("Content-Type")) && s.InboundPipeline.HasStreamingResponders() && resp.Body != nil { - if s.InboundPipeline.WritesBody() { - slog.Warn("reverse-proxy: text/event-stream response with WritesBody plugin — falling back to buffered path", "host", pctx.Host) + if s.InboundPipeline.WritesRequestBody() { + slog.Warn("reverse-proxy: text/event-stream response with WritesRequestBody plugin — falling back to buffered path", "host", pctx.Host) } else { s.installStreamingResponseBody(resp, pctx) // Strip Content-Length — the framing reader doesn't know diff --git a/authbridge/authlib/listener/reverseproxy/server_test.go b/authbridge/authlib/listener/reverseproxy/server_test.go index ffe6b4136..f159e21d8 100644 --- a/authbridge/authlib/listener/reverseproxy/server_test.go +++ b/authbridge/authlib/listener/reverseproxy/server_test.go @@ -285,8 +285,8 @@ func (p *bodyRecorderPlugin) OnResponse(_ context.Context, _ *pipeline.Context) return pipeline.Action{Type: pipeline.Continue} } -// bodyMutatorPlugin declares WritesBody and rewrites the request body -// to a fixed payload. The pipeline validator requires WritesBody run +// bodyMutatorPlugin declares WritesRequestBody and rewrites the request body +// to a fixed payload. The pipeline validator requires WritesRequestBody run // after any ReadsBody plugin, which this satisfies by itself (no reader // present when used alone). type bodyMutatorPlugin struct { @@ -295,7 +295,7 @@ type bodyMutatorPlugin struct { func (p *bodyMutatorPlugin) Name() string { return "body-mutator" } func (p *bodyMutatorPlugin) Capabilities() pipeline.PluginCapabilities { - return pipeline.PluginCapabilities{WritesBody: true} + return pipeline.PluginCapabilities{WritesRequestBody: true} } func (p *bodyMutatorPlugin) OnRequest(_ context.Context, pctx *pipeline.Context) pipeline.Action { pctx.SetBody(p.newBody) @@ -305,7 +305,7 @@ func (p *bodyMutatorPlugin) OnResponse(_ context.Context, _ *pipeline.Context) p return pipeline.Action{Type: pipeline.Continue} } -// TestReverseProxy_RequestBodyMutation: a WritesBody plugin that +// TestReverseProxy_RequestBodyMutation: a WritesRequestBody plugin that // rewrites pctx.Body via SetBody must cause the upstream backend to // receive the new bytes with a correct Content-Length header. Confirms // that the reverseproxy request-path propagation is wired to the diff --git a/authbridge/authlib/listener/reverseproxy/streaming_test.go b/authbridge/authlib/listener/reverseproxy/streaming_test.go index 1a1f02fc0..1725c94aa 100644 --- a/authbridge/authlib/listener/reverseproxy/streaming_test.go +++ b/authbridge/authlib/listener/reverseproxy/streaming_test.go @@ -29,9 +29,9 @@ type streamingProbe struct { caps pipeline.PluginCapabilities } -func newStreamingProbe(writesBody bool) *streamingProbe { +func newStreamingProbe(writesRequestBody bool) *streamingProbe { return &streamingProbe{ - caps: pipeline.PluginCapabilities{ReadsBody: true, WritesBody: writesBody}, + caps: pipeline.PluginCapabilities{ReadsBody: true, WritesRequestBody: writesRequestBody}, } } diff --git a/authbridge/authlib/pipeline/bodymutation_test.go b/authbridge/authlib/pipeline/bodymutation_test.go index 0ea97a00f..54414c04f 100644 --- a/authbridge/authlib/pipeline/bodymutation_test.go +++ b/authbridge/authlib/pipeline/bodymutation_test.go @@ -7,7 +7,7 @@ import ( "testing" ) -// TestCapabilities_Normalize: WritesBody auto-promotes to ReadsBody so a +// TestCapabilities_Normalize: WritesRequestBody auto-promotes to ReadsBody so a // mutator always satisfies the "must have read" invariant. func TestCapabilities_Normalize(t *testing.T) { tests := []struct { @@ -17,8 +17,8 @@ func TestCapabilities_Normalize(t *testing.T) { wantWrites bool }{ { - name: "WritesBody implies ReadsBody", - in: PluginCapabilities{WritesBody: true}, + name: "WritesRequestBody implies ReadsBody", + in: PluginCapabilities{WritesRequestBody: true}, wantReads: true, wantWrites: true, }, @@ -40,41 +40,41 @@ func TestCapabilities_Normalize(t *testing.T) { if got.ReadsBody != tc.wantReads { t.Errorf("ReadsBody = %v, want %v", got.ReadsBody, tc.wantReads) } - if got.WritesBody != tc.wantWrites { - t.Errorf("WritesBody = %v, want %v", got.WritesBody, tc.wantWrites) + if got.WritesRequestBody != tc.wantWrites { + t.Errorf("WritesRequestBody = %v, want %v", got.WritesRequestBody, tc.wantWrites) } }) } } -// TestPipeline_NeedsBody_IncludesWritesBody: NeedsBody returns true even +// TestPipeline_NeedsBody_IncludesWritesRequestBody: NeedsBody returns true even // if the only body-touching plugin is a pure mutator. Listeners rely on // this to turn on buffering before the mutator sees (and rewrites) the // body. -func TestPipeline_NeedsBody_IncludesWritesBody(t *testing.T) { +func TestPipeline_NeedsBody_IncludesWritesRequestBody(t *testing.T) { p := mustBuild(t, &stubPlugin{ name: "mutator", - caps: PluginCapabilities{WritesBody: true}, + caps: PluginCapabilities{WritesRequestBody: true}, }) if !p.NeedsBody() { - t.Error("NeedsBody should be true when any plugin declares WritesBody") + t.Error("NeedsBody should be true when any plugin declares WritesRequestBody") } - if !p.WritesBody() { - t.Error("WritesBody should be true") + if !p.WritesRequestBody() { + t.Error("WritesRequestBody should be true") } } -// TestNew_RejectsTwoMutators: two WritesBody plugins in one pipeline +// TestNew_RejectsTwoMutators: two WritesRequestBody plugins in one pipeline // have ambiguous mutation ordering; Pipeline.New rejects the build and // the error names both plugins so an operator reading pod logs can // identify which two to reconcile. func TestNew_RejectsTwoMutators(t *testing.T) { _, err := New([]Plugin{ - &stubPlugin{name: "redactor-a", caps: PluginCapabilities{WritesBody: true}}, - &stubPlugin{name: "redactor-b", caps: PluginCapabilities{WritesBody: true}}, + &stubPlugin{name: "redactor-a", caps: PluginCapabilities{WritesRequestBody: true}}, + &stubPlugin{name: "redactor-b", caps: PluginCapabilities{WritesRequestBody: true}}, }) if err == nil { - t.Fatal("expected error for two WritesBody plugins") + t.Fatal("expected error for two WritesRequestBody plugins") } if !strings.Contains(err.Error(), "redactor-a") || !strings.Contains(err.Error(), "redactor-b") { t.Errorf("error should name both plugins, got %q", err.Error()) @@ -87,7 +87,7 @@ func TestNew_RejectsTwoMutators(t *testing.T) { // reader mutated content. func TestNew_RejectsReaderAfterMutator(t *testing.T) { _, err := New([]Plugin{ - &stubPlugin{name: "rewriter", caps: PluginCapabilities{WritesBody: true}}, + &stubPlugin{name: "rewriter", caps: PluginCapabilities{WritesRequestBody: true}}, &stubPlugin{name: "parser", caps: PluginCapabilities{ReadsBody: true}}, }) if err == nil { @@ -104,7 +104,7 @@ func TestNew_RejectsReaderAfterMutator(t *testing.T) { func TestNew_AcceptsReaderBeforeMutator(t *testing.T) { _, err := New([]Plugin{ &stubPlugin{name: "parser", caps: PluginCapabilities{ReadsBody: true}}, - &stubPlugin{name: "rewriter", caps: PluginCapabilities{WritesBody: true}}, + &stubPlugin{name: "rewriter", caps: PluginCapabilities{WritesRequestBody: true}}, }) if err != nil { t.Fatalf("reader-before-mutator should be valid, got %v", err) diff --git a/authbridge/authlib/pipeline/context.go b/authbridge/authlib/pipeline/context.go index 789c9ec18..dd882fca8 100644 --- a/authbridge/authlib/pipeline/context.go +++ b/authbridge/authlib/pipeline/context.go @@ -388,9 +388,9 @@ func (c *Context) DenyAndRecord(reason, code, message string) Action { } // SetBody replaces the request body with newBody. Only meaningful when -// the plugin declares WritesBody: true in its Capabilities — the +// the plugin declares WritesRequestBody: true in its Capabilities — the // listener consults pctx.BodyMutated() after Run to decide whether to -// emit the new bytes on the wire. Plugins without WritesBody that call +// emit the new bytes on the wire. Plugins without WritesRequestBody that call // SetBody mutate the in-memory Context (readers downstream see the // change), but the wire is unchanged. // diff --git a/authbridge/authlib/pipeline/holder.go b/authbridge/authlib/pipeline/holder.go index b5f8e018c..fb9f881c9 100644 --- a/authbridge/authlib/pipeline/holder.go +++ b/authbridge/authlib/pipeline/holder.go @@ -81,11 +81,11 @@ func (h *Holder) RunFinish(ctx context.Context, pctx *Context, outcome Outcome) // that decide whether to buffer the request/response body. func (h *Holder) NeedsBody() bool { return h.p.Load().NeedsBody() } -// WritesBody is equivalent to h.Load().WritesBody(). Listeners read this +// WritesRequestBody is equivalent to h.Load().WritesRequestBody(). Listeners read this // when deciding whether streaming responses are safe — a pipeline with // a body mutator can't stream because the proxy can't rewrite a body // it has already started forwarding. -func (h *Holder) WritesBody() bool { return h.p.Load().WritesBody() } +func (h *Holder) WritesRequestBody() bool { return h.p.Load().WritesRequestBody() } // Ready is equivalent to h.Load().Ready(). func (h *Holder) Ready() bool { return h.p.Load().Ready() } diff --git a/authbridge/authlib/pipeline/pipeline.go b/authbridge/authlib/pipeline/pipeline.go index f46cf6de4..133215c95 100644 --- a/authbridge/authlib/pipeline/pipeline.go +++ b/authbridge/authlib/pipeline/pipeline.go @@ -364,25 +364,25 @@ func (p *Pipeline) NotReadyPlugin() string { } // NeedsBody returns true if any plugin in the pipeline needs the body -// buffered — either to read it (ReadsBody) or to mutate it (WritesBody). +// buffered — either to read it (ReadsBody) or to mutate it (WritesRequestBody). func (p *Pipeline) NeedsBody() bool { for _, plugin := range p.plugins { caps := plugin.Capabilities().Normalize() - if caps.ReadsBody || caps.WritesBody { + if caps.ReadsBody || caps.WritesRequestBody { return true } } return false } -// WritesBody returns true if any plugin in the pipeline declares -// WritesBody. Listeners use this to decide whether to diff-and-emit a -// body mutation on the wire. A pipeline with no WritesBody plugins +// WritesRequestBody returns true if any plugin in the pipeline declares +// WritesRequestBody. Listeners use this to decide whether to diff-and-emit a +// body mutation on the wire. A pipeline with no WritesRequestBody plugins // bypasses the mutation path entirely — zero overhead for the common // read-only case. -func (p *Pipeline) WritesBody() bool { +func (p *Pipeline) WritesRequestBody() bool { for _, plugin := range p.plugins { - if plugin.Capabilities().Normalize().WritesBody { + if plugin.Capabilities().Normalize().WritesRequestBody { return true } } @@ -547,19 +547,19 @@ func (p *Pipeline) dispatchFinish(parent context.Context, name string, f Finishe } // validateCapabilities enforces body-mutation ordering rules: -// - At most one WritesBody plugin per pipeline — mutation ordering would +// - At most one WritesRequestBody plugin per pipeline — mutation ordering would // otherwise be ambiguous; downstream readers can't tell which version // they're seeing. -// - A body reader (ReadsBody) must not follow a body mutator (WritesBody) — +// - A body reader (ReadsBody) must not follow a body mutator (WritesRequestBody) — // the reader would silently see mutated bytes instead of the originals. func validateCapabilities(plugins []Plugin) error { var mutatorName string var readerAfterMutator string for _, plugin := range plugins { caps := plugin.Capabilities().Normalize() - if caps.WritesBody { + if caps.WritesRequestBody { if mutatorName != "" { - return fmt.Errorf("pipeline: two plugins declare WritesBody: %q and %q — mutation ordering would be ambiguous; at most one body mutator per pipeline is allowed", mutatorName, plugin.Name()) + return fmt.Errorf("pipeline: two plugins declare WritesRequestBody: %q and %q — mutation ordering would be ambiguous; at most one body mutator per pipeline is allowed", mutatorName, plugin.Name()) } mutatorName = plugin.Name() } else if caps.ReadsBody && mutatorName != "" && readerAfterMutator == "" { diff --git a/authbridge/authlib/pipeline/pipeline_test.go b/authbridge/authlib/pipeline/pipeline_test.go index 54548e212..f1d39be46 100644 --- a/authbridge/authlib/pipeline/pipeline_test.go +++ b/authbridge/authlib/pipeline/pipeline_test.go @@ -744,7 +744,7 @@ func TestPipelineRun_ObserveSynthesizesRecordWhenPluginSkipsIt(t *testing.T) { func TestSetBody_ObserveModeIsNoop(t *testing.T) { mutator := &stubPlugin{ name: "redactor", - caps: PluginCapabilities{WritesBody: true}, + caps: PluginCapabilities{WritesRequestBody: true}, onReq: func(_ context.Context, pctx *Context) Action { pctx.SetBody([]byte("REDACTED")) return Action{Type: Continue} diff --git a/authbridge/authlib/pipeline/plugin.go b/authbridge/authlib/pipeline/plugin.go index a00a1e5ac..a8d088498 100644 --- a/authbridge/authlib/pipeline/plugin.go +++ b/authbridge/authlib/pipeline/plugin.go @@ -24,18 +24,18 @@ type PluginCapabilities struct { // a read silently sees "no body." ReadsBody bool - // WritesBody: the plugin may mutate pctx.Body / pctx.ResponseBody + // WritesRequestBody: the plugin may mutate pctx.Body / pctx.ResponseBody // (call pctx.SetBody / pctx.SetResponseBody). Implies ReadsBody — // Normalize() auto-promotes. Listener propagates the mutation to // the wire (ext_proc BodyMutation, or the outbound http.Request / // downstream http.Response for proxy listeners). // - // Pipeline.New rejects a pipeline that has more than one WritesBody + // Pipeline.New rejects a pipeline that has more than one WritesRequestBody // plugin per direction — mutation ordering would be ambiguous. - // Waypoint mode (ext_authz) cannot support WritesBody at all: + // Waypoint mode (ext_authz) cannot support WritesRequestBody at all: // ext_authz has no body-mutation field. main.go enforces this at // process boot. - WritesBody bool + WritesRequestBody bool // Requires names plugins that MUST be present in the same chain // AND appear earlier (lower index). Matches are case-sensitive @@ -71,12 +71,12 @@ type PluginCapabilities struct { Description string } -// Normalize applies WritesBody-implies-ReadsBody promotion. +// Normalize applies WritesRequestBody-implies-ReadsBody promotion. // Called by Pipeline.New for every plugin's declared capabilities so the // rest of the framework reads a normalized form. Plugins never need to // call this themselves. func (c PluginCapabilities) Normalize() PluginCapabilities { - if c.WritesBody { + if c.WritesRequestBody { c.ReadsBody = true } return c diff --git a/authbridge/authlib/plugins/contextguru/build_test.go b/authbridge/authlib/plugins/contextguru/build_test.go index 4aa82494a..059741dfc 100644 --- a/authbridge/authlib/plugins/contextguru/build_test.go +++ b/authbridge/authlib/plugins/contextguru/build_test.go @@ -12,7 +12,7 @@ import ( ) // TestBuild_InChainAfterInferenceParser confirms the plugin assembles on the -// outbound chain when a parser precedes it (RequiresAny + the single-WritesBody +// outbound chain when a parser precedes it (RequiresAny + the single-WritesRequestBody // slot are accepted together). func TestBuild_InChainAfterInferenceParser(t *testing.T) { p, err := plugins.Build([]config.PluginEntry{ diff --git a/authbridge/authlib/plugins/contextguru/plugin.go b/authbridge/authlib/plugins/contextguru/plugin.go index c016cad08..6723bbf69 100644 --- a/authbridge/authlib/plugins/contextguru/plugin.go +++ b/authbridge/authlib/plugins/contextguru/plugin.go @@ -6,7 +6,7 @@ // etc.) it replaces the body via pctx.SetBody. OnResponse is a pass-through in // v1 — model-driven restoration/expand is a later integration. // -// It is the single outbound WritesBody plugin, so it is mutually exclusive with +// It is the single outbound WritesRequestBody plugin, so it is mutually exclusive with // SPARC on the outbound chain (the pipeline refuses to build with two). It // declares RequiresAny: [inference-parser] so a parser establishes the request // is an inference call before it runs. @@ -30,13 +30,13 @@ import ( "github.com/rossoctl/cortex/authbridge/authlib/pipeline" "github.com/rossoctl/cortex/authbridge/authlib/plugins" + bschemas "github.com/maximhq/bifrost/core/schemas" "github.com/rossoctl/context-guru/apply" cgcomponents "github.com/rossoctl/context-guru/components" _ "github.com/rossoctl/context-guru/components/offload" // register offload components _ "github.com/rossoctl/context-guru/components/reformat" // register reformat components cgconfig "github.com/rossoctl/context-guru/config" cgstore "github.com/rossoctl/context-guru/store" - bschemas "github.com/maximhq/bifrost/core/schemas" ) // sentinelHeader is set on the plugin's own outbound LLM calls (via llmclient) so @@ -156,10 +156,10 @@ func (p *ContextGuru) Name() string { return "context-guru" } func (p *ContextGuru) Capabilities() pipeline.PluginCapabilities { return pipeline.PluginCapabilities{ - ReadsBody: true, - WritesBody: true, // single outbound body-writer slot (mutually exclusive with SPARC) - RequiresAny: []string{"inference-parser"}, - Description: "Compacts the outbound LLM request context before forwarding (context-guru).", + ReadsBody: true, + WritesRequestBody: true, // single outbound body-writer slot (mutually exclusive with SPARC) + RequiresAny: []string{"inference-parser"}, + Description: "Compacts the outbound LLM request context before forwarding (context-guru).", } } diff --git a/authbridge/authlib/plugins/cpex/plugin.go b/authbridge/authlib/plugins/cpex/plugin.go index bf189ea27..da50bed35 100644 --- a/authbridge/authlib/plugins/cpex/plugin.go +++ b/authbridge/authlib/plugins/cpex/plugin.go @@ -103,10 +103,10 @@ func (p *CPEX) Name() string { return "cpex" } // Capabilities declares body access and content-source requirements. // -// - ReadsBody / WritesBody: CPEX policies routinely inspect and +// - ReadsBody / WritesRequestBody: CPEX policies routinely inspect and // mutate tool args, LLM messages, and HTTP headers, so the // plugin needs the body buffered and writable. (Normalize() -// auto-promotes ReadsBody from WritesBody, so this is belt and +// auto-promotes ReadsBody from WritesRequestBody, so this is belt and // suspenders.) // // - RequiresAny: the plugin reads through pctx.ContentSources() @@ -118,10 +118,10 @@ func (p *CPEX) Name() string { return "cpex" } // surfaces in the catalog. func (p *CPEX) Capabilities() pipeline.PluginCapabilities { return pipeline.PluginCapabilities{ - ReadsBody: true, - WritesBody: true, - RequiresAny: []string{"mcp-parser", "inference-parser", "a2a-parser"}, - Description: "CPEX bridge: APL DSL + named CPEX plugins (Cedar, PII, audit, …) over a single chain step.", + ReadsBody: true, + WritesRequestBody: true, + RequiresAny: []string{"mcp-parser", "inference-parser", "a2a-parser"}, + Description: "CPEX bridge: APL DSL + named CPEX plugins (Cedar, PII, audit, …) over a single chain step.", } } diff --git a/authbridge/authlib/plugins/cpex/plugin_test.go b/authbridge/authlib/plugins/cpex/plugin_test.go index 34c66357e..1e4378e17 100644 --- a/authbridge/authlib/plugins/cpex/plugin_test.go +++ b/authbridge/authlib/plugins/cpex/plugin_test.go @@ -256,8 +256,8 @@ func TestName(t *testing.T) { func TestCapabilities_RequiresAnyParser(t *testing.T) { caps := NewCPEX().Capabilities() - if !caps.ReadsBody || !caps.WritesBody { - t.Fatal("ReadsBody/WritesBody must be true: CPEX policies routinely mutate payloads") + if !caps.ReadsBody || !caps.WritesRequestBody { + t.Fatal("ReadsBody/WritesRequestBody must be true: CPEX policies routinely mutate payloads") } want := []string{"mcp-parser", "inference-parser", "a2a-parser"} if len(caps.RequiresAny) != len(want) { diff --git a/authbridge/authlib/plugins/registry.go b/authbridge/authlib/plugins/registry.go index fc4cc0a8f..96abe1e90 100644 --- a/authbridge/authlib/plugins/registry.go +++ b/authbridge/authlib/plugins/registry.go @@ -209,11 +209,11 @@ func cloneCatalog(in []CatalogEntry) []CatalogEntry { out[i] = CatalogEntry{ Name: in[i].Name, Capabilities: pipeline.PluginCapabilities{ - ReadsBody: caps.ReadsBody, - WritesBody: caps.WritesBody, - Description: caps.Description, - Requires: append([]string(nil), caps.Requires...), - RequiresAny: append([]string(nil), caps.RequiresAny...), + ReadsBody: caps.ReadsBody, + WritesRequestBody: caps.WritesRequestBody, + Description: caps.Description, + Requires: append([]string(nil), caps.Requires...), + RequiresAny: append([]string(nil), caps.RequiresAny...), }, Fields: cloneFieldSchemas(in[i].Fields), } diff --git a/authbridge/authlib/plugins/sparc/plugin.go b/authbridge/authlib/plugins/sparc/plugin.go index f0a4ebd20..2c2e1f883 100644 --- a/authbridge/authlib/plugins/sparc/plugin.go +++ b/authbridge/authlib/plugins/sparc/plugin.go @@ -209,10 +209,10 @@ func (p *SPARC) Capabilities() pipeline.PluginCapabilities { // conversation + tool specs (both modes); mcp-parser provides the tool // call (mcp mode). RequiresAny is a static "at least one" check; the // per-mode runtime requirements are validated/handled below. - RequiresAny: []string{"inference-parser", "mcp-parser"}, - ReadsBody: true, - WritesBody: true, // MCP result (mcp mode) / completion rewrite (inference mode) - Description: "SPARC pre-tool reflection: blocks ungrounded/hallucinated tool calls.", + RequiresAny: []string{"inference-parser", "mcp-parser"}, + ReadsBody: true, + WritesRequestBody: true, // MCP result (mcp mode) / completion rewrite (inference mode) + Description: "SPARC pre-tool reflection: blocks ungrounded/hallucinated tool calls.", } } diff --git a/authbridge/authlib/plugins/sparc/plugin_test.go b/authbridge/authlib/plugins/sparc/plugin_test.go index 51c6cc02a..5457e580e 100644 --- a/authbridge/authlib/plugins/sparc/plugin_test.go +++ b/authbridge/authlib/plugins/sparc/plugin_test.go @@ -337,8 +337,8 @@ func TestInference_MCPModeOnResponseIsNoop(t *testing.T) { func TestCapabilities(t *testing.T) { caps := NewSPARC().Capabilities() - if !caps.WritesBody || !caps.ReadsBody { - t.Error("expected ReadsBody+WritesBody") + if !caps.WritesRequestBody || !caps.ReadsBody { + t.Error("expected ReadsBody+WritesRequestBody") } if len(caps.RequiresAny) == 0 { t.Error("expected RequiresAny parsers") diff --git a/authbridge/demos/context-guru/README.md b/authbridge/demos/context-guru/README.md index 2996361d7..77a1a3053 100644 --- a/authbridge/demos/context-guru/README.md +++ b/authbridge/demos/context-guru/README.md @@ -42,7 +42,7 @@ the request body before it leaves the pod. ``` The pipeline is `inference-parser → context-guru`. context-guru is the single -outbound `WritesBody` plugin (mutually exclusive with `sparc`). +outbound `WritesRequestBody` plugin (mutually exclusive with `sparc`). ## The engine: 2 deterministic reducers + extract-code @@ -150,7 +150,7 @@ inject a second sidecar). The extract-code key lives in the `cg-model-key` Secre compacted request` log line without altering the request. - **collapse stays gentle** (`head/tail: 12`); `extract` (query-aware) is the primary reducer that preserves the mid-log needle. Very aggressive collapse can drop it. -- **context-guru + SPARC are mutually exclusive** on the outbound chain (one WritesBody slot). +- **context-guru + SPARC are mutually exclusive** on the outbound chain (one WritesRequestBody slot). ## Files diff --git a/authbridge/demos/context-guru/k8s/authbridge-config.yaml b/authbridge/demos/context-guru/k8s/authbridge-config.yaml index 521931621..718b66009 100644 --- a/authbridge/demos/context-guru/k8s/authbridge-config.yaml +++ b/authbridge/demos/context-guru/k8s/authbridge-config.yaml @@ -2,7 +2,7 @@ # # Outbound chain: inference-parser (parses the OpenAI /v1/chat/completions body) # -> context-guru (compacts the agent's growing tool-output context before it is -# forwarded to the LLM). context-guru is the single outbound WritesBody plugin +# forwarded to the LLM). context-guru is the single outbound WritesRequestBody plugin # (mutually exclusive with sparc) and requires a parser ahead of it. # # THREE MODES via the context-guru entry's `on_error`: diff --git a/authbridge/docs/cpex-plugin.md b/authbridge/docs/cpex-plugin.md index bc12f3b72..354ef5f09 100644 --- a/authbridge/docs/cpex-plugin.md +++ b/authbridge/docs/cpex-plugin.md @@ -178,8 +178,8 @@ At least one must appear earlier in the chain so the parser has populated `pctx.Extensions.MCP` / `.Inference` / `.A2A` before cpex extracts CMF content. `Pipeline.Build` rejects misordered chains at boot. -cpex also declares `ReadsBody: true, WritesBody: true`. Only one -`WritesBody` plugin is permitted per direction; chaining cpex with +cpex also declares `ReadsBody: true, WritesRequestBody: true`. Only one +`WritesRequestBody` plugin is permitted per direction; chaining cpex with another mutator (e.g. an inline transformer) will fail at boot. A typical inbound chain: diff --git a/authbridge/docs/framework-architecture.md b/authbridge/docs/framework-architecture.md index fcd2d4dc0..44cfdd468 100644 --- a/authbridge/docs/framework-architecture.md +++ b/authbridge/docs/framework-architecture.md @@ -74,7 +74,7 @@ type PluginCapabilities struct { Reads []string // extension slot names this plugin reads Writes []string // extension slot names this plugin writes ReadsBody bool // plugin reads pctx.Body / pctx.ResponseBody - WritesBody bool // plugin mutates body via pctx.SetBody / pctx.SetResponseBody + WritesRequestBody bool // plugin mutates body via pctx.SetBody / pctx.SetResponseBody BodyAccess bool // deprecated: alias for ReadsBody (folded by Normalize) } ``` @@ -87,7 +87,7 @@ plugin "guardrail" reads slot "mcp" but no earlier plugin writes it `ReadsBody: true` (or the legacy `BodyAccess` alias) on *any* plugin in a chain causes `Pipeline.NeedsBody()` to return true, which the **listener** uses to negotiate Envoy's `ProcessingMode` (BUFFERED vs HEADERS-only). Without this, the gRPC ext_proc server never asks for the body and parsers see `pctx.Body == nil`. -`WritesBody: true` declares that the plugin may rewrite the body via `pctx.SetBody` / `pctx.SetResponseBody`; the listener propagates the mutation to the wire. See §6, "Body mutation" for the full body-mutation contract (capability rules, ordering constraints, Content-Encoding policy). +`WritesRequestBody: true` declares that the plugin may rewrite the body via `pctx.SetBody` / `pctx.SetResponseBody`; the listener propagates the mutation to the wire. See §6, "Body mutation" for the full body-mutation contract (capability rules, ordering constraints, Content-Encoding policy). ### `OnRequest(ctx, pctx) Action` Called when a request is entering the pipeline. Plugins typically read request headers / body, mutate one or more extension slots, and return `Continue` or `Reject`. @@ -131,7 +131,7 @@ type Context struct { - Plugins **read** any field they declared in `Capabilities.Reads`. - Plugins **write** fields they declared in `Capabilities.Writes`. By convention each extension slot has exactly one writer (the parser plugin). - Plugins read `pctx.Body` / `pctx.ResponseBody` only if they declared `ReadsBody: true` (or the deprecated `BodyAccess: true`). -- Plugins mutate body content via `pctx.SetBody(newBytes)` / `pctx.SetResponseBody(newBytes)`, and only if they declared `WritesBody: true`. Direct assignment (`pctx.Body = ...`) compiles but bypasses listener propagation and misses the Invocation + body-mutation event emission — see §6, "Body mutation." +- Plugins mutate body content via `pctx.SetBody(newBytes)` / `pctx.SetResponseBody(newBytes)`, and only if they declared `WritesRequestBody: true`. Direct assignment (`pctx.Body = ...`) compiles but bypasses listener propagation and misses the Invocation + body-mutation event emission — see §6, "Body mutation." - `Identity` is populated by whichever auth plugin ran (jwt-validation ships a `claimsIdentity` adapter around `validation.Claims`; a SAML / mTLS / custom plugin publishes its own adapter). The framework reads it through the `Identity` interface (`Subject()` / `ClientID()` / `Scopes()`) so no plugin-specific type leaks into `pipeline/`. - `Agent`, `Session` are populated by the listener before `Run`. Plugins treat them as read-only. - `ResponseBody` appears between `Run` and `RunResponse` — plugins must not read it in `OnRequest`. @@ -584,20 +584,20 @@ Always sequential. No priority / mode / fire-and-forget semantics yet. This is t ### Body mutation -A plugin that declares `WritesBody: true` may rewrite the request or response body. The framework owns the propagation to the wire; plugins only call `pctx.SetBody(newBytes)` / `pctx.SetResponseBody(newBytes)`. +A plugin that declares `WritesRequestBody: true` may rewrite the request or response body. The framework owns the propagation to the wire; plugins only call `pctx.SetBody(newBytes)` / `pctx.SetResponseBody(newBytes)`. **Capability model.** Three booleans on `PluginCapabilities`: | Field | Meaning | Listener effect | |---|---|---| | `ReadsBody` | plugin reads `pctx.Body` / `pctx.ResponseBody` | buffers the body; plugin sees the bytes | -| `WritesBody` | plugin may call `pctx.SetBody` / `pctx.SetResponseBody` | implies `ReadsBody`; propagates mutations | +| `WritesRequestBody` | plugin may call `pctx.SetBody` / `pctx.SetResponseBody` | implies `ReadsBody`; propagates mutations | | `BodyAccess` (deprecated) | legacy alias for `ReadsBody` | folded by `Normalize()`, removed in a future release | `pipeline.New` enforces two rules at build time: -1. **At most one `WritesBody` plugin per pipeline.** Multiple mutators would have ambiguous ordering semantics; the error names both plugins so an operator debugging pod logs knows which two to reconcile. -2. **`WritesBody` cannot precede a `ReadsBody`-only plugin.** A reader expects to see the original bytes; putting a mutator before it would silently feed the reader the post-rewrite content. +1. **At most one `WritesRequestBody` plugin per pipeline.** Multiple mutators would have ambiguous ordering semantics; the error names both plugins so an operator debugging pod logs knows which two to reconcile. +2. **`WritesRequestBody` cannot precede a `ReadsBody`-only plugin.** A reader expects to see the original bytes; putting a mutator before it would silently feed the reader the post-rewrite content. **Mutation helpers.** `SetBody` / `SetResponseBody` replace the byte slice and flip an internal `bodyMutated` / `responseBodyMutated` flag that listeners read via `pctx.BodyMutated()` / `pctx.ResponseBodyMutated()`. They also auto-emit: @@ -802,7 +802,7 @@ The plugin interface is **not** semver-stable yet (AuthBridge is pre-1.0). Chang - **`pctx.Record` helpers**: `Allow` / `Skip` / `Observe` / `Modify` / `Record` / `DenyAndRecord` on `Context`. Framework-managed attribution (`currentPlugin`, `currentPhase`, `Path`) fills Invocation fields automatically. - **Open plugin registry**: plugins self-register from `init()` via `plugins.RegisterPlugin`. Third-party plugins in external modules drop in via a side-effect import. Closed `registry` map literal removed. - **Config hot-reload**: new `pipeline.Holder` (atomic wrapper) + `authlib/reloader` package (fsnotify-driven). Listeners receive `*Holder` instead of `*Pipeline`; the reloader atomically swaps the holder's contents when the config file changes. `mode` and `listener.*` edits are refused (pod restart required); any other change is picked up within the kubelet sync window (~60s). See §9. -- **Body mutation**: `PluginCapabilities.BodyAccess` split into `ReadsBody` / `WritesBody`. New `pctx.SetBody` / `pctx.SetResponseBody` helpers flip a mutation flag; all three listeners (extproc / forwardproxy / reverseproxy) propagate the rewrite to the upstream with correct `Content-Length` and cleared `Content-Encoding`. `BodyAccess` kept as deprecated alias. See §6, "Body mutation." +- **Body mutation**: `PluginCapabilities.BodyAccess` split into `ReadsBody` / `WritesRequestBody`. New `pctx.SetBody` / `pctx.SetResponseBody` helpers flip a mutation flag; all three listeners (extproc / forwardproxy / reverseproxy) propagate the rewrite to the upstream with correct `Content-Length` and cleared `Content-Encoding`. `BodyAccess` kept as deprecated alias. See §6, "Body mutation." - **Detyped framework**: `pipeline/` no longer imports plugin-specific packages. **Breaking**: `Context.Claims *validation.Claims` → `Context.Identity Identity` (interface with `Subject()`/`ClientID()`/`Scopes()`); plugins publish adapters. `Context.Route` removed (was dead code). `Invocation`'s nine jwt-validation + token-exchange specific fields (`ExpectedIssuer`, `TokenSubject`, `RouteHost`, `CacheHit`, etc.) collapsed into `Details map[string]string`; built-in plugins migrated to `Details["expected_issuer"]` etc. `SessionEvent.TargetAudience` removed (was only populated from dead `pctx.Route`). Third-party plugins get a clean diagnostic slot they can populate without framework edits. - **Single-owner packages relocated**: `authlib/validation` → `authlib/plugins/jwtvalidation/validation`. `authlib/exchange` / `authlib/cache` / `authlib/spiffe` → `authlib/plugins/tokenexchange/{exchange,cache,spiffe}`. Each plugin now lives in its own directory (`plugins/jwtvalidation/plugin.go`, `plugins/tokenexchange/plugin.go`) and self-registers via its own init(). `authlib/bypass`, `authlib/routing`, `authlib/auth` stay shared. - **Plugin relationship declarations**: `PluginCapabilities` extended with four chain-scoped fields — `Requires` (all-must-be-earlier), `RequiresAny` (at-least-one-earlier), `After` (soft ordering), `Claims` (mutex on a semantic resource). Validated at `plugins.Build` time (startup + hot-reload); all errors per chain are collected into one report. `authlib/contracts/claims.go` ships `ClaimAuthorizationHeader` as the initial canonical claim constant. `token-exchange` and `token-broker` migrated to declare it, so configuring both on the same outbound chain now fails startup instead of silently clobbering each other's Authorization header. See [`plugin-reference.md` "Declaring plugin relationships"](./plugin-reference.md#declaring-plugin-relationships). @@ -821,9 +821,9 @@ Breaking changes will be announced in `authbridge/CHANGELOG.md` (TBD) before a 1 **Package sources:** -- `pipeline.go` — `Pipeline` type, `New`, `Run`, `RunResponse`, `Start`, `Stop`, `Plugins`, `NeedsBody`, `WritesBody`. +- `pipeline.go` — `Pipeline` type, `New`, `Run`, `RunResponse`, `Start`, `Stop`, `Plugins`, `NeedsBody`, `WritesRequestBody`. - `holder.go` — `Holder`, the atomic slot listeners hold in place of a raw `*Pipeline`. -- `plugin.go` — `Plugin` interface, `PluginCapabilities` (with `ReadsBody` / `WritesBody` / deprecated `BodyAccess` + `Normalize()`; chain-scoped relationship fields `Requires` / `RequiresAny` / `After` / `Claims`), `Configurable`, `Initializer`, `Shutdowner`, `Readier`, `Finisher`. +- `plugin.go` — `Plugin` interface, `PluginCapabilities` (with `ReadsBody` / `WritesRequestBody` / deprecated `BodyAccess` + `Normalize()`; chain-scoped relationship fields `Requires` / `RequiresAny` / `After` / `Claims`), `Configurable`, `Initializer`, `Shutdowner`, `Readier`, `Finisher`. - `outcome.go` — `Outcome` struct + `OutcomeAction` (allow / deny / error) for `Finisher` consumers; `Context.Outcome()` getter. - `action.go` — `Action`, `ActionType`, `Violation`, helper constructors (`Deny`, `DenyStatus`, `DenyWithDetails`, `Challenge`, `RateLimited`), `StatusFromCode`. - `context.go` — `Context`, `Direction`, `AgentIdentity`, the `pctx.Record` / `Allow` / `Skip` / `Observe` / `Modify` / `DenyAndRecord` helpers, and `pctx.SetBody` / `SetResponseBody` / `BodyMutated` / `ResponseBodyMutated` for body mutation. diff --git a/authbridge/docs/plugin-reference.md b/authbridge/docs/plugin-reference.md index bd68a476e..f4a855339 100644 --- a/authbridge/docs/plugin-reference.md +++ b/authbridge/docs/plugin-reference.md @@ -217,7 +217,7 @@ fail loud before serving traffic. ```go type PluginCapabilities struct { ReadsBody bool - WritesBody bool + WritesRequestBody bool Requires []string // ALL must be present + earlier (hard) RequiresAny []string // AT LEAST ONE must be present + run after it (hard) @@ -684,7 +684,7 @@ in it. ## Body mutation Plugins that need to rewrite request or response bodies declare -`WritesBody: true` and call the `pctx.SetBody` / `pctx.SetResponseBody` +`WritesRequestBody: true` and call the `pctx.SetBody` / `pctx.SetResponseBody` helpers. The framework propagates the rewrite to the wire, emits a `modify`-action Invocation, and publishes a `body-mutation/event` entry in `pctx.Extensions.Custom` with length delta + sha256 @@ -700,25 +700,25 @@ before/after (never the raw body). ```go type PluginCapabilities struct { ReadsBody bool // plugin reads pctx.Body / pctx.ResponseBody - WritesBody bool // plugin may call pctx.SetBody / pctx.SetResponseBody + WritesRequestBody bool // plugin may call pctx.SetBody / pctx.SetResponseBody } ``` - `ReadsBody`: listener buffers the body; plugin sees bytes. -- `WritesBody`: implies `ReadsBody`. Listener propagates `pctx.SetBody` +- `WritesRequestBody`: implies `ReadsBody`. Listener propagates `pctx.SetBody` rewrites to the upstream (and `pctx.SetResponseBody` to the downstream client). ### Build-time validation (enforced by `pipeline.New`) -- At most **one** `WritesBody` plugin per pipeline. Two mutators in +- At most **one** `WritesRequestBody` plugin per pipeline. Two mutators in the same direction would produce ambiguous ordering; `New` rejects with an error naming both plugins. -- A `WritesBody` plugin cannot precede a `ReadsBody`-only plugin. The +- A `WritesRequestBody` plugin cannot precede a `ReadsBody`-only plugin. The reader must see the original bytes. - Waypoint mode (ext_authz listener) cannot propagate body mutations — the ext_authz API has no body-mutation field. Do not combine - `WritesBody: true` plugins with `mode: waypoint`. + `WritesRequestBody: true` plugins with `mode: waypoint`. ### Mutation helpers diff --git a/authbridge/docs/plugin-tutorial.md b/authbridge/docs/plugin-tutorial.md index 201e4923a..b6060cddd 100644 --- a/authbridge/docs/plugin-tutorial.md +++ b/authbridge/docs/plugin-tutorial.md @@ -209,12 +209,12 @@ both stay nil even if you try to read them. ### Mutating the body If your plugin needs to **rewrite** the body — prompt-redaction, output -filtering, content transformation — declare `WritesBody` and call +filtering, content transformation — declare `WritesRequestBody` and call `pctx.SetBody` / `pctx.SetResponseBody`: ```go func (p *Redactor) Capabilities() pipeline.PluginCapabilities { - return pipeline.PluginCapabilities{WritesBody: true} // implies ReadsBody + return pipeline.PluginCapabilities{WritesRequestBody: true} // implies ReadsBody } func (p *Redactor) OnRequest(_ context.Context, pctx *pipeline.Context) pipeline.Action { @@ -232,9 +232,9 @@ for `SetResponseBody`) with a correct `Content-Length` and a cleared (never the raw body content). **Rules enforced by `pipeline.New`:** -- At most one `WritesBody` plugin per pipeline. Two mutators = ambiguous +- At most one `WritesRequestBody` plugin per pipeline. Two mutators = ambiguous ordering → build fails at startup. -- A `WritesBody` plugin must run **after** any `ReadsBody`-only plugin. +- A `WritesRequestBody` plugin must run **after** any `ReadsBody`-only plugin. Readers see the original bytes; a mutator in front would silently feed them post-rewrite content. From 6440927cb8ba6bdb3484b0e59f60e3ffd311628b Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 11:20:53 -0400 Subject: [PATCH 04/28] feat(authbridge): Split body-write capability by direction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PluginCapabilities.WritesRequestBody was doing double duty: both proxy listeners consulted it to decide whether an SSE response could be relayed incrementally, so a plugin that rewrites only the *request* body disabled *response* streaming for bytes it never touches. The cost is latency and feel rather than correctness — the buffered path restores the body verbatim — but a long completion arriving in one lump after a silent wait is the first thing anyone notices. Add WritesResponseBody and make it the streaming predicate: - Pipeline.WritesResponseBody() (and the Holder delegate) is what the forward and reverse proxies now gate the buffered fallback on. WritesRequestBody keeps gating request propagation. - Normalize() promotes ReadsBody from either write flag. - validateCapabilities is direction-aware: at most one mutator per direction, while reader-ordering still trips on either flag. Every configuration that exists in-tree today validates exactly as before, since all current mutators write requests. - sparc and cpex declare both flags and keep the buffered path. context-guru is request-only and regains streaming it never needed to lose. Two adjacent fixes the split exposed: - cloneCatalog copied capability fields one at a time, silently dropping any field added later from /v1/plugins. Replaced with a struct copy plus explicit slice reallocation, covered by a reflection round-trip that fails if a future field is missed. - SetBody's godoc claimed an undeclared mutation stays in-memory. It does not: bodyMutated is set unconditionally outside observe mode and the listeners gate purely on it. Corrected to describe actual behaviour and flag the divergence, rather than adding enforcement that would silently break out-of-tree plugins. Left as it was, the comment made "just don't declare the capability" look like a legitimate way to keep streaming. Tests: a truth table over the four plugin shapes; direction tests in both listeners asserting a request-only writer still receives one frame per SSE event plus a final (4 calls) while a response writer receives a single buffered delivery (1 call) — the frame count is what discriminates the paths, which the previous fallback test did not do; validateCapabilities table assertions; and the cloneCatalog round-trip. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- authbridge/authlib/listener/extproc/server.go | 2 +- .../extproc/server_contentlength_test.go | 2 +- .../authlib/listener/forwardproxy/server.go | 10 +- .../forwardproxy/streaming_direction_test.go | 101 +++++++++++ .../listener/forwardproxy/streaming_test.go | 28 +-- .../authlib/listener/reverseproxy/server.go | 4 +- .../reverseproxy/streaming_direction_test.go | 118 +++++++++++++ .../authlib/pipeline/bodydirection_test.go | 165 ++++++++++++++++++ authbridge/authlib/pipeline/context.go | 24 ++- authbridge/authlib/pipeline/holder.go | 14 +- authbridge/authlib/pipeline/pipeline.go | 51 +++++- authbridge/authlib/pipeline/plugin.go | 39 +++-- authbridge/authlib/plugins/cpex/plugin.go | 9 +- authbridge/authlib/plugins/registry.go | 16 +- .../plugins/registry_capsclone_test.go | 84 +++++++++ authbridge/authlib/plugins/sparc/plugin.go | 9 +- 16 files changed, 612 insertions(+), 64 deletions(-) create mode 100644 authbridge/authlib/listener/forwardproxy/streaming_direction_test.go create mode 100644 authbridge/authlib/listener/reverseproxy/streaming_direction_test.go create mode 100644 authbridge/authlib/pipeline/bodydirection_test.go create mode 100644 authbridge/authlib/plugins/registry_capsclone_test.go diff --git a/authbridge/authlib/listener/extproc/server.go b/authbridge/authlib/listener/extproc/server.go index dcb8c1348..3cfe5ed67 100644 --- a/authbridge/authlib/listener/extproc/server.go +++ b/authbridge/authlib/listener/extproc/server.go @@ -657,7 +657,7 @@ func (s *Server) handleResponseBody(ctx context.Context, body []byte, pctx *pipe s.recordOutboundResponseSession(pctx) } - // A plugin that declared WritesRequestBody: true and called pctx.SetResponseBody + // A plugin that declared WritesResponseBody: true and called pctx.SetResponseBody // flips the ResponseBodyMutated flag. Emit the replacement bytes via // BodyMutation so Envoy rewrites the downstream response; otherwise // pass through with no mutation. The flag avoids the O(n) string diff --git a/authbridge/authlib/listener/extproc/server_contentlength_test.go b/authbridge/authlib/listener/extproc/server_contentlength_test.go index 17401ab0f..d7b91b4af 100644 --- a/authbridge/authlib/listener/extproc/server_contentlength_test.go +++ b/authbridge/authlib/listener/extproc/server_contentlength_test.go @@ -60,7 +60,7 @@ type responseMutator struct{ newBody []byte } func (*responseMutator) Name() string { return "response-mutator" } func (*responseMutator) Capabilities() pipeline.PluginCapabilities { - return pipeline.PluginCapabilities{WritesRequestBody: true} + return pipeline.PluginCapabilities{WritesResponseBody: true} } func (*responseMutator) OnRequest(context.Context, *pipeline.Context) pipeline.Action { return pipeline.Action{Type: pipeline.Continue} diff --git a/authbridge/authlib/listener/forwardproxy/server.go b/authbridge/authlib/listener/forwardproxy/server.go index b4fce2d4b..9d822be93 100644 --- a/authbridge/authlib/listener/forwardproxy/server.go +++ b/authbridge/authlib/listener/forwardproxy/server.go @@ -404,10 +404,12 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, isBridge // forwarding is incompatible with streaming) — fall back to // buffered with a warning log instead. if isEventStream(resp.Header.Get("Content-Type")) && resp.Body != nil { - if s.OutboundPipeline.WritesRequestBody() { - // A body mutator needs the whole body to rewrite it, so it - // can't stream — fall back to the buffered path with a warning. - slog.Warn("forward-proxy: text/event-stream response with WritesRequestBody plugin — falling back to buffered path", "host", r.Host) + if s.OutboundPipeline.WritesResponseBody() { + // A response mutator needs the whole response to rewrite it, so + // it can't stream — fall back to the buffered path with a warning. + // A request-only mutator does NOT land here: it never touches + // these bytes, so the relay stays incremental. + slog.Warn("forward-proxy: text/event-stream response with WritesResponseBody plugin — falling back to buffered path", "host", r.Host) } else if s.OutboundPipeline.HasStreamingResponders() { // Streaming-aware plugins (inference-parser, a2a-parser) parse // each SSE frame; handleStreamingResponse re-frames via sseframe. diff --git a/authbridge/authlib/listener/forwardproxy/streaming_direction_test.go b/authbridge/authlib/listener/forwardproxy/streaming_direction_test.go new file mode 100644 index 000000000..6b0a5a760 --- /dev/null +++ b/authbridge/authlib/listener/forwardproxy/streaming_direction_test.go @@ -0,0 +1,101 @@ +package forwardproxy + +import ( + "bytes" + "fmt" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" +) + +// newResponseWritingProbe is a probe shaped like sparc / cpex: it rewrites the +// upstream response, so it genuinely cannot stream. +func newResponseWritingProbe() *streamingProbe { + return &streamingProbe{ + caps: pipeline.PluginCapabilities{ + ReadsBody: true, + WritesResponseBody: true, + }, + } +} + +// TestForwardProxy_Streaming_RequestOnlyWriterKeepsStreaming is the point of +// the directional split. Before it, any plugin declaring the single WritesBody +// flag forfeited incremental SSE relay — including a plugin that only ever +// rewrites the *request*, for bytes it never touches. tool-prune and +// context-guru are exactly that shape. +// +// The assertion is the frame count: the streaming path delivers one call per +// frame plus a final last=true (4 for 3 frames), where the buffered path +// delivers a single last=true call. A regression that reattached the fallback +// to the request flag would collapse this to 1. +func TestForwardProxy_Streaming_RequestOnlyWriterKeepsStreaming(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher := w.(http.Flusher) + for i := 1; i <= 3; i++ { + fmt.Fprintf(w, "data: {\"id\":%d}\n\n", i) + flusher.Flush() + time.Sleep(20 * time.Millisecond) + } + })) + defer upstream.Close() + + probe := newStreamingProbe(true) // WritesRequestBody only + if probe.caps.WritesResponseBody { + t.Fatal("probe must not declare WritesResponseBody for this test") + } + pipe, err := pipeline.New([]pipeline.Plugin{probe}) + if err != nil { + t.Fatalf("New pipeline: %v", err) + } + if !pipe.WritesRequestBody() { + t.Fatal("pipeline should report WritesRequestBody") + } + if pipe.WritesResponseBody() { + t.Fatal("pipeline must NOT report WritesResponseBody") + } + + srv, err := NewServer(pipeline.NewHolder(pipe), nil, nil) + if err != nil { + t.Fatalf("NewServer: %v", err) + } + proxy := httptest.NewServer(srv.Handler()) + defer proxy.Close() + + proxyClient := &http.Client{ + Transport: &http.Transport{Proxy: http.ProxyURL(mustParseURL(proxy.URL))}, + } + req, _ := http.NewRequest("GET", upstream.URL+"/stream", nil) + resp, err := proxyClient.Do(req) + if err != nil { + t.Fatalf("Do: %v", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("ReadAll: %v", err) + } + if got := bytes.Count(body, []byte("data:")); got != 3 { + t.Errorf("body has %d data: lines, want 3 — body=%q", got, body) + } + + frames, lasts := probe.snapshot() + if len(frames) != 4 { + t.Fatalf("plugin saw %d calls, want 4 (3 frames + final) — a request-only writer must keep streaming; lasts=%v", len(frames), lasts) + } + for i := 0; i < 3; i++ { + if lasts[i] { + t.Errorf("frame %d last=true, want false", i) + } + } + if !lasts[3] { + t.Error("final call last=false, want true") + } +} diff --git a/authbridge/authlib/listener/forwardproxy/streaming_test.go b/authbridge/authlib/listener/forwardproxy/streaming_test.go index c14eaa822..8c7c84177 100644 --- a/authbridge/authlib/listener/forwardproxy/streaming_test.go +++ b/authbridge/authlib/listener/forwardproxy/streaming_test.go @@ -158,13 +158,16 @@ func TestForwardProxy_Streaming_FramesFlowThrough(t *testing.T) { } } -// TestForwardProxy_Streaming_WritesRequestBodyFallsBackToBuffered asserts the -// safety guard: a pipeline with a WritesRequestBody plugin can't take the -// streaming path (the plugin can't rewrite a body we've already +// TestForwardProxy_Streaming_WritesResponseBodyFallsBackToBuffered asserts the +// safety guard: a pipeline with a WritesResponseBody plugin can't take the +// streaming path (the plugin can't rewrite a response we've already // started forwarding). The proxy logs a warning and falls back to // buffered, so the response is delivered correctly even though it // loses the streaming property. -func TestForwardProxy_Streaming_WritesRequestBodyFallsBackToBuffered(t *testing.T) { +// +// Only the response flag does this. The request-only case is covered by +// TestForwardProxy_Streaming_RequestOnlyWriterKeepsStreaming. +func TestForwardProxy_Streaming_WritesResponseBodyFallsBackToBuffered(t *testing.T) { upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) @@ -174,7 +177,7 @@ func TestForwardProxy_Streaming_WritesRequestBodyFallsBackToBuffered(t *testing. })) defer upstream.Close() - probe := newStreamingProbe(true) // WritesRequestBody=true → buffered fallback + probe := newResponseWritingProbe() // WritesResponseBody=true → buffered fallback pipe, err := pipeline.New([]pipeline.Plugin{probe}) if err != nil { t.Fatalf("New pipeline: %v", err) @@ -201,11 +204,16 @@ func TestForwardProxy_Streaming_WritesRequestBodyFallsBackToBuffered(t *testing. if !bytes.Contains(body, []byte(`{"id":1}`)) { t.Errorf("body did not contain expected payload: %q", body) } - // Buffered path: streaming-aware plugins still see one last=true - // frame carrying the whole body. Sanity-check. - _, lasts := probe.snapshot() - if len(lasts) == 0 || !lasts[len(lasts)-1] { - t.Errorf("last call lasts = %v; expected final last=true on buffered fallback", lasts) + // Buffered path: streaming-aware plugins see exactly ONE last=true + // delivery carrying the whole body. The count is what discriminates + // buffered from streaming — the streaming path would deliver one call + // per frame plus a final — so assert it rather than just the flag. + frames, lasts := probe.snapshot() + if len(frames) != 1 { + t.Fatalf("plugin saw %d calls, want exactly 1 on the buffered path — lasts=%v", len(frames), lasts) + } + if !lasts[0] { + t.Errorf("buffered delivery lasts = %v, want [true]", lasts) } } diff --git a/authbridge/authlib/listener/reverseproxy/server.go b/authbridge/authlib/listener/reverseproxy/server.go index 88516a37c..81b4d3988 100644 --- a/authbridge/authlib/listener/reverseproxy/server.go +++ b/authbridge/authlib/listener/reverseproxy/server.go @@ -446,8 +446,8 @@ func (s *Server) modifyResponse(resp *http.Response) error { if isEventStream(resp.Header.Get("Content-Type")) && s.InboundPipeline.HasStreamingResponders() && resp.Body != nil { - if s.InboundPipeline.WritesRequestBody() { - slog.Warn("reverse-proxy: text/event-stream response with WritesRequestBody plugin — falling back to buffered path", "host", pctx.Host) + if s.InboundPipeline.WritesResponseBody() { + slog.Warn("reverse-proxy: text/event-stream response with WritesResponseBody plugin — falling back to buffered path", "host", pctx.Host) } else { s.installStreamingResponseBody(resp, pctx) // Strip Content-Length — the framing reader doesn't know diff --git a/authbridge/authlib/listener/reverseproxy/streaming_direction_test.go b/authbridge/authlib/listener/reverseproxy/streaming_direction_test.go new file mode 100644 index 000000000..de4e12f92 --- /dev/null +++ b/authbridge/authlib/listener/reverseproxy/streaming_direction_test.go @@ -0,0 +1,118 @@ +package reverseproxy + +import ( + "bytes" + "fmt" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" +) + +// newResponseWritingProbe is a probe shaped like sparc / cpex: it rewrites the +// upstream response, so it genuinely cannot stream. +func newResponseWritingProbe() *streamingProbe { + return &streamingProbe{ + caps: pipeline.PluginCapabilities{ + ReadsBody: true, + WritesResponseBody: true, + }, + } +} + +func sseBackend(t *testing.T, frames int) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher := w.(http.Flusher) + for i := 1; i <= frames; i++ { + fmt.Fprintf(w, "data: {\"event\":%d}\n\n", i) + flusher.Flush() + time.Sleep(20 * time.Millisecond) + } + })) +} + +func serveWith(t *testing.T, p *streamingProbe, backendURL string) *httptest.Server { + t.Helper() + pipe, err := pipeline.New([]pipeline.Plugin{p}) + if err != nil { + t.Fatalf("New pipeline: %v", err) + } + srv, err := NewServer(pipeline.NewHolder(pipe), nil, backendURL, nil) + if err != nil { + t.Fatalf("NewServer: %v", err) + } + return httptest.NewServer(srv.Handler()) +} + +// TestReverseProxy_Streaming_RequestOnlyWriterKeepsStreaming mirrors the +// forward-proxy case: a plugin that rewrites only the request must not cost +// the inbound listener incremental SSE relay. +func TestReverseProxy_Streaming_RequestOnlyWriterKeepsStreaming(t *testing.T) { + backend := sseBackend(t, 3) + defer backend.Close() + + probe := newStreamingProbe(true) // WritesRequestBody only + proxy := serveWith(t, probe, backend.URL) + defer proxy.Close() + + resp, err := http.Get(proxy.URL + "/stream") + if err != nil { + t.Fatalf("Get: %v", err) + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("ReadAll: %v", err) + } + if got := bytes.Count(body, []byte("data:")); got != 3 { + t.Errorf("body has %d data: lines, want 3 — body=%q", got, body) + } + + frames, lasts := probe.snapshot() + if len(frames) != 4 { + t.Fatalf("plugin saw %d calls, want 4 (3 frames + final) — a request-only writer must keep streaming; lasts=%v", len(frames), lasts) + } + if !lasts[3] { + t.Error("final call last=false, want true") + } +} + +// TestReverseProxy_Streaming_WritesResponseBodyFallsBackToBuffered asserts the +// safety guard still holds for the direction that actually needs it: a +// response mutator forfeits streaming and receives one buffered delivery. +func TestReverseProxy_Streaming_WritesResponseBodyFallsBackToBuffered(t *testing.T) { + backend := sseBackend(t, 3) + defer backend.Close() + + probe := newResponseWritingProbe() + proxy := serveWith(t, probe, backend.URL) + defer proxy.Close() + + resp, err := http.Get(proxy.URL + "/stream") + if err != nil { + t.Fatalf("Get: %v", err) + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("ReadAll: %v", err) + } + // The buffered path still delivers every byte, just not incrementally. + if got := bytes.Count(body, []byte("data:")); got != 3 { + t.Errorf("body has %d data: lines, want 3 — body=%q", got, body) + } + + frames, lasts := probe.snapshot() + if len(frames) != 1 { + t.Fatalf("plugin saw %d calls, want exactly 1 on the buffered path — lasts=%v", len(frames), lasts) + } + if !lasts[0] { + t.Errorf("buffered delivery lasts = %v, want [true]", lasts) + } +} diff --git a/authbridge/authlib/pipeline/bodydirection_test.go b/authbridge/authlib/pipeline/bodydirection_test.go new file mode 100644 index 000000000..28bb8f016 --- /dev/null +++ b/authbridge/authlib/pipeline/bodydirection_test.go @@ -0,0 +1,165 @@ +package pipeline + +import ( + "strings" + "testing" +) + +// TestCapabilities_Normalize_EitherWriteImpliesReadsBody: both write flags +// promote ReadsBody, so a mutator of either direction always satisfies the +// "must have read the body" invariant. +func TestCapabilities_Normalize_EitherWriteImpliesReadsBody(t *testing.T) { + tests := []struct { + name string + in PluginCapabilities + }{ + {"request writer", PluginCapabilities{WritesRequestBody: true}}, + {"response writer", PluginCapabilities{WritesResponseBody: true}}, + {"both", PluginCapabilities{WritesRequestBody: true, WritesResponseBody: true}}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if !tc.in.Normalize().ReadsBody { + t.Error("Normalize() should promote ReadsBody") + } + }) + } + if (PluginCapabilities{}).Normalize().ReadsBody { + t.Error("empty capabilities must not gain ReadsBody") + } +} + +// TestPipeline_BodyWritePredicates_TruthTable pins the two predicates across +// all four plugin shapes. WritesResponseBody is the SSE streaming predicate, +// so a request-only writer reporting true here would silently cost every +// caller incremental relay — the exact defect this split fixes. +func TestPipeline_BodyWritePredicates_TruthTable(t *testing.T) { + tests := []struct { + name string + caps PluginCapabilities + wantReq, wantResp bool + wantNeedsBody bool + }{ + { + name: "request-only writer (tool-prune, context-guru)", + caps: PluginCapabilities{WritesRequestBody: true}, + wantReq: true, + wantResp: false, + wantNeedsBody: true, + }, + { + name: "response-only writer", + caps: PluginCapabilities{WritesResponseBody: true}, + wantReq: false, + wantResp: true, + wantNeedsBody: true, + }, + { + name: "both directions (sparc, cpex)", + caps: PluginCapabilities{WritesRequestBody: true, WritesResponseBody: true}, + wantReq: true, + wantResp: true, + wantNeedsBody: true, + }, + { + name: "neither (pure reader)", + caps: PluginCapabilities{ReadsBody: true}, + wantReq: false, + wantResp: false, + wantNeedsBody: true, + }, + { + name: "neither, no body at all", + caps: PluginCapabilities{}, + wantReq: false, + wantResp: false, + wantNeedsBody: false, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + p := mustBuild(t, &stubPlugin{name: "p", caps: tc.caps}) + if got := p.WritesRequestBody(); got != tc.wantReq { + t.Errorf("WritesRequestBody() = %v, want %v", got, tc.wantReq) + } + if got := p.WritesResponseBody(); got != tc.wantResp { + t.Errorf("WritesResponseBody() = %v, want %v", got, tc.wantResp) + } + if got := p.NeedsBody(); got != tc.wantNeedsBody { + t.Errorf("NeedsBody() = %v, want %v", got, tc.wantNeedsBody) + } + }) + } +} + +// TestValidateCapabilities_Directional: the mutator-exclusivity rule is +// per-direction, and reader-ordering is triggered by either write flag. +// Crucially, every combination that exists in-tree today validates exactly +// as it did before the split. +func TestValidateCapabilities_Directional(t *testing.T) { + req := PluginCapabilities{WritesRequestBody: true} + resp := PluginCapabilities{WritesResponseBody: true} + both := PluginCapabilities{WritesRequestBody: true, WritesResponseBody: true} + reader := PluginCapabilities{ReadsBody: true} + + tests := []struct { + name string + plugins []Plugin + wantErr string + }{ + { + name: "two request writers rejected", + plugins: []Plugin{&stubPlugin{name: "a", caps: req}, &stubPlugin{name: "b", caps: req}}, + wantErr: "WritesRequestBody", + }, + { + name: "two response writers rejected", + plugins: []Plugin{&stubPlugin{name: "a", caps: resp}, &stubPlugin{name: "b", caps: resp}}, + wantErr: "WritesResponseBody", + }, + { + name: "one of each direction is fine — they never collide", + plugins: []Plugin{&stubPlugin{name: "a", caps: req}, &stubPlugin{name: "b", caps: resp}}, + }, + { + name: "two both-direction writers rejected on the request rule first", + plugins: []Plugin{&stubPlugin{name: "a", caps: both}, &stubPlugin{name: "b", caps: both}}, + wantErr: "WritesRequestBody", + }, + { + name: "reader before mutator is fine", + plugins: []Plugin{&stubPlugin{name: "r", caps: reader}, &stubPlugin{name: "m", caps: req}}, + }, + { + name: "reader after request mutator rejected", + plugins: []Plugin{&stubPlugin{name: "m", caps: req}, &stubPlugin{name: "r", caps: reader}}, + wantErr: "reads body after mutator", + }, + { + name: "reader after response mutator rejected too", + plugins: []Plugin{&stubPlugin{name: "m", caps: resp}, &stubPlugin{name: "r", caps: reader}}, + wantErr: "reads body after mutator", + }, + { + name: "today's shape: parser then single both-direction mutator", + plugins: []Plugin{&stubPlugin{name: "inference-parser", caps: reader}, &stubPlugin{name: "sparc", caps: both}}, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + err := validateCapabilities(tc.plugins) + if tc.wantErr == "" { + if err != nil { + t.Fatalf("validateCapabilities() = %v, want nil", err) + } + return + } + if err == nil { + t.Fatalf("validateCapabilities() = nil, want error containing %q", tc.wantErr) + } + if !strings.Contains(err.Error(), tc.wantErr) { + t.Errorf("error %q does not contain %q", err.Error(), tc.wantErr) + } + }) + } +} diff --git a/authbridge/authlib/pipeline/context.go b/authbridge/authlib/pipeline/context.go index dd882fca8..4d93c9170 100644 --- a/authbridge/authlib/pipeline/context.go +++ b/authbridge/authlib/pipeline/context.go @@ -387,12 +387,20 @@ func (c *Context) DenyAndRecord(reason, code, message string) Action { return Deny(code, message) } -// SetBody replaces the request body with newBody. Only meaningful when -// the plugin declares WritesRequestBody: true in its Capabilities — the -// listener consults pctx.BodyMutated() after Run to decide whether to -// emit the new bytes on the wire. Plugins without WritesRequestBody that call -// SetBody mutate the in-memory Context (readers downstream see the -// change), but the wire is unchanged. +// SetBody replaces the request body with newBody. A plugin that calls it +// must declare WritesRequestBody: true in its Capabilities — the listener +// consults pctx.BodyMutated() after Run to decide whether to emit the new +// bytes on the wire. +// +// NOTE — the capability is a contract, not an enforcement. SetBody sets +// bodyMutated unconditionally outside observe mode, and the listeners gate +// purely on pctx.BodyMutated(), so a plugin that calls SetBody WITHOUT +// declaring the capability still reaches the wire. This divergence is +// documented rather than closed: adding the enforcement silently would +// break any out-of-tree plugin relying on today's behaviour, so it needs +// its own compatibility review. Do not read it as licence to skip the +// declaration in order to keep response streaming — declaring +// WritesRequestBody costs no streaming (see PluginCapabilities). // // Under ErrorPolicyObserve (shadow mode) SetBody is a NO-OP on bytes: // the in-memory body is not replaced, bodyMutated stays false, and @@ -432,6 +440,10 @@ func (c *Context) SetBody(newBody []byte) { // Invocation + body-mutation/event emitted; never logs the body — // and the same observe-mode suppression: under ErrorPolicyObserve the // response body is untouched and the Invocation is marked Shadow=true. +// +// A plugin that calls this must declare WritesResponseBody: true. That +// declaration is what makes listeners buffer the response instead of +// relaying SSE frames incrementally, so it must not be omitted. func (c *Context) SetResponseBody(newBody []byte) { if c.inFinish { slog.Warn("pipeline: plugin called pctx.SetResponseBody during OnFinish — dropped (response already sent)", diff --git a/authbridge/authlib/pipeline/holder.go b/authbridge/authlib/pipeline/holder.go index fb9f881c9..0201edec9 100644 --- a/authbridge/authlib/pipeline/holder.go +++ b/authbridge/authlib/pipeline/holder.go @@ -81,12 +81,18 @@ func (h *Holder) RunFinish(ctx context.Context, pctx *Context, outcome Outcome) // that decide whether to buffer the request/response body. func (h *Holder) NeedsBody() bool { return h.p.Load().NeedsBody() } -// WritesRequestBody is equivalent to h.Load().WritesRequestBody(). Listeners read this -// when deciding whether streaming responses are safe — a pipeline with -// a body mutator can't stream because the proxy can't rewrite a body -// it has already started forwarding. +// WritesRequestBody is equivalent to h.Load().WritesRequestBody(). +// Listeners read this when deciding whether to propagate a rewritten +// request body to the wire. func (h *Holder) WritesRequestBody() bool { return h.p.Load().WritesRequestBody() } +// WritesResponseBody is equivalent to h.Load().WritesResponseBody(). +// Listeners read this when deciding whether streaming responses are safe +// — a pipeline with a response mutator can't stream, because the proxy +// can't rewrite a body it has already started forwarding. A request-only +// mutator does not disable streaming. +func (h *Holder) WritesResponseBody() bool { return h.p.Load().WritesResponseBody() } + // Ready is equivalent to h.Load().Ready(). func (h *Holder) Ready() bool { return h.p.Load().Ready() } diff --git a/authbridge/authlib/pipeline/pipeline.go b/authbridge/authlib/pipeline/pipeline.go index 133215c95..496b41fd8 100644 --- a/authbridge/authlib/pipeline/pipeline.go +++ b/authbridge/authlib/pipeline/pipeline.go @@ -368,7 +368,7 @@ func (p *Pipeline) NotReadyPlugin() string { func (p *Pipeline) NeedsBody() bool { for _, plugin := range p.plugins { caps := plugin.Capabilities().Normalize() - if caps.ReadsBody || caps.WritesRequestBody { + if caps.ReadsBody || caps.WritesRequestBody || caps.WritesResponseBody { return true } } @@ -389,6 +389,23 @@ func (p *Pipeline) WritesRequestBody() bool { return false } +// WritesResponseBody returns true if any plugin in the pipeline declares +// WritesResponseBody. This is the SSE streaming predicate: a response +// mutator needs the whole response to rewrite it, so listeners fall back +// from incremental relay to the buffered path only when this is true. +// +// A request-only mutator (tool-prune, context-guru) keeps streaming: the +// request body is already complete before dispatch, so rewriting it has +// no bearing on how the response is relayed. +func (p *Pipeline) WritesResponseBody() bool { + for _, plugin := range p.plugins { + if plugin.Capabilities().Normalize().WritesResponseBody { + return true + } + } + return false +} + // Start invokes Init on every plugin that implements the Initializer // interface, in declaration order. Returns the first error encountered; // on error, later plugins are not initialized. Plugins without Init are @@ -553,21 +570,39 @@ func (p *Pipeline) dispatchFinish(parent context.Context, name string, f Finishe // - A body reader (ReadsBody) must not follow a body mutator (WritesRequestBody) — // the reader would silently see mutated bytes instead of the originals. func validateCapabilities(plugins []Plugin) error { - var mutatorName string - var readerAfterMutator string + // Each direction admits at most one mutator. The rules are per-direction + // because ordering is only ambiguous between two plugins rewriting the + // same bytes; a request mutator and a response mutator never collide. + var requestMutator, responseMutator string + var firstMutator, readerAfterMutator string for _, plugin := range plugins { caps := plugin.Capabilities().Normalize() if caps.WritesRequestBody { - if mutatorName != "" { - return fmt.Errorf("pipeline: two plugins declare WritesRequestBody: %q and %q — mutation ordering would be ambiguous; at most one body mutator per pipeline is allowed", mutatorName, plugin.Name()) + if requestMutator != "" { + return fmt.Errorf("pipeline: two plugins declare WritesRequestBody: %q and %q — mutation ordering would be ambiguous; at most one request-body mutator per pipeline is allowed", requestMutator, plugin.Name()) + } + requestMutator = plugin.Name() + } + if caps.WritesResponseBody { + if responseMutator != "" { + return fmt.Errorf("pipeline: two plugins declare WritesResponseBody: %q and %q — mutation ordering would be ambiguous; at most one response-body mutator per pipeline is allowed", responseMutator, plugin.Name()) } - mutatorName = plugin.Name() - } else if caps.ReadsBody && mutatorName != "" && readerAfterMutator == "" { + responseMutator = plugin.Name() + } + if caps.WritesRequestBody || caps.WritesResponseBody { + if firstMutator == "" { + firstMutator = plugin.Name() + } + continue + } + // Reader-ordering is triggered by either write flag: a reader placed + // after any mutator would no longer see the original bytes. + if caps.ReadsBody && firstMutator != "" && readerAfterMutator == "" { readerAfterMutator = plugin.Name() } } if readerAfterMutator != "" { - return fmt.Errorf("pipeline: plugin %q reads body after mutator %q — body readers must precede the mutator so they see the original bytes", readerAfterMutator, mutatorName) + return fmt.Errorf("pipeline: plugin %q reads body after mutator %q — body readers must precede the mutator so they see the original bytes", readerAfterMutator, firstMutator) } return nil } diff --git a/authbridge/authlib/pipeline/plugin.go b/authbridge/authlib/pipeline/plugin.go index a8d088498..a924c6896 100644 --- a/authbridge/authlib/pipeline/plugin.go +++ b/authbridge/authlib/pipeline/plugin.go @@ -24,19 +24,36 @@ type PluginCapabilities struct { // a read silently sees "no body." ReadsBody bool - // WritesRequestBody: the plugin may mutate pctx.Body / pctx.ResponseBody - // (call pctx.SetBody / pctx.SetResponseBody). Implies ReadsBody — - // Normalize() auto-promotes. Listener propagates the mutation to - // the wire (ext_proc BodyMutation, or the outbound http.Request / - // downstream http.Response for proxy listeners). + // WritesRequestBody: the plugin may mutate pctx.Body (call + // pctx.SetBody). Implies ReadsBody — Normalize() auto-promotes. + // Listener propagates the mutation to the wire (ext_proc + // BodyMutation, or the outbound http.Request for proxy listeners). // - // Pipeline.New rejects a pipeline that has more than one WritesRequestBody - // plugin per direction — mutation ordering would be ambiguous. - // Waypoint mode (ext_authz) cannot support WritesRequestBody at all: - // ext_authz has no body-mutation field. main.go enforces this at - // process boot. + // Pipeline.New rejects a pipeline that has more than one + // WritesRequestBody plugin per direction — mutation ordering would + // be ambiguous. Waypoint mode (ext_authz) cannot support body + // mutation at all: ext_authz has no body-mutation field. main.go + // enforces this at process boot. + // + // Declaring this does NOT cost response streaming. Requests are + // never streamed — they arrive complete with a Content-Length and + // are read end to end before dispatch — so rewriting one says + // nothing about whether the response may be relayed incrementally. WritesRequestBody bool + // WritesResponseBody: the plugin may mutate pctx.ResponseBody (call + // pctx.SetResponseBody). Implies ReadsBody — Normalize() auto-promotes. + // + // This is the streaming predicate. A plugin that rewrites a response + // needs the whole response to rewrite it, so listeners fall back from + // incremental SSE relay to the buffered path when — and only when — + // some plugin in the chain declares this. See + // Pipeline.WritesResponseBody. + // + // Pipeline.New rejects more than one WritesResponseBody plugin per + // direction, for the same ordering reason as the request side. + WritesResponseBody bool + // Requires names plugins that MUST be present in the same chain // AND appear earlier (lower index). Matches are case-sensitive // plugin Name() strings. A missing or misordered name causes @@ -76,7 +93,7 @@ type PluginCapabilities struct { // rest of the framework reads a normalized form. Plugins never need to // call this themselves. func (c PluginCapabilities) Normalize() PluginCapabilities { - if c.WritesRequestBody { + if c.WritesRequestBody || c.WritesResponseBody { c.ReadsBody = true } return c diff --git a/authbridge/authlib/plugins/cpex/plugin.go b/authbridge/authlib/plugins/cpex/plugin.go index da50bed35..1f4987d2b 100644 --- a/authbridge/authlib/plugins/cpex/plugin.go +++ b/authbridge/authlib/plugins/cpex/plugin.go @@ -118,10 +118,11 @@ func (p *CPEX) Name() string { return "cpex" } // surfaces in the catalog. func (p *CPEX) Capabilities() pipeline.PluginCapabilities { return pipeline.PluginCapabilities{ - ReadsBody: true, - WritesRequestBody: true, - RequiresAny: []string{"mcp-parser", "inference-parser", "a2a-parser"}, - Description: "CPEX bridge: APL DSL + named CPEX plugins (Cedar, PII, audit, …) over a single chain step.", + ReadsBody: true, + WritesRequestBody: true, + WritesResponseBody: true, // cmf_body / cmf_a2a / cmf_inference rewrite responses + RequiresAny: []string{"mcp-parser", "inference-parser", "a2a-parser"}, + Description: "CPEX bridge: APL DSL + named CPEX plugins (Cedar, PII, audit, …) over a single chain step.", } } diff --git a/authbridge/authlib/plugins/registry.go b/authbridge/authlib/plugins/registry.go index 96abe1e90..8d209717b 100644 --- a/authbridge/authlib/plugins/registry.go +++ b/authbridge/authlib/plugins/registry.go @@ -205,17 +205,15 @@ func cloneCatalog(in []CatalogEntry) []CatalogEntry { } out := make([]CatalogEntry, len(in)) for i := range in { + // Struct copy, then reallocate the slices. Copying field-by-field + // silently drops any capability added later; this picks them up. caps := in[i].Capabilities + caps.Requires = append([]string(nil), in[i].Capabilities.Requires...) + caps.RequiresAny = append([]string(nil), in[i].Capabilities.RequiresAny...) out[i] = CatalogEntry{ - Name: in[i].Name, - Capabilities: pipeline.PluginCapabilities{ - ReadsBody: caps.ReadsBody, - WritesRequestBody: caps.WritesRequestBody, - Description: caps.Description, - Requires: append([]string(nil), caps.Requires...), - RequiresAny: append([]string(nil), caps.RequiresAny...), - }, - Fields: cloneFieldSchemas(in[i].Fields), + Name: in[i].Name, + Capabilities: caps, + Fields: cloneFieldSchemas(in[i].Fields), } } return out diff --git a/authbridge/authlib/plugins/registry_capsclone_test.go b/authbridge/authlib/plugins/registry_capsclone_test.go new file mode 100644 index 000000000..e0011577e --- /dev/null +++ b/authbridge/authlib/plugins/registry_capsclone_test.go @@ -0,0 +1,84 @@ +package plugins + +import ( + "reflect" + "testing" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" +) + +// nonZeroCaps fills every field of PluginCapabilities with a non-zero value, +// driven by reflection rather than a hand-written literal. A field added to +// the struct with a kind this helper does not handle fails the test loudly, +// which is the point: it is impossible to add a capability and forget it here. +func nonZeroCaps(t *testing.T) pipeline.PluginCapabilities { + t.Helper() + var c pipeline.PluginCapabilities + v := reflect.ValueOf(&c).Elem() + for i := 0; i < v.NumField(); i++ { + name := v.Type().Field(i).Name + f := v.Field(i) + switch f.Kind() { + case reflect.Bool: + f.SetBool(true) + case reflect.String: + f.SetString("value-" + name) + case reflect.Slice: + if f.Type().Elem().Kind() != reflect.String { + t.Fatalf("PluginCapabilities.%s is a slice of %s — extend nonZeroCaps", name, f.Type().Elem().Kind()) + } + f.Set(reflect.ValueOf([]string{"elem-" + name})) + default: + t.Fatalf("PluginCapabilities.%s has unhandled kind %s — extend nonZeroCaps", name, f.Kind()) + } + } + return c +} + +// TestCloneCatalog_PreservesEveryCapabilityField is the regression test for the +// field-by-field copy that cloneCatalog used to do: it silently dropped any +// capability added later, so /v1/plugins under-reported. A struct copy picks +// new fields up automatically, and this test proves it for every field the +// struct has — now and after future additions. +func TestCloneCatalog_PreservesEveryCapabilityField(t *testing.T) { + caps := nonZeroCaps(t) + in := []CatalogEntry{{ + Name: "probe", + Capabilities: caps, + Fields: []pipeline.FieldSchema{{Name: "f"}}, + }} + + out := cloneCatalog(in) + if len(out) != 1 { + t.Fatalf("cloneCatalog returned %d entries, want 1", len(out)) + } + if !reflect.DeepEqual(out[0].Capabilities, caps) { + t.Errorf("capabilities not round-tripped:\n got %+v\nwant %+v", out[0].Capabilities, caps) + } + if out[0].Name != "probe" { + t.Errorf("Name = %q, want %q", out[0].Name, "probe") + } +} + +// TestCloneCatalog_DeepCopiesSlices: the clone must not alias the caller's +// slices, or a mutation through /v1/plugins would reach into the registry. +func TestCloneCatalog_DeepCopiesSlices(t *testing.T) { + in := []CatalogEntry{{ + Name: "probe", + Capabilities: pipeline.PluginCapabilities{ + Requires: []string{"a"}, + RequiresAny: []string{"b"}, + }, + }} + out := cloneCatalog(in) + + out[0].Capabilities.Requires[0] = "mutated" + out[0].Capabilities.RequiresAny[0] = "mutated" + + if in[0].Capabilities.Requires[0] != "a" { + t.Error("Requires aliases the input slice") + } + if in[0].Capabilities.RequiresAny[0] != "b" { + t.Error("RequiresAny aliases the input slice") + } +} diff --git a/authbridge/authlib/plugins/sparc/plugin.go b/authbridge/authlib/plugins/sparc/plugin.go index 2c2e1f883..7d0d72618 100644 --- a/authbridge/authlib/plugins/sparc/plugin.go +++ b/authbridge/authlib/plugins/sparc/plugin.go @@ -209,10 +209,11 @@ func (p *SPARC) Capabilities() pipeline.PluginCapabilities { // conversation + tool specs (both modes); mcp-parser provides the tool // call (mcp mode). RequiresAny is a static "at least one" check; the // per-mode runtime requirements are validated/handled below. - RequiresAny: []string{"inference-parser", "mcp-parser"}, - ReadsBody: true, - WritesRequestBody: true, // MCP result (mcp mode) / completion rewrite (inference mode) - Description: "SPARC pre-tool reflection: blocks ungrounded/hallucinated tool calls.", + RequiresAny: []string{"inference-parser", "mcp-parser"}, + ReadsBody: true, + WritesRequestBody: true, // MCP result (mcp mode) / completion rewrite (inference mode) + WritesResponseBody: true, // respond.go rewrites the upstream response + Description: "SPARC pre-tool reflection: blocks ungrounded/hallucinated tool calls.", } } From 90bd2fa184fd57bd8367710dcd5f1870c5d3b90b Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 11:26:09 -0400 Subject: [PATCH 05/28] feat(authbridge): Add a plugin metrics channel surfaced in abctl MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Plugins had no way to report operator-facing counters. plugins.StatsSource and auth.Stats exist but are auth-shaped — typed approval and denial enums plus a custom MarshalJSON — so carrying something like "bytes removed" through them would distort their meaning. Add pipeline.Metric and pipeline.MetricsProvider: a plugin that implements Metrics() has its counters picked up by describePipeline and rendered by abctl. Named interfaces rather than inline literals at the call site, for the same reason as RawConfigProvider — a greppable contract, and signature drift becomes a compile error instead of a silently-failing assertion. - authlib/pipeline/metrics.go: Metric{Name,Value,Unit,Note} and the MetricsProvider interface. Value is float64 so a ratio or per-request average needs no second type. Note carries the caveat a derived number needs to be read honestly — above all the sample size behind an estimate. - sessionapi: Metrics on the pipeline plugin view, populated by a type assertion mirroring the RawConfigProvider case. Omitted from the payload when a plugin reports none, so a consumer can tell "no such channel" from "channel with nothing in it". - configuredPlugin forwards Metrics() explicitly. Go does not promote optional interfaces through the wrapper's embedded Plugin, which is why Initializer/Shutdowner/Finisher/Readier are each forwarded by hand — MetricsProvider needs the same treatment. Without it, metrics are invisible for every plugin that HAS config, which is every plugin an operator actually configures. Unlike StreamingResponder this can be forwarded unconditionally: no dispatch path selects on it, and a non-provider returns nil, which omitempty drops. - abctl: PluginMetric mirrors the wire type locally (the PluginFieldEntry convention, with a decode test guarding the tags), and the plugin detail pane grows a Metrics section between the dependency sections and Config. Values right-align into one column so they can be compared by eye; Note renders in styleHint. The section header is drawn even when empty, following the deliberate always-newline convention that keeps the layout from shifting as you navigate between plugins. Tests cover the unwrapped provider, the wrapped provider, and both non-provider cases — the wrapped path specifically, because a test using an unconfigured stub cannot see the forwarding bug at all. Generic on purpose: it lands before any plugin uses it, so it is reviewed on its own merits rather than as scaffolding for one caller. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- authbridge/authlib/pipeline/configured.go | 19 ++ authbridge/authlib/pipeline/metrics.go | 34 ++++ authbridge/authlib/sessionapi/metrics_test.go | 170 ++++++++++++++++++ authbridge/authlib/sessionapi/server.go | 7 + authbridge/cmd/abctl/apiclient/client.go | 12 ++ .../abctl/apiclient/metrics_decode_test.go | 70 ++++++++ .../cmd/abctl/tui/plugin_detail_pane.go | 13 ++ authbridge/cmd/abctl/tui/plugin_metrics.go | 60 +++++++ .../cmd/abctl/tui/plugin_metrics_test.go | 77 ++++++++ 9 files changed, 462 insertions(+) create mode 100644 authbridge/authlib/pipeline/metrics.go create mode 100644 authbridge/authlib/sessionapi/metrics_test.go create mode 100644 authbridge/cmd/abctl/apiclient/metrics_decode_test.go create mode 100644 authbridge/cmd/abctl/tui/plugin_metrics.go create mode 100644 authbridge/cmd/abctl/tui/plugin_metrics_test.go diff --git a/authbridge/authlib/pipeline/configured.go b/authbridge/authlib/pipeline/configured.go index 5bcdbd3c4..2ac94a7b6 100644 --- a/authbridge/authlib/pipeline/configured.go +++ b/authbridge/authlib/pipeline/configured.go @@ -133,6 +133,25 @@ func (c *configuredPlugin) OnFinish(ctx context.Context, pctx *Context) { } } +// Metrics forwards to the wrapped plugin if it implements MetricsProvider; +// otherwise returns nil. Required for the same reason as the four above: +// MetricsProvider is an optional interface, so it is not promoted through the +// embedded Plugin, and without this every Configurable plugin — which is to +// say every plugin an operator actually configures — would silently report no +// metrics at all on /v1/pipeline. +// +// This does make every wrapped plugin satisfy MetricsProvider, but unlike +// StreamingResponder that costs nothing: no dispatch path selects on it, and +// a non-provider returns nil, which the session API omits. "No such channel" +// and "channel with nothing in it" therefore still look identical on the wire, +// which is what abctl renders as "(none)". +func (c *configuredPlugin) Metrics() []Metric { + if mp, ok := c.Plugin.(MetricsProvider); ok { + return mp.Metrics() + } + return nil +} + // Ready forwards to the wrapped plugin if it implements Readier; otherwise // returns true. This matches the existing semantics in Pipeline.Ready // (pipeline.go:287-289): plugins without Readier are considered always-ready. diff --git a/authbridge/authlib/pipeline/metrics.go b/authbridge/authlib/pipeline/metrics.go new file mode 100644 index 000000000..1fd5322d5 --- /dev/null +++ b/authbridge/authlib/pipeline/metrics.go @@ -0,0 +1,34 @@ +package pipeline + +// Metric is one operator-facing counter reported by a plugin. Values are +// float64 so a plugin can report a ratio or a per-request average without a +// second type; counts are whole numbers that happen to fit exactly. +// +// Unit is advisory and drives display, not arithmetic: "count", "bytes", +// "tokens", "ratio". Note carries the caveat a number needs to be read +// honestly — most importantly the sample size behind an estimate, e.g. +// "estimate, n=1284". A derived figure with no Note is read as measured, so +// anything inferred must say so here. +type Metric struct { + Name string `json:"name"` + Value float64 `json:"value"` + Unit string `json:"unit,omitempty"` + Note string `json:"note,omitempty"` +} + +// MetricsProvider is implemented by plugins that expose counters for operator +// display. describePipeline calls it on demand while serving /v1/pipeline, so +// implementations must be safe for concurrent use with the request path and +// must not block — take a mutex, copy, release. Returning nil is fine and +// renders as "(none)". +// +// This is deliberately separate from plugins.StatsSource / auth.Stats, which +// are auth-shaped: they carry typed approval and denial enums and a custom +// MarshalJSON, so routing "bytes removed" through them would distort their +// meaning. Naming the interface here (rather than asserting an inline literal +// at the call site) gives callers a greppable contract and turns future +// signature drift into a compile error rather than a silently-failing +// type assertion — the same reasoning as RawConfigProvider. +type MetricsProvider interface { + Metrics() []Metric +} diff --git a/authbridge/authlib/sessionapi/metrics_test.go b/authbridge/authlib/sessionapi/metrics_test.go new file mode 100644 index 000000000..adb132590 --- /dev/null +++ b/authbridge/authlib/sessionapi/metrics_test.go @@ -0,0 +1,170 @@ +package sessionapi + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" + "github.com/rossoctl/cortex/authbridge/authlib/session" +) + +// meteredPlugin implements pipeline.MetricsProvider on top of fakePlugin's shape. +type meteredPlugin struct { + name string + metrics []pipeline.Metric +} + +func (m *meteredPlugin) Name() string { return m.name } +func (m *meteredPlugin) Capabilities() pipeline.PluginCapabilities { + return pipeline.PluginCapabilities{} +} +func (m *meteredPlugin) OnRequest(_ context.Context, _ *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +func (m *meteredPlugin) OnResponse(_ context.Context, _ *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +func (m *meteredPlugin) Metrics() []pipeline.Metric { return m.metrics } + +func pipelineJSON(t *testing.T, outbound []pipeline.Plugin) (string, []pipelinePluginView) { + t.Helper() + pipe, err := pipeline.New(outbound) + if err != nil { + t.Fatalf("pipeline.New: %v", err) + } + store := session.New(5*time.Minute, 100, 0) + defer store.Close() + srv := New(":0", store, WithPipelines(nil, pipeline.NewHolder(pipe))) + ts := httptest.NewServer(srv.server.Handler) + defer ts.Close() + + resp, err := http.Get(ts.URL + "/v1/pipeline") + if err != nil { + t.Fatalf("Get: %v", err) + } + defer resp.Body.Close() + raw, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("ReadAll: %v", err) + } + var body struct { + Outbound []pipelinePluginView `json:"outbound"` + } + if err := json.Unmarshal(raw, &body); err != nil { + t.Fatalf("Unmarshal: %v — raw=%s", err, raw) + } + return string(raw), body.Outbound +} + +// TestPipelineView_OmitsMetricsForNonProviders: a plugin that does not +// implement MetricsProvider must not emit a metrics key at all. abctl relies +// on absence-vs-empty to render "(none)" rather than an empty table, and an +// always-present null would also churn every existing golden payload. +func TestPipelineView_OmitsMetricsForNonProviders(t *testing.T) { + raw, views := pipelineJSON(t, []pipeline.Plugin{&fakePlugin{name: "token-exchange"}}) + if len(views) != 1 { + t.Fatalf("got %d plugins, want 1", len(views)) + } + if views[0].Metrics != nil { + t.Errorf("Metrics = %v, want nil for a non-provider", views[0].Metrics) + } + if strings.Contains(raw, "metrics") { + t.Errorf("payload should not mention metrics at all:\n%s", raw) + } +} + +// TestPipelineView_CarriesProviderMetrics: values, units and notes survive the +// round trip, including the Note that labels an estimate as one. +func TestPipelineView_CarriesProviderMetrics(t *testing.T) { + want := []pipeline.Metric{ + {Name: "requests seen", Value: 1284, Unit: "count"}, + {Name: "bytes removed", Value: 9389184, Unit: "bytes"}, + {Name: "tokens saved / request", Value: 1830.5, Unit: "tokens", Note: "estimate, n=1284"}, + } + _, views := pipelineJSON(t, []pipeline.Plugin{&meteredPlugin{name: "tool-prune", metrics: want}}) + if len(views) != 1 { + t.Fatalf("got %d plugins, want 1", len(views)) + } + got := views[0].Metrics + if len(got) != len(want) { + t.Fatalf("got %d metrics, want %d: %+v", len(got), len(want), got) + } + for i := range want { + if got[i] != want[i] { + t.Errorf("metric %d = %+v, want %+v", i, got[i], want[i]) + } + } +} + +// TestPipelineView_ProviderReturningNilOmitsKey: a provider that has nothing +// to report yet behaves like a non-provider on the wire, so a freshly started +// plugin doesn't render an empty table. +func TestPipelineView_ProviderReturningNilOmitsKey(t *testing.T) { + _, views := pipelineJSON(t, []pipeline.Plugin{&meteredPlugin{name: "tool-prune", metrics: nil}}) + if len(views) != 1 { + t.Fatalf("got %d plugins, want 1", len(views)) + } + if views[0].Metrics != nil { + t.Errorf("Metrics = %v, want nil", views[0].Metrics) + } +} + +// TestPipelineView_CarriesMetricsThroughConfiguredWrapper is the regression test +// for a bug the tests above could not see. A plugin that has config is wrapped +// by pipeline.WrapConfigured, and Go does not promote optional interfaces +// through the wrapper's embedded Plugin — so MetricsProvider has to be +// forwarded explicitly, exactly as Initializer/Shutdowner/Finisher/Readier are. +// +// Every plugin an operator actually configures takes this path, so before the +// forwarding existed, metrics were invisible in every real deployment while the +// unconfigured case above passed happily. End-to-end verification caught it; +// this test keeps it caught. +func TestPipelineView_CarriesMetricsThroughConfiguredWrapper(t *testing.T) { + want := []pipeline.Metric{ + {Name: "requests pruned", Value: 3, Unit: "count"}, + {Name: "bytes removed", Value: 825, Unit: "bytes"}, + } + inner := &meteredPlugin{name: "tool-prune", metrics: want} + wrapped := pipeline.WrapConfigured(inner, json.RawMessage(`{"remove":["NotebookEdit"]}`)) + + _, views := pipelineJSON(t, []pipeline.Plugin{wrapped}) + if len(views) != 1 { + t.Fatalf("got %d plugins, want 1", len(views)) + } + got := views[0].Metrics + if len(got) != len(want) { + t.Fatalf("got %d metrics through the wrapper, want %d — MetricsProvider is not being forwarded", len(got), len(want)) + } + for i := range want { + if got[i] != want[i] { + t.Errorf("metric %d = %+v, want %+v", i, got[i], want[i]) + } + } + // The wrapper must still surface config, so the two channels coexist. + if len(views[0].Config) == 0 { + t.Error("wrapped plugin lost its config") + } +} + +// TestPipelineView_WrappedNonProviderStillOmitsMetrics: forwarding makes every +// wrapped plugin satisfy MetricsProvider, so confirm that does not turn into an +// empty metrics table for plugins that report nothing. +func TestPipelineView_WrappedNonProviderStillOmitsMetrics(t *testing.T) { + wrapped := pipeline.WrapConfigured(&fakePlugin{name: "token-exchange"}, json.RawMessage(`{"a":1}`)) + raw, views := pipelineJSON(t, []pipeline.Plugin{wrapped}) + if len(views) != 1 { + t.Fatalf("got %d plugins, want 1", len(views)) + } + if views[0].Metrics != nil { + t.Errorf("Metrics = %v, want nil for a wrapped non-provider", views[0].Metrics) + } + if strings.Contains(raw, "metrics") { + t.Errorf("payload should omit metrics entirely:\n%s", raw) + } +} diff --git a/authbridge/authlib/sessionapi/server.go b/authbridge/authlib/sessionapi/server.go index 1f45aad2a..ccba81017 100644 --- a/authbridge/authlib/sessionapi/server.go +++ b/authbridge/authlib/sessionapi/server.go @@ -169,6 +169,10 @@ type pipelinePluginView struct { RequiresAny []string `json:"requiresAny,omitempty"` Description string `json:"description,omitempty"` Config json.RawMessage `json:"config,omitempty"` + // Metrics is populated for plugins implementing pipeline.MetricsProvider. + // Omitted entirely when a plugin reports none, so abctl can distinguish + // "no such channel" from "channel with nothing in it". + Metrics []pipeline.Metric `json:"metrics,omitempty"` } // handlePipeline returns the composition of the inbound and outbound @@ -234,6 +238,9 @@ func describePipeline(h *pipeline.Holder, direction string) []pipelinePluginView if rc, ok := pl.(pipeline.RawConfigProvider); ok { view.Config = redact.JSON(rc.RawConfig()) } + if mp, ok := pl.(pipeline.MetricsProvider); ok { + view.Metrics = mp.Metrics() + } out[i] = view } return out diff --git a/authbridge/cmd/abctl/apiclient/client.go b/authbridge/cmd/abctl/apiclient/client.go index 15a4e2be0..85f7f36ed 100644 --- a/authbridge/cmd/abctl/apiclient/client.go +++ b/authbridge/cmd/abctl/apiclient/client.go @@ -91,6 +91,18 @@ type PipelinePlugin struct { RequiresAny []string `json:"requiresAny,omitempty"` Description string `json:"description,omitempty"` Config json.RawMessage `json:"config,omitempty"` + Metrics []PluginMetric `json:"metrics,omitempty"` +} + +// PluginMetric mirrors authlib/pipeline.Metric on the wire. Kept as a local +// type rather than importing the server struct, matching PluginFieldEntry: +// the client owns its decode shape, and a decode test guards the tags +// against drift. +type PluginMetric struct { + Name string `json:"name"` + Value float64 `json:"value"` + Unit string `json:"unit,omitempty"` + Note string `json:"note,omitempty"` } // GetPipeline fetches /v1/pipeline. diff --git a/authbridge/cmd/abctl/apiclient/metrics_decode_test.go b/authbridge/cmd/abctl/apiclient/metrics_decode_test.go new file mode 100644 index 000000000..56a13cf0d --- /dev/null +++ b/authbridge/cmd/abctl/apiclient/metrics_decode_test.go @@ -0,0 +1,70 @@ +package apiclient + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" +) + +// TestGetPipeline_DecodesPluginMetrics guards against tag drift between +// server-side pipelinePluginView.Metrics (authlib/sessionapi/server.go) and +// client-side PluginMetric here. The payload below is the exact shape the +// server emits; if a key stops decoding, the abctl metrics pane silently +// renders zeros, which is worse than rendering nothing. +func TestGetPipeline_DecodesPluginMetrics(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/pipeline" { + http.NotFound(w, r) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "inbound": [], + "outbound": [ + { + "name": "tool-prune", + "direction": "outbound", + "position": 4, + "readsBody": true, + "metrics": [ + {"name": "requests seen", "value": 1284, "unit": "count"}, + {"name": "bytes removed", "value": 9389184, "unit": "bytes"}, + {"name": "tokens saved / request", "value": 1830.5, + "unit": "tokens", "note": "estimate, n=1284"} + ] + }, + {"name": "mcp-parser", "direction": "outbound", "position": 2} + ] + }`)) + })) + defer ts.Close() + + c := New(ts.URL) + view, err := c.GetPipeline(context.Background()) + if err != nil { + t.Fatalf("GetPipeline: %v", err) + } + if len(view.Outbound) != 2 { + t.Fatalf("got %d outbound plugins, want 2", len(view.Outbound)) + } + + got := view.Outbound[0].Metrics + if len(got) != 3 { + t.Fatalf("got %d metrics, want 3: %+v", len(got), got) + } + if got[0].Name != "requests seen" || got[0].Value != 1284 || got[0].Unit != "count" { + t.Errorf("metrics[0] = %+v", got[0]) + } + if got[2].Value != 1830.5 { + t.Errorf("metrics[2].Value = %v, want 1830.5 (fractional values must survive)", got[2].Value) + } + if got[2].Note != "estimate, n=1284" { + t.Errorf("metrics[2].Note = %q — the estimate caveat must decode", got[2].Note) + } + // A plugin with no metrics key decodes to nil, which the pane renders + // as "(none)" rather than an empty table. + if view.Outbound[1].Metrics != nil { + t.Errorf("mcp-parser Metrics = %+v, want nil", view.Outbound[1].Metrics) + } +} diff --git a/authbridge/cmd/abctl/tui/plugin_detail_pane.go b/authbridge/cmd/abctl/tui/plugin_detail_pane.go index 3bed7b03e..efd12ac2b 100644 --- a/authbridge/cmd/abctl/tui/plugin_detail_pane.go +++ b/authbridge/cmd/abctl/tui/plugin_detail_pane.go @@ -63,6 +63,19 @@ func (m *model) showPluginDetail(p *apiclient.PipelinePlugin) { b.WriteString("\n") } } + // Metrics section, for plugins implementing pipeline.MetricsProvider. + // Same always-newline treatment as Config below: the header is drawn + // whether or not there are rows, so navigating between a plugin that + // reports counters and one that does not doesn't shift the layout. + fmt.Fprintln(&b) + b.WriteString(styleMuted.Render("Metrics:")) + b.WriteString("\n") + if len(p.Metrics) == 0 { + b.WriteString(" (none)\n") + } else { + b.WriteString(formatPluginMetrics(p.Metrics)) + } + fmt.Fprintln(&b) // Always-newline format keeps the visual layout consistent whether // the plugin is Configurable (JSON body, multi-line) or not ("(none)", diff --git a/authbridge/cmd/abctl/tui/plugin_metrics.go b/authbridge/cmd/abctl/tui/plugin_metrics.go new file mode 100644 index 000000000..db2c78b17 --- /dev/null +++ b/authbridge/cmd/abctl/tui/plugin_metrics.go @@ -0,0 +1,60 @@ +package tui + +import ( + "fmt" + "math" + "strconv" + "strings" + + "github.com/rossoctl/cortex/authbridge/cmd/abctl/apiclient" +) + +// formatMetricValue renders a metric value without lying about precision. +// Counters and byte totals are whole numbers and print as integers; a derived +// figure (a ratio, a per-request average) keeps two decimals. Very large +// values fall back to %g rather than printing 20 digits of float noise. +func formatMetricValue(v float64) string { + switch { + case math.IsNaN(v) || math.IsInf(v, 0): + return "—" + case math.Abs(v) >= 1e15: + return strconv.FormatFloat(v, 'g', 6, 64) + case v == math.Trunc(v): + return strconv.FormatInt(int64(v), 10) + default: + return strconv.FormatFloat(v, 'f', 2, 64) + } +} + +// formatPluginMetrics lays out metric rows as name / right-aligned value / +// unit / note. Columns are sized to the widest entry so the numbers line up +// and can be compared by eye, which is the whole reason an operator opens +// this pane. Note renders in styleHint, as Description does in the header — +// it carries the caveat (sample size, "estimate") that keeps a derived +// number from being read as a measurement. +func formatPluginMetrics(metrics []apiclient.PluginMetric) string { + nameW, valW := 0, 0 + vals := make([]string, len(metrics)) + for i, m := range metrics { + vals[i] = formatMetricValue(m.Value) + if len(m.Name) > nameW { + nameW = len(m.Name) + } + if len(vals[i]) > valW { + valW = len(vals[i]) + } + } + + var b strings.Builder + for i, m := range metrics { + fmt.Fprintf(&b, " %-*s %*s", nameW, m.Name, valW, vals[i]) + if m.Unit != "" { + fmt.Fprintf(&b, " %s", styleMuted.Render(m.Unit)) + } + if m.Note != "" { + fmt.Fprintf(&b, " %s", styleHint.Render(m.Note)) + } + b.WriteString("\n") + } + return b.String() +} diff --git a/authbridge/cmd/abctl/tui/plugin_metrics_test.go b/authbridge/cmd/abctl/tui/plugin_metrics_test.go new file mode 100644 index 000000000..fef61710a --- /dev/null +++ b/authbridge/cmd/abctl/tui/plugin_metrics_test.go @@ -0,0 +1,77 @@ +package tui + +import ( + "strings" + "testing" + + "github.com/rossoctl/cortex/authbridge/cmd/abctl/apiclient" +) + +func TestFormatMetricValue(t *testing.T) { + tests := []struct { + in float64 + want string + }{ + {1284, "1284"}, // counter: no decimal noise + {9389184, "9389184"}, // byte total + {0, "0"}, // a fresh counter is still a number + {1830.5, "1830.50"}, // derived figure keeps precision + {0.126, "0.13"}, // ratio rounds up + {0.125, "0.12"}, // exact tie: strconv rounds half-to-even + {-3, "-3"}, // negative whole + } + for _, tc := range tests { + if got := formatMetricValue(tc.in); got != tc.want { + t.Errorf("formatMetricValue(%v) = %q, want %q", tc.in, got, tc.want) + } + } +} + +// TestFormatPluginMetrics_AlignsColumns: the point of the pane is comparing +// numbers by eye, so values must right-align into one column regardless of +// name length. +func TestFormatPluginMetrics_AlignsColumns(t *testing.T) { + out := formatPluginMetrics([]apiclient.PluginMetric{ + {Name: "requests seen", Value: 7, Unit: "count"}, + {Name: "bytes removed / request", Value: 7312, Unit: "bytes"}, + }) + lines := strings.Split(strings.TrimRight(out, "\n"), "\n") + if len(lines) != 2 { + t.Fatalf("got %d lines, want 2:\n%s", len(lines), out) + } + // Both value strings must end at the same column. + col0 := strings.Index(lines[0], "7") + col1 := strings.Index(lines[1], "7312") + if col0 < 0 || col1 < 0 { + t.Fatalf("values not found in output:\n%s", out) + } + if end0, end1 := col0+len("7"), col1+len("7312"); end0 != end1 { + t.Errorf("values not right-aligned: %q ends at %d, %q ends at %d\n%s", + "7", end0, "7312", end1, out) + } + for i, l := range lines { + if !strings.HasPrefix(l, " ") { + t.Errorf("line %d not indented: %q", i, l) + } + } +} + +// TestFormatPluginMetrics_RendersNote: a derived number without its caveat +// reads as a measurement. The note must appear on the row. +func TestFormatPluginMetrics_RendersNote(t *testing.T) { + out := formatPluginMetrics([]apiclient.PluginMetric{ + {Name: "tokens saved / request", Value: 1830, Unit: "tokens", Note: "estimate, n=1284"}, + }) + if !strings.Contains(out, "estimate, n=1284") { + t.Errorf("note missing from output: %q", out) + } + if !strings.Contains(out, "tokens") { + t.Errorf("unit missing from output: %q", out) + } +} + +func TestFormatPluginMetrics_EmptyIsEmptyString(t *testing.T) { + if got := formatPluginMetrics(nil); got != "" { + t.Errorf("formatPluginMetrics(nil) = %q, want empty (pane renders (none))", got) + } +} From 356809e229d97ce1b2dd3dbfba5ab90214fbab90 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 12:47:14 -0400 Subject: [PATCH 06/28] feat(authbridge): Add the tool-prune plugin and abctl tools scan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A Claude Code turn carries the full tool manifest on every request — tens of thousands of tokens of JSON schema, billed each time, largely for tools the agent will never call in a given deployment. The manifest is assembled by the client, so the proxy is the only place to trim it without changing every client. tool-prune deletes configured tool definitions from the outbound manifest. The verdict is entirely configuration: `remove` names the tools, and there is no learning, no state and no storage dependency. It declares WritesRequestBody only, so responses still stream incrementally — which is what the directional capability split in the earlier commit bought. Measure before enforcing. Under `on_error: observe` the plugin computes exactly what it would remove and counts it while the bytes on the wire stay untouched: SetBody is a no-op on bytes and leaves BodyMutated() false, which is how the plugin knows which counter to increment. One registration serves both modes, selected by one word of config, and the readout says "requests projected" rather than "requests pruned" so a projection is never mistaken for a realised saving. Safety is one-directional throughout — removing a tool the model needs is the harmful failure, carrying extra definitions is merely expensive: - Names are resolved from the raw request bytes, not from the parsed manifest, because inference-parser drops unnamed tools and manifest position therefore does not map back to array position. Covers both the Anthropic (tools.i.name) and OpenAI (tools.i.function.name) dialects. - Deletions run descending so an earlier one never shifts a later index. - Every byte outside the removed elements is preserved, key order and whitespace included. - Removing every tool drops the `tools` and `tool_choice` keys rather than leaving `tools: []`, which OpenAI rejects. - Input must be valid JSON and the result is re-validated for JSON validity and expected tool count before it reaches the wire. gjson parses leniently, so without the input guard a truncated body resolved `tools` and sjson rewrote the fragment down to `{` — caught in testing. - Malformed bodies, absent manifests, non-shrinking rewrites and panics all forward the original bytes. Metrics use the channel added in the previous commit: request counts, tools and bytes removed, per-tool attribution, and a tokens-saved estimate calibrated on the operator's own traffic from the response usage block via OnFinish, reported with its sample size rather than a bundled tokenizer or an assumed bytes-per-token constant. `abctl tools scan` derives a candidate list from ~/.claude/projects transcripts: literal prefilter before any JSON parsing, tool calls deduplicated by tool_use block id (a transcript is rewritten on every resume), windowed by --days. Transcripts record tools that were *called*, never tools that were *offered*, so the scan intersects "known Claude Code built-ins" with "never called" and keeps anything it does not recognise. An implies table covers indirect use (Agent implying SendMessage). --write patches the remove: list in place, line-based and idempotent, so the operator's comments and hand-tuned entries survive byte-for-byte. Shipped inert: present in the --demo pipeline with an empty remove list and on_error: observe, placed last because it is the body mutator. install-demo.sh offers the scan only for a config that already exists, so a first run never rewrites a file it just created. Excludable via exclude_plugin_toolprune, added to the authbridge-lite tag set in both workflows. Verified end to end against a running proxy, not only in unit tests: enforce prunes exactly the configured tools and nothing else, observe leaves the upstream bytes identical while still counting the projection, the all-removed path drops both keys, and the metrics readout appears on /v1/pipeline. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- .github/workflows/build.yaml | 2 +- .github/workflows/ci.yaml | 1 + CLAUDE.md | 2 +- authbridge/CLAUDE.md | 22 +- authbridge/authlib/go.mod | 4 +- .../authlib/plugins/toolprune/metrics.go | 159 +++++++ .../authlib/plugins/toolprune/plugin.go | 311 +++++++++++++ .../authlib/plugins/toolprune/plugin_test.go | 416 ++++++++++++++++++ authbridge/cmd/abctl/cmd_tools.go | 88 ++++ authbridge/cmd/abctl/main.go | 14 + authbridge/cmd/abctl/toolscan/known.go | 67 +++ authbridge/cmd/abctl/toolscan/patch.go | 82 ++++ authbridge/cmd/abctl/toolscan/patch_test.go | 175 ++++++++ authbridge/cmd/abctl/toolscan/scan.go | 185 ++++++++ authbridge/cmd/abctl/toolscan/scan_test.go | 208 +++++++++ .../cmd/authbridge-envoy/plugins_toolprune.go | 5 + authbridge/cmd/authbridge-proxy/demo.go | 18 +- authbridge/cmd/authbridge-proxy/demo_test.go | 26 +- .../cmd/authbridge-proxy/plugins_toolprune.go | 5 + authbridge/docs/framework-architecture.md | 20 +- authbridge/docs/plugin-catalog.md | 21 + authbridge/docs/plugin-reference.md | 86 +++- authbridge/docs/plugin-tutorial.md | 20 +- authbridge/docs/tool-prune-plugin.md | 165 +++++++ authbridge/install-demo.sh | 13 + 25 files changed, 2087 insertions(+), 28 deletions(-) create mode 100644 authbridge/authlib/plugins/toolprune/metrics.go create mode 100644 authbridge/authlib/plugins/toolprune/plugin.go create mode 100644 authbridge/authlib/plugins/toolprune/plugin_test.go create mode 100644 authbridge/cmd/abctl/cmd_tools.go create mode 100644 authbridge/cmd/abctl/toolscan/known.go create mode 100644 authbridge/cmd/abctl/toolscan/patch.go create mode 100644 authbridge/cmd/abctl/toolscan/patch_test.go create mode 100644 authbridge/cmd/abctl/toolscan/scan.go create mode 100644 authbridge/cmd/abctl/toolscan/scan_test.go create mode 100644 authbridge/cmd/authbridge-envoy/plugins_toolprune.go create mode 100644 authbridge/cmd/authbridge-proxy/plugins_toolprune.go create mode 100644 authbridge/docs/tool-prune-plugin.md diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index f712755f2..7fa2965aa 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -56,7 +56,7 @@ jobs: context: ./authbridge dockerfile: cmd/authbridge-proxy/Dockerfile build_args: | - GO_BUILD_TAGS=exclude_plugin_a2aparser,exclude_plugin_ibac,exclude_plugin_inferenceparser,exclude_plugin_mcpparser,exclude_plugin_opa,exclude_plugin_sparc,exclude_plugin_tokenbroker + GO_BUILD_TAGS=exclude_plugin_a2aparser,exclude_plugin_ibac,exclude_plugin_inferenceparser,exclude_plugin_mcpparser,exclude_plugin_opa,exclude_plugin_sparc,exclude_plugin_tokenbroker,exclude_plugin_toolprune # AuthBridge proxy-sidecar CPEX image — authbridge-proxy built # with -tags cpex (links libcpex_ffi.a from a pinned CPEX diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index ede476b26..f9d55baed 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -109,6 +109,7 @@ jobs: run: | TAGS="exclude_plugin_a2aparser,exclude_plugin_ibac,exclude_plugin_inferenceparser" TAGS="$TAGS,exclude_plugin_mcpparser,exclude_plugin_opa,exclude_plugin_sparc,exclude_plugin_tokenbroker" + TAGS="$TAGS,exclude_plugin_toolprune" go build -v -tags "$TAGS" ./... go test -v -race -cover -tags "$TAGS" ./... diff --git a/CLAUDE.md b/CLAUDE.md index 0fca75204..ff4dec73e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -254,7 +254,7 @@ cd authbridge && podman build -f cmd/authbridge-proxy/Dockerfile -t authbridge:l cd authbridge && podman build -f cmd/authbridge-envoy/Dockerfile -t authbridge-envoy:latest . # authbridge-lite: same proxy Dockerfile, built with exclude_plugin_* tags (auth-only) cd authbridge && podman build -f cmd/authbridge-proxy/Dockerfile \ - --build-arg GO_BUILD_TAGS="exclude_plugin_a2aparser,exclude_plugin_ibac,exclude_plugin_inferenceparser,exclude_plugin_mcpparser,exclude_plugin_opa,exclude_plugin_sparc,exclude_plugin_tokenbroker" \ + --build-arg GO_BUILD_TAGS="exclude_plugin_a2aparser,exclude_plugin_ibac,exclude_plugin_inferenceparser,exclude_plugin_mcpparser,exclude_plugin_opa,exclude_plugin_sparc,exclude_plugin_tokenbroker,exclude_plugin_toolprune" \ -t authbridge-lite:latest . ``` diff --git a/authbridge/CLAUDE.md b/authbridge/CLAUDE.md index d79a4bb8f..3fa90a671 100644 --- a/authbridge/CLAUDE.md +++ b/authbridge/CLAUDE.md @@ -11,7 +11,8 @@ binaries with shared auth logic in `authlib/`: - `cmd/authbridge-proxy/` — proxy-sidecar mode (default). HTTP forward + reverse proxies. Compiles in every plugin by default (jwt-validation, token-exchange, - a2a-parser, mcp-parser, inference-parser, opa, sparc, ibac, token-broker). + a2a-parser, mcp-parser, inference-parser, opa, sparc, ibac, token-broker, + tool-prune). **Every** plugin is excludable via `-tags exclude_plugin_` — one `plugins_.go` file per plugin, gated by `//go:build !exclude_plugin_`; `main.go` imports no plugin package directly. **Exception:** `context-guru` is @@ -154,6 +155,23 @@ wants to register. - `authlib/pipeline/` -- Plugin interface + lifecycle (`Configurable`, `Initializer`, `Shutdowner`); see [`docs/framework-architecture.md`](docs/framework-architecture.md) - `authlib/plugins/` -- The concrete plugins + registry; see [`docs/plugin-reference.md`](docs/plugin-reference.md) for the per-plugin config convention +**Directional body capabilities.** `PluginCapabilities` declares body writes +per direction: `WritesRequestBody` (calls `pctx.SetBody`) and +`WritesResponseBody` (calls `pctx.SetResponseBody`). `WritesResponseBody` is the +SSE streaming predicate — both proxy listeners fall back from incremental relay +to the buffered path only when some plugin declares it. A request-only mutator +(`tool-prune`, `context-guru`) therefore keeps streaming, because requests are +never streamed in the first place. `pipeline.New` allows at most one mutator per +direction, and no mutator of either direction may precede a `ReadsBody`-only +plugin. See [`docs/plugin-reference.md`](docs/plugin-reference.md#capability-fields). + +**Plugin metrics.** Plugins that implement `pipeline.MetricsProvider` have their +counters surfaced on `GET /v1/pipeline` and rendered in abctl's plugin pane. +Optional interfaces are not promoted through `configuredPlugin`'s embedded +`Plugin`, so a new one must be forwarded there explicitly or it is invisible for +every plugin that has config. Counters are per-process and reset on restart +**and on config hot-reload**. + **Plugin classification.** Protocol parsers (`mcp-parser`, `a2a-parser`, `inference-parser`) populate an `IsAction bool` field on their respective extensions to classify each request as either a user-meaningful action or protocol mechanics. Default-false means "not classified as action" — guardrails treat it as bypass. Parsers explicitly set `IsAction = true` for the small set of action methods (`tools/call` / `prompts/get` / `resources/read` for MCP; `message/send` / `message/stream` for A2A; every populated case for inference). Guardrails (`ibac` today; future rate limiters, audit loggers, etc.) read the aggregated verdict via `pctx.Classification()` which returns `(anyAction, anyBypass)`. A defense-in-depth guardrail skips on `anyBypass`, passes through on `!anyAction` (no parser claimed this traffic), and judges only when `anyAction && !anyBypass`. This puts the protocol-specific bypass-vs-action vocabulary in each parser — adding a new guardrail or new protocol does not multiply work at the guardrail layer. See [`docs/plugin-reference.md` "Classifying requests"](docs/plugin-reference.md#classifying-requests-as-actions-vs-protocol-mechanics) for the contract. ### init-iptables.sh @@ -377,7 +395,7 @@ podman build -f cmd/authbridge-proxy/Dockerfile -t authbridge:latest . # p podman build -f cmd/authbridge-envoy/Dockerfile -t authbridge-envoy:latest . # envoy-sidecar # authbridge-lite: the proxy Dockerfile built with exclude_plugin_* tags (auth-only) podman build -f cmd/authbridge-proxy/Dockerfile \ - --build-arg GO_BUILD_TAGS="exclude_plugin_a2aparser,exclude_plugin_ibac,exclude_plugin_inferenceparser,exclude_plugin_mcpparser,exclude_plugin_opa,exclude_plugin_sparc,exclude_plugin_tokenbroker" \ + --build-arg GO_BUILD_TAGS="exclude_plugin_a2aparser,exclude_plugin_ibac,exclude_plugin_inferenceparser,exclude_plugin_mcpparser,exclude_plugin_opa,exclude_plugin_sparc,exclude_plugin_tokenbroker,exclude_plugin_toolprune" \ -t authbridge-lite:latest . kind load docker-image authbridge:latest --name rossoctl kind load docker-image authbridge-envoy:latest --name rossoctl diff --git a/authbridge/authlib/go.mod b/authbridge/authlib/go.mod index d17b267cf..191d9a8f4 100644 --- a/authbridge/authlib/go.mod +++ b/authbridge/authlib/go.mod @@ -12,6 +12,8 @@ require ( github.com/open-policy-agent/opa v1.20.1 github.com/rossoctl/context-guru v0.1.0 github.com/spiffe/go-spiffe/v2 v2.8.1 + github.com/tidwall/gjson v1.18.0 + github.com/tidwall/sjson v1.2.5 golang.org/x/net v0.58.0 golang.org/x/sync v0.22.0 golang.org/x/sys v0.47.0 @@ -90,10 +92,8 @@ require ( github.com/spf13/cast v1.10.0 // indirect github.com/tchap/go-patricia/v2 v2.3.3 // indirect github.com/tetratelabs/wazero v1.12.0 // indirect - github.com/tidwall/gjson v1.18.0 // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.1 // indirect - github.com/tidwall/sjson v1.2.5 // indirect github.com/tiktoken-go/tokenizer v0.7.0 // indirect github.com/tree-sitter/go-tree-sitter v0.25.0 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect diff --git a/authbridge/authlib/plugins/toolprune/metrics.go b/authbridge/authlib/plugins/toolprune/metrics.go new file mode 100644 index 000000000..62bd2f5d0 --- /dev/null +++ b/authbridge/authlib/plugins/toolprune/metrics.go @@ -0,0 +1,159 @@ +package toolprune + +import ( + "fmt" + "sort" + "sync" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" +) + +// metrics holds the plugin's counters. In-memory and per-process by design: +// this targets the single-laptop case, and staying free of a storage backend is +// what keeps the plugin dependency-free. Counters reset on restart, which is +// why every derived figure is reported alongside the sample behind it. +type metrics struct { + mu sync.Mutex + + requestsSeen uint64 // matched the path gate and carried a manifest + requestsPruned uint64 // body actually rewritten (enforce) + requestsProjected uint64 // would have been rewritten (observe) + + toolsRemoved uint64 + perTool map[string]uint64 + + bytesRemoved uint64 + + // Calibration sample for bytes -> tokens, gathered from response usage. + promptTokens uint64 + requestBytes uint64 + requestsWithUsage uint64 +} + +func (m *metrics) seen() { + m.mu.Lock() + m.requestsSeen++ + m.mu.Unlock() +} + +func (m *metrics) pruned(names []string, bytesRemoved int) { + m.mu.Lock() + m.requestsPruned++ + m.record(names, bytesRemoved) + m.mu.Unlock() +} + +func (m *metrics) projected(names []string, bytesRemoved int) { + m.mu.Lock() + m.requestsProjected++ + m.record(names, bytesRemoved) + m.mu.Unlock() +} + +// record must be called with mu held. +func (m *metrics) record(names []string, bytesRemoved int) { + if m.perTool == nil { + m.perTool = make(map[string]uint64) + } + for _, n := range names { + m.perTool[n]++ + } + m.toolsRemoved += uint64(len(names)) + if bytesRemoved > 0 { + m.bytesRemoved += uint64(bytesRemoved) + } +} + +func (m *metrics) observeUsage(promptTokens, requestBytes int) { + m.mu.Lock() + m.promptTokens += uint64(promptTokens) + m.requestBytes += uint64(requestBytes) + m.requestsWithUsage++ + m.mu.Unlock() +} + +// snapshot renders the counters as operator-facing metrics. Every derived row +// carries the sample it was computed from, so a figure can never be read as +// more certain than it is. +func (m *metrics) snapshot() []pipeline.Metric { + m.mu.Lock() + defer m.mu.Unlock() + + if m.requestsSeen == 0 && m.requestsPruned == 0 && m.requestsProjected == 0 { + return nil + } + + out := []pipeline.Metric{ + {Name: "requests seen", Value: float64(m.requestsSeen), Unit: "count"}, + } + // Enforce and observe are mutually exclusive in practice (one policy per + // plugin instance), but report whichever has fired so a mid-flight policy + // change is visible rather than silently blended. + if m.requestsPruned > 0 || m.requestsProjected == 0 { + out = append(out, pipeline.Metric{ + Name: "requests pruned", Value: float64(m.requestsPruned), Unit: "count", + }) + } + if m.requestsProjected > 0 { + out = append(out, pipeline.Metric{ + Name: "requests projected", + Value: float64(m.requestsProjected), + Unit: "count", + Note: "observe mode — body unchanged", + }) + } + out = append(out, + pipeline.Metric{Name: "tools removed", Value: float64(m.toolsRemoved), Unit: "count"}, + pipeline.Metric{Name: "bytes removed", Value: float64(m.bytesRemoved), Unit: "bytes"}, + ) + + acted := m.requestsPruned + m.requestsProjected + if acted > 0 { + perReq := float64(m.bytesRemoved) / float64(acted) + out = append(out, pipeline.Metric{ + Name: "bytes removed / request", Value: perReq, Unit: "bytes", + }) + // Calibrate bytes -> tokens on the operator's own traffic instead of + // bundling a tokenizer or assuming a constant. With no usage sample + // yet, report zero rather than dividing by zero. + if m.requestBytes > 0 && m.promptTokens > 0 { + ratio := float64(m.promptTokens) / float64(m.requestBytes) + out = append(out, pipeline.Metric{ + Name: "tokens saved / request", + Value: perReq * ratio, + Unit: "tokens", + Note: fmt.Sprintf("estimate, n=%d", m.requestsWithUsage), + }) + } else { + out = append(out, pipeline.Metric{ + Name: "tokens saved / request", + Value: 0, + Unit: "tokens", + Note: "no usage sample yet", + }) + } + } + + // Per-tool attribution, sorted by count then name so the readout is + // stable across calls and the biggest contributors come first. + type kv struct { + name string + n uint64 + } + tools := make([]kv, 0, len(m.perTool)) + for k, v := range m.perTool { + tools = append(tools, kv{k, v}) + } + sort.Slice(tools, func(i, j int) bool { + if tools[i].n != tools[j].n { + return tools[i].n > tools[j].n + } + return tools[i].name < tools[j].name + }) + for _, t := range tools { + out = append(out, pipeline.Metric{ + Name: "removed: " + t.name, Value: float64(t.n), Unit: "count", + }) + } + return out +} diff --git a/authbridge/authlib/plugins/toolprune/plugin.go b/authbridge/authlib/plugins/toolprune/plugin.go new file mode 100644 index 000000000..89d4f1a1d --- /dev/null +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -0,0 +1,311 @@ +// Package toolprune removes unused tool definitions from outbound inference +// requests. +// +// A Claude Code request carries the full tool manifest on every turn — tens of +// thousands of tokens of JSON schema, billed each time and largely for tools +// the agent will never call in a given deployment. The manifest is assembled by +// the client, so the only place to trim it without touching every client is in +// the proxy. +// +// The verdict is entirely configuration: `remove` names the tools to drop. +// There is no learning, no state and no storage dependency. `abctl tools scan` +// produces a candidate list from local transcripts, but the plugin itself only +// ever does what it was told. +// +// Safety is one-directional. Removing a tool the model needs is the harmful +// failure; carrying a few extra definitions is not. So every error path fails +// open, forwarding the original bytes untouched: a cost optimisation must never +// be able to break a request. +package toolprune + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "log/slog" + "strings" + "sync" + + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" + "github.com/rossoctl/cortex/authbridge/authlib/plugins" +) + +// defaultPaths are the inference endpoints the plugin acts on, matched by +// suffix as context-guru does. +var defaultPaths = []string{"/v1/chat/completions", "/v1/completions", "/v1/messages"} + +type config struct { + // Remove names the tools to delete from the manifest. Names not present + // in a given request are ignored; names the plugin never observes are + // reported as drift rather than failing. + Remove []string `json:"remove" description:"Tool names to remove from the outbound manifest."` + + // Paths are the request paths this plugin acts on, matched exactly or by + // suffix. Defaults to the three inference endpoints. + Paths []string `json:"paths" description:"Request paths to act on (exact or suffix match)."` +} + +func (c *config) applyDefaults() { + if len(c.Paths) == 0 { + c.Paths = append([]string(nil), defaultPaths...) + } +} + +// ToolPrune is the plugin. Counters live in metrics, guarded by its own mutex; +// everything else is read-only after Configure. +type ToolPrune struct { + cfg config + raw json.RawMessage + remove map[string]struct{} + + m metrics + driftOnce sync.Once +} + +func New() *ToolPrune { return &ToolPrune{} } + +func init() { + plugins.RegisterPlugin("tool-prune", func() pipeline.Plugin { return New() }) +} + +func (p *ToolPrune) Name() string { return "tool-prune" } + +func (p *ToolPrune) Capabilities() pipeline.PluginCapabilities { + return pipeline.PluginCapabilities{ + // Request-only: the response is never touched, so SSE relay stays + // incremental. That distinction is the reason WritesResponseBody + // exists as a separate capability. + WritesRequestBody: true, + RequiresAny: []string{"inference-parser"}, + Description: "Removes unused tool definitions from inference requests.", + } +} + +// ConfigSchema implements pipeline.SchemaProvider. +func (p *ToolPrune) ConfigSchema() []pipeline.FieldSchema { + return pipeline.SchemaOf(config{}) +} + +// RawConfig implements pipeline.RawConfigProvider. +func (p *ToolPrune) RawConfig() json.RawMessage { return p.raw } + +func (p *ToolPrune) Configure(raw json.RawMessage) error { + var c config + if len(raw) > 0 { + dec := json.NewDecoder(bytes.NewReader(raw)) + dec.DisallowUnknownFields() + if err := dec.Decode(&c); err != nil { + return fmt.Errorf("tool-prune config: %w", err) + } + } + c.applyDefaults() + + p.cfg = c + p.raw = raw + p.remove = make(map[string]struct{}, len(c.Remove)) + for _, n := range c.Remove { + if n != "" { + p.remove[n] = struct{}{} + } + } + if len(p.remove) == 0 { + slog.Info("tool-prune: configured with an empty remove list — no-op until names are added", + "hint", "abctl tools scan") + } + return nil +} + +// gated reports whether the request path is one the plugin acts on. +func (p *ToolPrune) gated(path string) bool { + for _, s := range p.cfg.Paths { + if path == s || strings.HasSuffix(path, s) { + return true + } + } + return false +} + +// toolNameAt extracts a tool's name from raw manifest element i, covering both +// dialects: Anthropic puts it at tools.i.name, OpenAI at tools.i.function.name. +func toolNameAt(body []byte, i int) string { + if n := gjson.GetBytes(body, fmt.Sprintf("tools.%d.name", i)); n.Exists() { + return n.String() + } + return gjson.GetBytes(body, fmt.Sprintf("tools.%d.function.name", i)).String() +} + +// OnRequest prunes the manifest. Every failure path returns Continue with the +// body untouched. +func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action pipeline.Action) { + action = pipeline.Action{Type: pipeline.Continue} + if len(p.remove) == 0 { + return action + } + // A panic here would fail a request to save tokens. Never worth it. + defer func() { + if r := recover(); r != nil { + slog.Warn("tool-prune: recovered, forwarding original body", "panic", r) + action = pipeline.Action{Type: pipeline.Continue} + } + }() + + if !p.gated(pctx.Path) { + pctx.Record(pipeline.Invocation{Action: pipeline.ActionSkip, Reason: "path_not_inference"}) + return action + } + // inference-parser establishes that this is an inference call at all. Its + // absence means the chain is misconfigured; RequiresAny catches that at + // build time, so treat it as a skip rather than an error. + if pctx.Extensions.Inference == nil { + pctx.Record(pipeline.Invocation{Action: pipeline.ActionSkip, Reason: "no_inference_extension"}) + return action + } + body := pctx.Body + if len(body) == 0 { + pctx.Record(pipeline.Invocation{Action: pipeline.ActionSkip, Reason: "no_body"}) + return action + } + // gjson parses leniently: on a truncated document it still resolves + // fields, and sjson then rewrites the fragment into garbage. Refuse to + // touch anything that is not well-formed JSON to begin with. + if !gjson.ValidBytes(body) { + pctx.Record(pipeline.Invocation{Action: pipeline.ActionSkip, Reason: "invalid_json"}) + return action + } + tools := gjson.GetBytes(body, "tools") + if !tools.IsArray() { + pctx.Record(pipeline.Invocation{Action: pipeline.ActionSkip, Reason: "no_tool_manifest"}) + return action + } + + raw := tools.Array() + p.noteDrift(pctx.Extensions.Inference.Tools) + p.m.seen() + + // Resolve indices from the raw bytes rather than from the parsed manifest: + // inference-parser drops unnamed tools, so manifest position does not + // reliably map back to array position. + var victims []int + names := make([]string, 0, len(raw)) + for i := range raw { + name := toolNameAt(body, i) + if name == "" { + continue + } + if _, ok := p.remove[name]; ok { + victims = append(victims, i) + names = append(names, name) + } + } + if len(victims) == 0 { + pctx.Record(pipeline.Invocation{Action: pipeline.ActionSkip, Reason: "no_configured_tool_present"}) + return action + } + + out := body + var err error + if len(victims) == len(raw) { + // Emptying the array is not safe — OpenAI rejects `tools: []`, and + // tool_choice without tools. Drop both keys instead. + if out, err = sjson.DeleteBytes(out, "tools"); err != nil { + slog.Warn("tool-prune: delete tools failed, forwarding original", "err", err) + return action + } + if gjson.GetBytes(out, "tool_choice").Exists() { + if out, err = sjson.DeleteBytes(out, "tool_choice"); err != nil { + slog.Warn("tool-prune: delete tool_choice failed, forwarding original", "err", err) + return action + } + } + } else { + // Descending, so an earlier deletion never shifts a later index. + for i := len(victims) - 1; i >= 0; i-- { + if out, err = sjson.DeleteBytes(out, fmt.Sprintf("tools.%d", victims[i])); err != nil { + slog.Warn("tool-prune: delete failed, forwarding original", "index", victims[i], "err", err) + return action + } + } + } + if len(out) >= len(body) { + // Nothing shrank: treat as a no-op rather than emitting a rewrite. + pctx.Record(pipeline.Invocation{Action: pipeline.ActionSkip, Reason: "no_bytes_removed"}) + return action + } + // Post-conditions. The edit is surgical, so verify it actually did what + // was intended before putting it on the wire: still valid JSON, and + // exactly the intended number of tools left standing. + if !gjson.ValidBytes(out) { + slog.Warn("tool-prune: rewrite produced invalid JSON, forwarding original") + return action + } + want := len(raw) - len(victims) + if got := len(gjson.GetBytes(out, "tools").Array()); got != want { + slog.Warn("tool-prune: unexpected tool count after rewrite, forwarding original", + "got", got, "want", want) + return action + } + + removedBytes := len(body) - len(out) + pctx.SetBody(out) + // Under ErrorPolicyObserve, SetBody is a no-op on bytes and leaves + // bodyMutated false — so this same code path measures without enforcing, + // and the counter it lands in is what distinguishes the two. + if pctx.BodyMutated() { + p.m.pruned(names, removedBytes) + } else { + p.m.projected(names, removedBytes) + } + return action +} + +func (p *ToolPrune) OnResponse(_ context.Context, _ *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} + +// OnFinish calibrates the bytes-to-tokens ratio on the operator's own traffic, +// rather than bundling a tokenizer or hardcoding a constant. inference-parser +// is a StreamingResponder, so RunResponse skips its OnResponse — OnFinish is +// the hook where response-derived usage is reliably available. +func (p *ToolPrune) OnFinish(_ context.Context, pctx *pipeline.Context) { + if pctx.Extensions.Inference == nil { + return + } + prompt := pctx.Extensions.Inference.PromptTokens + if prompt <= 0 || len(pctx.Body) == 0 { + return + } + p.m.observeUsage(prompt, len(pctx.Body)) +} + +// noteDrift logs, once, any configured name absent from the first manifest the +// plugin actually sees. A stale list costs savings rather than correctness, so +// it surfaces as a warning instead of a failure. +func (p *ToolPrune) noteDrift(observed []pipeline.InferenceTool) { + p.driftOnce.Do(func() { + if len(observed) == 0 { + return + } + present := make(map[string]struct{}, len(observed)) + for _, t := range observed { + present[t.Name] = struct{}{} + } + var missing []string + for _, n := range p.cfg.Remove { + if _, ok := present[n]; !ok { + missing = append(missing, n) + } + } + if len(missing) > 0 { + slog.Warn("tool-prune: configured tools not present in the observed manifest — list may be stale", + "missing", strings.Join(missing, ","), + "hint", "re-run abctl tools scan") + } + }) +} + +// Metrics implements pipeline.MetricsProvider. +func (p *ToolPrune) Metrics() []pipeline.Metric { return p.m.snapshot() } diff --git a/authbridge/authlib/plugins/toolprune/plugin_test.go b/authbridge/authlib/plugins/toolprune/plugin_test.go new file mode 100644 index 000000000..33e23d3cb --- /dev/null +++ b/authbridge/authlib/plugins/toolprune/plugin_test.go @@ -0,0 +1,416 @@ +package toolprune + +import ( + "context" + "encoding/json" + "strings" + "sync" + "testing" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" +) + +// anthropicBody is deliberately awkward: unsorted keys, odd indentation, a +// trailing field after tools. Byte-exactness assertions below depend on it +// staying awkward, because the whole safety claim is "every byte outside the +// deleted elements is unchanged". +const anthropicBody = `{"model":"claude-opus-5", + "tools":[ + {"name":"Read","description":"read a file","input_schema":{"type":"object"}}, + {"name":"NotebookEdit","description":"edit a notebook","input_schema":{"type":"object"}}, + {"name":"Bash","description":"run a command","input_schema":{"type":"object"}} + ], + "max_tokens":1024,"stream":true}` + +func configured(t *testing.T, remove ...string) *ToolPrune { + t.Helper() + p := New() + raw, err := json.Marshal(map[string]any{"remove": remove}) + if err != nil { + t.Fatal(err) + } + if err := p.Configure(raw); err != nil { + t.Fatalf("Configure: %v", err) + } + return p +} + +func inferenceCtx(path, body string, toolNames ...string) *pipeline.Context { + pctx := &pipeline.Context{Path: path, Body: []byte(body)} + tools := make([]pipeline.InferenceTool, 0, len(toolNames)) + for _, n := range toolNames { + tools = append(tools, pipeline.InferenceTool{Name: n}) + } + pctx.Extensions.Inference = &pipeline.InferenceExtension{Tools: tools} + return pctx +} + +func run(t *testing.T, p *ToolPrune, pctx *pipeline.Context, policies ...pipeline.ErrorPolicy) { + t.Helper() + var opts []pipeline.Option + if len(policies) > 0 { + opts = append(opts, pipeline.WithPolicies(policies...)) + } + pipe, err := pipeline.New([]pipeline.Plugin{p}, opts...) + if err != nil { + t.Fatalf("pipeline.New: %v", err) + } + if act := pipe.Run(context.Background(), pctx); act.Type != pipeline.Continue { + t.Fatalf("action = %v, want Continue — tool-prune must never block a request", act.Type) + } +} + +// TestPrune_LeavesEveryOtherByteIntact is the core safety claim. Deleting a +// tool must not reformat the document, reorder keys, or disturb whitespace: the +// request that reaches the model has to be the one the client sent, minus +// exactly the elements named. +func TestPrune_LeavesEveryOtherByteIntact(t *testing.T) { + p := configured(t, "NotebookEdit") + pctx := inferenceCtx("/v1/messages", anthropicBody, "Read", "NotebookEdit", "Bash") + run(t, p, pctx) + + if !pctx.BodyMutated() { + t.Fatal("expected the body to be rewritten") + } + got := string(pctx.Body) + if strings.Contains(got, "NotebookEdit") { + t.Errorf("removed tool still present:\n%s", got) + } + for _, keep := range []string{ + `"model":"claude-opus-5"`, + `"name":"Read"`, + `"name":"Bash"`, + `"max_tokens":1024`, + `"stream":true`, + } { + if !strings.Contains(got, keep) { + t.Errorf("expected %s to survive verbatim:\n%s", keep, got) + } + } + // The only difference from the original must be the removed element. + if len(got) >= len(anthropicBody) { + t.Errorf("body did not shrink: %d -> %d", len(anthropicBody), len(got)) + } +} + +// TestPrune_DescendingDeletion: removing several tools by index only works if +// the deletions run high-to-low. An ascending loop would shift the array under +// itself and delete the wrong elements — here it would leave "Bash" and remove +// something else, so the assertion catches exactly that bug. +func TestPrune_DescendingDeletion(t *testing.T) { + body := `{"tools":[{"name":"A"},{"name":"B"},{"name":"C"},{"name":"D"},{"name":"E"}]}` + p := configured(t, "A", "B", "D") + pctx := inferenceCtx("/v1/messages", body, "A", "B", "C", "D", "E") + run(t, p, pctx) + + got := string(pctx.Body) + for _, gone := range []string{`"A"`, `"B"`, `"D"`} { + if strings.Contains(got, gone) { + t.Errorf("tool %s should be gone: %s", gone, got) + } + } + for _, kept := range []string{`"C"`, `"E"`} { + if !strings.Contains(got, kept) { + t.Errorf("tool %s should remain: %s", kept, got) + } + } +} + +// TestPrune_OpenAIDialect: OpenAI nests the name under function, Anthropic puts +// it at the top level. Both must resolve, since the plugin reads names out of +// the raw bytes rather than trusting manifest ordering. +func TestPrune_OpenAIDialect(t *testing.T) { + body := `{"tools":[{"type":"function","function":{"name":"Read"}},` + + `{"type":"function","function":{"name":"NotebookEdit"}}]}` + p := configured(t, "NotebookEdit") + pctx := inferenceCtx("/v1/chat/completions", body, "Read", "NotebookEdit") + run(t, p, pctx) + + got := string(pctx.Body) + if strings.Contains(got, "NotebookEdit") { + t.Errorf("removed tool still present: %s", got) + } + if !strings.Contains(got, "Read") { + t.Errorf("kept tool missing: %s", got) + } +} + +// TestPrune_RemovingEveryToolDropsTheKeys: an empty tools array is not a safe +// output — OpenAI rejects `tools: []`, and tool_choice without tools. Drop both +// keys instead, so an over-broad remove list still yields a valid request. +func TestPrune_RemovingEveryToolDropsTheKeys(t *testing.T) { + body := `{"model":"m","tools":[{"name":"A"},{"name":"B"}],"tool_choice":"auto"}` + p := configured(t, "A", "B") + pctx := inferenceCtx("/v1/chat/completions", body, "A", "B") + run(t, p, pctx) + + got := string(pctx.Body) + if strings.Contains(got, "tools") { + t.Errorf("tools key should be gone entirely, not left empty: %s", got) + } + if strings.Contains(got, "tool_choice") { + t.Errorf("tool_choice is invalid without tools; should be dropped: %s", got) + } + if !strings.Contains(got, `"model":"m"`) { + t.Errorf("unrelated fields must survive: %s", got) + } +} + +// TestPrune_UnknownNamesIgnored: a name absent from this request's manifest is +// simply not there — not an error. Drift in the configured list costs savings, +// never correctness. +func TestPrune_UnknownNamesIgnored(t *testing.T) { + body := `{"tools":[{"name":"Read"}]}` + p := configured(t, "ToolThatDoesNotExist") + pctx := inferenceCtx("/v1/messages", body, "Read") + run(t, p, pctx) + + if pctx.BodyMutated() { + t.Error("no configured tool was present; body must be untouched") + } + if string(pctx.Body) != body { + t.Errorf("body = %s, want unchanged", pctx.Body) + } +} + +// TestPrune_FailsOpen: malformed, truncated and manifest-less bodies all +// forward the original bytes. A cost optimisation must never break a request. +func TestPrune_FailsOpen(t *testing.T) { + cases := []struct { + name string + body string + }{ + {"malformed json", `{"tools":[{"name":"NotebookEdit"}`}, + {"truncated mid-string", `{"tools":[{"name":"Notebook`}, + {"tools is not an array", `{"tools":"NotebookEdit"}`}, + {"tools absent", `{"model":"m"}`}, + {"empty body", ``}, + {"empty tools array", `{"tools":[]}`}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + p := configured(t, "NotebookEdit") + pctx := inferenceCtx("/v1/messages", tc.body, "NotebookEdit") + run(t, p, pctx) + if pctx.BodyMutated() { + t.Errorf("body was mutated; must fail open on %s", tc.name) + } + if string(pctx.Body) != tc.body { + t.Errorf("body = %q, want original %q", pctx.Body, tc.body) + } + }) + } +} + +// TestPrune_PathGate: only inference paths are touched, so an unrelated POST +// through the same proxy is never rewritten. +func TestPrune_PathGate(t *testing.T) { + body := `{"tools":[{"name":"NotebookEdit"}]}` + p := configured(t, "NotebookEdit") + pctx := inferenceCtx("/v1/some/other/api", body, "NotebookEdit") + run(t, p, pctx) + if pctx.BodyMutated() { + t.Error("non-inference path must not be pruned") + } +} + +// TestPrune_EmptyRemoveListIsNoop: the shipped default is an empty list, so the +// plugin must be inert until an operator fills it in. +func TestPrune_EmptyRemoveListIsNoop(t *testing.T) { + p := configured(t) + pctx := inferenceCtx("/v1/messages", anthropicBody, "Read", "NotebookEdit") + run(t, p, pctx) + if pctx.BodyMutated() { + t.Error("empty remove list must not touch the body") + } + if p.Metrics() != nil { + t.Errorf("no requests acted on; Metrics should be nil, got %+v", p.Metrics()) + } +} + +// TestPrune_ObserveModeIsProjection: under on_error: observe the plugin computes +// exactly what it would remove and counts it, while the bytes on the wire stay +// untouched and the invocation is marked Shadow. That is what makes measure-only +// mode possible with one registration and no separate code path. +func TestPrune_ObserveModeIsProjection(t *testing.T) { + p := configured(t, "NotebookEdit") + pctx := inferenceCtx("/v1/messages", anthropicBody, "Read", "NotebookEdit", "Bash") + run(t, p, pctx, pipeline.ErrorPolicyObserve) + + if pctx.BodyMutated() { + t.Error("observe mode must leave the wire untouched") + } + if string(pctx.Body) != anthropicBody { + t.Errorf("body changed under observe:\n%s", pctx.Body) + } + if pctx.Extensions.Invocations == nil { + t.Fatal("expected invocations to be recorded") + } + var sawShadowModify bool + for _, inv := range pctx.Extensions.Invocations.Inbound { + if inv.Shadow && inv.Reason == "body_rewritten" { + sawShadowModify = true + } + } + if !sawShadowModify { + t.Errorf("expected a Shadow=true body_rewritten invocation, got %+v", + pctx.Extensions.Invocations.Inbound) + } + + // The projection must still be countable, and must be reported as a + // projection rather than a realised saving. + if p.m.requestsProjected != 1 { + t.Errorf("requestsProjected = %d, want 1", p.m.requestsProjected) + } + if p.m.requestsPruned != 0 { + t.Errorf("requestsPruned = %d, want 0 under observe", p.m.requestsPruned) + } + if p.m.bytesRemoved == 0 { + t.Error("bytesRemoved must accumulate under observe — that is the projection") + } + if !hasMetric(p.Metrics(), "requests projected") { + t.Errorf("readout should say 'requests projected': %+v", p.Metrics()) + } +} + +// TestPrune_EnforceCountsPruned is the enforce-mode counterpart. +func TestPrune_EnforceCountsPruned(t *testing.T) { + p := configured(t, "NotebookEdit") + pctx := inferenceCtx("/v1/messages", anthropicBody, "Read", "NotebookEdit", "Bash") + run(t, p, pctx, pipeline.ErrorPolicyEnforce) + + if p.m.requestsPruned != 1 { + t.Errorf("requestsPruned = %d, want 1", p.m.requestsPruned) + } + if p.m.requestsProjected != 0 { + t.Errorf("requestsProjected = %d, want 0 under enforce", p.m.requestsProjected) + } + if p.m.toolsRemoved != 1 { + t.Errorf("toolsRemoved = %d, want 1", p.m.toolsRemoved) + } + if !hasMetric(p.Metrics(), "removed: NotebookEdit") { + t.Errorf("per-tool attribution missing: %+v", p.Metrics()) + } +} + +// TestMetrics_NoUsageSampleReportsZeroNotNaN: the bytes-to-tokens ratio divides +// by a sample that starts empty. Report a zero-valued estimate with a note +// rather than NaN or a panic. +func TestMetrics_NoUsageSampleReportsZeroNotNaN(t *testing.T) { + p := configured(t, "NotebookEdit") + pctx := inferenceCtx("/v1/messages", anthropicBody, "Read", "NotebookEdit") + run(t, p, pctx) + + m := findMetric(t, p.Metrics(), "tokens saved / request") + if m.Value != 0 { + t.Errorf("value = %v, want 0 with no usage sample", m.Value) + } + if m.Note != "no usage sample yet" { + t.Errorf("note = %q, want the missing-sample caveat", m.Note) + } +} + +// TestMetrics_TokenEstimateCalibratesOnObservedUsage: once OnFinish has seen a +// response usage block, the estimate is derived from the operator's own +// traffic and labelled with its sample size. +func TestMetrics_TokenEstimateCalibratesOnObservedUsage(t *testing.T) { + p := configured(t, "NotebookEdit") + pctx := inferenceCtx("/v1/messages", anthropicBody, "Read", "NotebookEdit") + run(t, p, pctx) + + // 1 prompt token per 4 body bytes. + pctx.Extensions.Inference.PromptTokens = len(pctx.Body) / 4 + p.OnFinish(context.Background(), pctx) + + m := findMetric(t, p.Metrics(), "tokens saved / request") + if m.Value <= 0 { + t.Errorf("value = %v, want a positive estimate", m.Value) + } + if !strings.HasPrefix(m.Note, "estimate, n=") { + t.Errorf("note = %q, want it labelled an estimate with its sample size", m.Note) + } + perReq := findMetric(t, p.Metrics(), "bytes removed / request") + if want := perReq.Value / 4; m.Value < want*0.9 || m.Value > want*1.1 { + t.Errorf("estimate %v not within 10%% of calibrated %v", m.Value, want) + } +} + +// TestMetrics_ConcurrentAccess exercises Metrics() against live counter updates. +// describePipeline calls it from the HTTP handler while requests are in flight, +// so it must be safe under -race. +func TestMetrics_ConcurrentAccess(t *testing.T) { + p := configured(t, "NotebookEdit") + var wg sync.WaitGroup + for i := 0; i < 8; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 50; j++ { + pctx := inferenceCtx("/v1/messages", anthropicBody, "Read", "NotebookEdit") + pipe, err := pipeline.New([]pipeline.Plugin{p}) + if err != nil { + t.Error(err) + return + } + pipe.Run(context.Background(), pctx) + } + }() + } + for i := 0; i < 4; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 100; j++ { + _ = p.Metrics() + } + }() + } + wg.Wait() + if p.m.requestsPruned != 8*50 { + t.Errorf("requestsPruned = %d, want %d", p.m.requestsPruned, 8*50) + } +} + +func TestConfigure_RejectsUnknownFields(t *testing.T) { + p := New() + err := p.Configure(json.RawMessage(`{"remove":["A"],"nope":1}`)) + if err == nil { + t.Fatal("expected an error for an unknown config field") + } + if !strings.Contains(err.Error(), "tool-prune config") { + t.Errorf("error should name the plugin: %v", err) + } +} + +func TestCapabilities_RequestOnlySoStreamingSurvives(t *testing.T) { + caps := New().Capabilities() + if !caps.WritesRequestBody { + t.Error("must declare WritesRequestBody") + } + if caps.WritesResponseBody { + t.Error("must NOT declare WritesResponseBody — it would cost SSE streaming for nothing") + } + if len(caps.RequiresAny) != 1 || caps.RequiresAny[0] != "inference-parser" { + t.Errorf("RequiresAny = %v, want [inference-parser]", caps.RequiresAny) + } +} + +func hasMetric(ms []pipeline.Metric, name string) bool { + for _, m := range ms { + if m.Name == name { + return true + } + } + return false +} + +func findMetric(t *testing.T, ms []pipeline.Metric, name string) pipeline.Metric { + t.Helper() + for _, m := range ms { + if m.Name == name { + return m + } + } + t.Fatalf("metric %q not found in %+v", name, ms) + return pipeline.Metric{} +} diff --git a/authbridge/cmd/abctl/cmd_tools.go b/authbridge/cmd/abctl/cmd_tools.go new file mode 100644 index 000000000..784a92a4a --- /dev/null +++ b/authbridge/cmd/abctl/cmd_tools.go @@ -0,0 +1,88 @@ +package main + +import ( + "flag" + "fmt" + "io" + "strings" + + "github.com/rossoctl/cortex/authbridge/cmd/abctl/toolscan" +) + +const toolsUsage = `abctl tools scan — derive a tool-prune remove list from local transcripts + +Usage: + abctl tools scan [--days N] [--keep Name,Name] [--dir PATH] [--write CONFIG] + +Flags: + --days N window in days to consider a tool "used" (default 30) + --keep LIST comma-separated tool names to withhold from the candidate list + --dir PATH transcript directory (default ~/.claude/projects) + --write CONFIG patch the remove: list of the tool-prune entry in CONFIG in + place; without it, the YAML block is printed for you to paste + +Transcripts record tools that were called, never tools that were offered, so a +name abctl does not recognise is never proposed for removal. +` + +// runTools handles the `tools` subcommand. Returns the process exit code. +func runTools(args []string, stdout, stderr io.Writer) int { + if len(args) == 0 || args[0] != "scan" { + fmt.Fprint(stderr, toolsUsage) + return 2 + } + + fs := flag.NewFlagSet("tools scan", flag.ContinueOnError) + fs.SetOutput(stderr) + days := fs.Int("days", 30, "window in days") + keep := fs.String("keep", "", "comma-separated tool names to keep") + dir := fs.String("dir", "", "transcript directory (default ~/.claude/projects)") + write := fs.String("write", "", "patch the tool-prune remove: list in this config file") + if err := fs.Parse(args[1:]); err != nil { + return 2 + } + if *days <= 0 { + fmt.Fprintln(stderr, "abctl: --days must be positive") + return 2 + } + + scanDir := *dir + if scanDir == "" { + d, err := toolscan.DefaultProjectsDir() + if err != nil { + fmt.Fprintf(stderr, "abctl: locating transcripts: %v\n", err) + return 1 + } + scanDir = d + } + + res, err := toolscan.Scan(scanDir, *days, strings.Split(*keep, ",")) + if err != nil { + fmt.Fprintf(stderr, "abctl: scanning %s: %v\n", scanDir, err) + return 1 + } + if res.Files == 0 { + fmt.Fprintf(stderr, "abctl: no transcripts found under %s — nothing to infer from\n", scanDir) + return 1 + } + + fmt.Fprint(stdout, res.Summary(*days)) + if *write == "" { + fmt.Fprintln(stdout) + fmt.Fprint(stdout, res.YAMLBlock()) + return 0 + } + + changed, err := toolscan.PatchConfig(*write, res.Candidates) + if err != nil { + fmt.Fprintf(stderr, "abctl: %v\n", err) + return 1 + } + if changed { + fmt.Fprintf(stdout, "\nUpdated remove: list in %s (%d tool(s)).\n", *write, len(res.Candidates)) + fmt.Fprintln(stdout, "The config is hot-reloaded; no restart needed.") + } else { + fmt.Fprintf(stdout, "\n%s already up to date.\n", *write) + } + return 0 +} diff --git a/authbridge/cmd/abctl/main.go b/authbridge/cmd/abctl/main.go index 754863221..8565e1e83 100644 --- a/authbridge/cmd/abctl/main.go +++ b/authbridge/cmd/abctl/main.go @@ -13,6 +13,7 @@ import ( "os" "os/exec" "os/signal" + "strings" "syscall" "github.com/rossoctl/cortex/authbridge/cmd/abctl/cluster" @@ -25,6 +26,19 @@ import ( var version = "dev" func main() { + // Subcommand dispatch happens before flag.Parse: a non-flag first + // argument selects a subcommand, and anything else falls through to the + // terminal UI, preserving the original flags-only invocation. + if len(os.Args) > 1 && !strings.HasPrefix(os.Args[1], "-") { + switch os.Args[1] { + case "tools": + os.Exit(runTools(os.Args[2:], os.Stdout, os.Stderr)) + default: + fmt.Fprintf(os.Stderr, "abctl: unknown subcommand %q (known: tools)\n", os.Args[1]) + os.Exit(2) + } + } + endpoint := flag.String("endpoint", "", "AuthBridge session API URL (e.g. http://localhost:9094). When omitted, abctl opens a Namespaces → Pods picker.") showVersion := flag.Bool("version", false, "print version and exit") diff --git a/authbridge/cmd/abctl/toolscan/known.go b/authbridge/cmd/abctl/toolscan/known.go new file mode 100644 index 000000000..bf93becef --- /dev/null +++ b/authbridge/cmd/abctl/toolscan/known.go @@ -0,0 +1,67 @@ +// Package toolscan derives a tool-prune candidate list from local Claude Code +// transcripts. +// +// The central limitation is structural: transcripts record tools that were +// *called*, never tools that were *offered*. A configured-but-never-invoked +// tool leaves no trace at all. So the scan cannot enumerate the manifest — it +// can only intersect "tools we know Claude Code ships" with "tools this user +// never called". +// +// That shapes the safety rule: a name the scan has never heard of is always +// kept. Removing a tool the model needs is the harmful direction of failure; +// carrying a few extra definitions is merely expensive. Drift in the table +// below therefore costs savings, never correctness. +package toolscan + +// knownTools is the set of Claude Code built-in tool names the scanner is +// willing to propose for removal. Membership is a claim that the tool is +// bundled and that its absence is safe when it is never called. +// +// Deliberately conservative. Tools that gate control flow (ExitPlanMode), +// carry state the model relies on (TodoWrite), or are the primary means of +// doing work (Bash, Read, Edit, Write, Glob, Grep) are omitted entirely, so +// they can never be proposed however long they sit unused in a window. +var knownTools = []string{ + "Artifact", + "AskUserQuestion", + "BashOutput", + "CronCreate", + "CronDelete", + "CronList", + "DesignSync", + "EndConversation", + "EnterWorktree", + "ExitWorktree", + "KillShell", + "LSP", + "ListAgents", + "Monitor", + "NotebookEdit", + "PushNotification", + "ReportFindings", + "ScheduleWakeup", + "SendFeedback", + "SendMessage", + "SlashCommand", + "TaskOutput", + "TaskStop", + "WebFetch", + "WebSearch", + "Workflow", +} + +// implies covers tools whose use is indirect: the transcript shows the driver +// being called, not the tool it depends on. Keeping the right-hand side +// whenever the left-hand side was called prevents the scan from proposing a +// tool that is reachable but never appears by name. +var implies = map[string][]string{ + "Agent": {"SendMessage", "ListAgents", "TaskOutput", "TaskStop"}, + "Task": {"SendMessage", "ListAgents", "TaskOutput", "TaskStop"}, + "Monitor": {"TaskOutput", "TaskStop"}, + "Bash": {"BashOutput", "KillShell"}, + "Workflow": {"TaskOutput", "TaskStop"}, + "Artifact": {"DesignSync"}, +} + +// KnownTools returns a copy of the candidate universe. +func KnownTools() []string { return append([]string(nil), knownTools...) } diff --git a/authbridge/cmd/abctl/toolscan/patch.go b/authbridge/cmd/abctl/toolscan/patch.go new file mode 100644 index 000000000..154d903ef --- /dev/null +++ b/authbridge/cmd/abctl/toolscan/patch.go @@ -0,0 +1,82 @@ +package toolscan + +import ( + "fmt" + "os" + "regexp" + "strings" +) + +var ( + toolPruneEntry = regexp.MustCompile(`^(\s*)-\s+name:\s*tool-prune\s*$`) + listItem = regexp.MustCompile(`^\s*-\s`) + removeKey = regexp.MustCompile(`^(\s*)remove:\s*.*$`) +) + +// PatchConfig rewrites the remove: list of the tool-prune entry in the YAML at +// path, in place, and reports whether the file changed. +// +// Line-based on purpose. Round-tripping through a YAML library would reformat +// the whole document — dropping the comments that explain each plugin and +// reflowing entries the operator hand-tuned. The edit here touches exactly one +// line, so everything else in the file survives byte-for-byte, and re-running +// with the same candidates is a no-op. +func PatchConfig(path string, candidates []string) (changed bool, err error) { + orig, err := os.ReadFile(path) //nolint:gosec // operator-supplied config path + if err != nil { + return false, err + } + lines := strings.Split(string(orig), "\n") + + start := -1 + var entryIndent string + for i, l := range lines { + if m := toolPruneEntry.FindStringSubmatch(l); m != nil { + start, entryIndent = i, m[1] + break + } + } + if start < 0 { + return false, fmt.Errorf("no `- name: tool-prune` entry in %s — add the plugin to a pipeline first", path) + } + + // The entry ends at the next list item indented no deeper than this one. + end := len(lines) + for i := start + 1; i < len(lines); i++ { + if listItem.MatchString(lines[i]) && leadingSpaces(lines[i]) <= len(entryIndent) { + end = i + break + } + } + + want := "remove: []" + if len(candidates) > 0 { + want = fmt.Sprintf("remove: [%s]", strings.Join(candidates, ", ")) + } + for i := start + 1; i < end; i++ { + m := removeKey.FindStringSubmatch(lines[i]) + if m == nil { + continue + } + replacement := m[1] + want + if lines[i] == replacement { + return false, nil // already current — idempotent + } + lines[i] = replacement + out := strings.Join(lines, "\n") + if err := os.WriteFile(path, []byte(out), 0o600); err != nil { + return false, err + } + return true, nil + } + return false, fmt.Errorf("tool-prune entry in %s has no `remove:` key under config: — add `remove: []` and re-run", path) +} + +func leadingSpaces(s string) int { + for i, r := range s { + if r != ' ' && r != '\t' { + return i + } + } + return len(s) +} diff --git a/authbridge/cmd/abctl/toolscan/patch_test.go b/authbridge/cmd/abctl/toolscan/patch_test.go new file mode 100644 index 000000000..e04f3dd3e --- /dev/null +++ b/authbridge/cmd/abctl/toolscan/patch_test.go @@ -0,0 +1,175 @@ +package toolscan + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +const sampleConfig = `mode: proxy-sidecar +pipeline: + outbound: + plugins: + # Parses the inference request so downstream plugins see a manifest. + - name: inference-parser + - name: tool-prune + on_error: observe # measure only; switch to enforce when trusted + config: + remove: [] + - name: token-exchange + config: + keycloak_url: http://keycloak:8080 +` + +func writeConfig(t *testing.T, body string) string { + t.Helper() + p := filepath.Join(t.TempDir(), "demo.yaml") + if err := os.WriteFile(p, []byte(body), 0o600); err != nil { + t.Fatal(err) + } + return p +} + +// TestPatchConfig_TouchesOnlyTheRemoveLine: the operator's comments and +// hand-tuned entries must survive. This is why the patch is line-based rather +// than a YAML round-trip. +func TestPatchConfig_TouchesOnlyTheRemoveLine(t *testing.T) { + p := writeConfig(t, sampleConfig) + changed, err := PatchConfig(p, []string{"NotebookEdit", "WebSearch"}) + if err != nil { + t.Fatal(err) + } + if !changed { + t.Fatal("expected the file to change") + } + out, err := os.ReadFile(p) + if err != nil { + t.Fatal(err) + } + got := string(out) + if !strings.Contains(got, " remove: [NotebookEdit, WebSearch]") { + t.Errorf("remove line not patched (indentation must be preserved):\n%s", got) + } + for _, keep := range []string{ + "# Parses the inference request so downstream plugins see a manifest.", + "on_error: observe # measure only; switch to enforce when trusted", + "keycloak_url: http://keycloak:8080", + "mode: proxy-sidecar", + "- name: token-exchange", + } { + if !strings.Contains(got, keep) { + t.Errorf("patch disturbed unrelated content, missing %q:\n%s", keep, got) + } + } + // Exactly one line differs. + var diffs int + origLines := strings.Split(sampleConfig, "\n") + newLines := strings.Split(got, "\n") + if len(origLines) != len(newLines) { + t.Fatalf("line count changed: %d -> %d", len(origLines), len(newLines)) + } + for i := range origLines { + if origLines[i] != newLines[i] { + diffs++ + } + } + if diffs != 1 { + t.Errorf("%d lines changed, want exactly 1", diffs) + } +} + +// TestPatchConfig_Idempotent: install-demo.sh may run the scan on every +// invocation, so re-writing the same candidates must not report a change or +// rewrite the file. +func TestPatchConfig_Idempotent(t *testing.T) { + p := writeConfig(t, sampleConfig) + if _, err := PatchConfig(p, []string{"NotebookEdit"}); err != nil { + t.Fatal(err) + } + before, err := os.ReadFile(p) + if err != nil { + t.Fatal(err) + } + changed, err := PatchConfig(p, []string{"NotebookEdit"}) + if err != nil { + t.Fatal(err) + } + if changed { + t.Error("second identical patch reported a change; must be idempotent") + } + after, err := os.ReadFile(p) + if err != nil { + t.Fatal(err) + } + if string(before) != string(after) { + t.Error("file rewritten despite no change") + } +} + +func TestPatchConfig_EmptyCandidatesWritesEmptyList(t *testing.T) { + p := writeConfig(t, strings.Replace(sampleConfig, "remove: []", "remove: [NotebookEdit]", 1)) + changed, err := PatchConfig(p, nil) + if err != nil { + t.Fatal(err) + } + if !changed { + t.Fatal("expected a change back to an empty list") + } + out, _ := os.ReadFile(p) + if !strings.Contains(string(out), "remove: []") { + t.Errorf("want an empty list:\n%s", out) + } +} + +// TestPatchConfig_ErrorsWhenPluginAbsent: silently doing nothing would leave the +// operator believing the list was written. +func TestPatchConfig_ErrorsWhenPluginAbsent(t *testing.T) { + p := writeConfig(t, "pipeline:\n outbound:\n plugins:\n - name: token-exchange\n") + _, err := PatchConfig(p, []string{"NotebookEdit"}) + if err == nil { + t.Fatal("expected an error when the tool-prune entry is missing") + } + if !strings.Contains(err.Error(), "tool-prune") { + t.Errorf("error should name the missing entry: %v", err) + } +} + +func TestPatchConfig_ErrorsWhenRemoveKeyAbsent(t *testing.T) { + p := writeConfig(t, "pipeline:\n outbound:\n plugins:\n - name: tool-prune\n on_error: observe\n - name: token-exchange\n") + _, err := PatchConfig(p, []string{"NotebookEdit"}) + if err == nil { + t.Fatal("expected an error when remove: is missing") + } + if !strings.Contains(err.Error(), "remove:") { + t.Errorf("error should name the missing key: %v", err) + } +} + +// TestPatchConfig_DoesNotEscapeTheEntry: a remove: key belonging to a different +// plugin further down the file must not be hijacked. +func TestPatchConfig_DoesNotEscapeTheEntry(t *testing.T) { + cfg := `pipeline: + outbound: + plugins: + - name: tool-prune + on_error: observe + config: + remove: [] + - name: other-plugin + config: + remove: [SomethingElse] +` + p := writeConfig(t, cfg) + if _, err := PatchConfig(p, []string{"NotebookEdit"}); err != nil { + t.Fatal(err) + } + out, _ := os.ReadFile(p) + got := string(out) + if !strings.Contains(got, "remove: [SomethingElse]") { + t.Errorf("another plugin's remove list was modified:\n%s", got) + } + if !strings.Contains(got, "remove: [NotebookEdit]") { + t.Errorf("tool-prune's list was not patched:\n%s", got) + } +} diff --git a/authbridge/cmd/abctl/toolscan/scan.go b/authbridge/cmd/abctl/toolscan/scan.go new file mode 100644 index 000000000..96d688bf3 --- /dev/null +++ b/authbridge/cmd/abctl/toolscan/scan.go @@ -0,0 +1,185 @@ +package toolscan + +import ( + "bufio" + "encoding/json" + "fmt" + "io/fs" + "os" + "path/filepath" + "sort" + "strings" + "time" +) + +// Result is what a scan found. +type Result struct { + Since time.Time + Files int + Lines int // lines that survived the literal prefilter + Called []string // tool names actually invoked in the window, sorted + CallCounts map[string]int + Candidates []string // known, never called, not kept, not implied — sorted + Kept []string // names withheld by --keep or the implies table +} + +// transcriptEntry is the minimum shape needed. Decoding only these fields keeps +// the parse cheap on 40MB+ transcripts. +type transcriptEntry struct { + Timestamp time.Time `json:"timestamp"` + Message struct { + Content []struct { + Type string `json:"type"` + ID string `json:"id"` + Name string `json:"name"` + } `json:"content"` + } `json:"message"` +} + +// DefaultProjectsDir is where Claude Code keeps per-project transcripts. +func DefaultProjectsDir() (string, error) { + home, err := os.UserHomeDir() + if err != nil { + return "", err + } + return filepath.Join(home, ".claude", "projects"), nil +} + +// Scan walks dir for *.jsonl transcripts and derives a candidate list. +// +// Tool calls are deduplicated by the unique tool_use block id: the same +// assistant turn is rewritten into the transcript on every resume, so counting +// raw occurrences would inflate heavily-resumed sessions. +func Scan(dir string, days int, keep []string) (*Result, error) { + since := time.Now().AddDate(0, 0, -days) + res := &Result{Since: since, CallCounts: map[string]int{}} + + seenIDs := make(map[string]struct{}) + err := filepath.WalkDir(dir, func(path string, d fs.DirEntry, err error) error { + if err != nil { + // An unreadable subtree should not abort the whole scan. + return nil //nolint:nilerr // best-effort walk + } + if d.IsDir() || !strings.HasSuffix(path, ".jsonl") { + return nil + } + res.Files++ + return scanFile(path, since, seenIDs, res) + }) + if err != nil { + return nil, err + } + + // Expand the keep set with anything implied by a tool that WAS called. + keepSet := make(map[string]struct{}, len(keep)) + for _, k := range keep { + if k = strings.TrimSpace(k); k != "" { + keepSet[k] = struct{}{} + } + } + for name := range res.CallCounts { + for _, dep := range implies[name] { + keepSet[dep] = struct{}{} + } + } + + for name := range res.CallCounts { + res.Called = append(res.Called, name) + } + sort.Strings(res.Called) + + for _, known := range knownTools { + if _, called := res.CallCounts[known]; called { + continue + } + if _, kept := keepSet[known]; kept { + res.Kept = append(res.Kept, known) + continue + } + res.Candidates = append(res.Candidates, known) + } + sort.Strings(res.Candidates) + sort.Strings(res.Kept) + return res, nil +} + +func scanFile(path string, since time.Time, seenIDs map[string]struct{}, res *Result) error { + f, err := os.Open(path) //nolint:gosec // operator-supplied transcript dir + if err != nil { + return nil //nolint:nilerr // skip unreadable file + } + defer f.Close() + + sc := bufio.NewScanner(f) + // Transcript lines routinely exceed the default 64KB (a single tool result + // can be hundreds of KB), so give the scanner room before it errors. + sc.Buffer(make([]byte, 0, 256*1024), 16*1024*1024) + + for sc.Scan() { + line := sc.Bytes() + // Hot path: the overwhelming majority of lines carry no tool call. + // A literal substring check is far cheaper than parsing them. + if !strings.Contains(string(line), `"tool_use"`) { + continue + } + res.Lines++ + + var e transcriptEntry + if err := json.Unmarshal(line, &e); err != nil { + continue + } + if !e.Timestamp.IsZero() && e.Timestamp.Before(since) { + continue + } + for _, c := range e.Message.Content { + if c.Type != "tool_use" || c.Name == "" { + continue + } + if c.ID != "" { + if _, dup := seenIDs[c.ID]; dup { + continue + } + seenIDs[c.ID] = struct{}{} + } + res.CallCounts[c.Name]++ + } + } + return nil +} + +// YAMLBlock renders the candidate list as the config fragment an operator +// pastes (or --write patches) into the tool-prune entry. +func (r *Result) YAMLBlock() string { + var b strings.Builder + b.WriteString(" - name: tool-prune\n") + b.WriteString(" on_error: observe # measure only; switch to enforce when trusted\n") + b.WriteString(" config:\n") + if len(r.Candidates) == 0 { + b.WriteString(" remove: []\n") + return b.String() + } + fmt.Fprintf(&b, " remove: [%s]\n", strings.Join(r.Candidates, ", ")) + return b.String() +} + +// Summary is the human-readable preamble printed above the YAML block. +func (r *Result) Summary(days int) string { + var b strings.Builder + fmt.Fprintf(&b, "Scanned %d transcript(s), %d tool-call line(s), window %d day(s) since %s.\n", + r.Files, r.Lines, days, r.Since.Format("2006-01-02")) + fmt.Fprintf(&b, "Called in window (%d): %s\n", len(r.Called), joinOrNone(r.Called)) + fmt.Fprintf(&b, "Removal candidates (%d): %s\n", len(r.Candidates), joinOrNone(r.Candidates)) + if len(r.Kept) > 0 { + fmt.Fprintf(&b, "Withheld by --keep / implied-by-usage (%d): %s\n", len(r.Kept), joinOrNone(r.Kept)) + } + b.WriteString("\nNames not in abctl's known-tool table are never proposed: removing a tool\n") + b.WriteString("the model needs is the harmful failure, carrying extra definitions is not.\n") + return b.String() +} + +func joinOrNone(v []string) string { + if len(v) == 0 { + return "(none)" + } + return strings.Join(v, ", ") +} diff --git a/authbridge/cmd/abctl/toolscan/scan_test.go b/authbridge/cmd/abctl/toolscan/scan_test.go new file mode 100644 index 000000000..08e790138 --- /dev/null +++ b/authbridge/cmd/abctl/toolscan/scan_test.go @@ -0,0 +1,208 @@ +package toolscan + +import ( + "fmt" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +// entry renders one transcript line containing a tool_use block. +func entry(ts time.Time, id, name string) string { + return fmt.Sprintf( + `{"type":"assistant","timestamp":%q,"message":{"role":"assistant","content":[{"type":"tool_use","id":%q,"name":%q,"input":{}}]}}`, + ts.Format(time.RFC3339), id, name) +} + +func writeTranscript(t *testing.T, dir, name string, lines ...string) { + t.Helper() + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatal(err) + } + path := filepath.Join(dir, name) + if err := os.WriteFile(path, []byte(strings.Join(lines, "\n")+"\n"), 0o600); err != nil { + t.Fatal(err) + } +} + +func contains(v []string, s string) bool { + for _, x := range v { + if x == s { + return true + } + } + return false +} + +// TestScan_DeduplicatesByToolUseID: a transcript is rewritten on every resume, +// so the same tool_use block appears many times. Counting raw occurrences would +// make a heavily-resumed session look busier than it was — and, worse, could +// make a tool look "used" on the strength of one ancient call replayed often. +func TestScan_DeduplicatesByToolUseID(t *testing.T) { + dir := t.TempDir() + now := time.Now() + writeTranscript(t, dir, "a.jsonl", + entry(now, "toolu_1", "WebFetch"), + entry(now, "toolu_1", "WebFetch"), // same id, replayed + entry(now, "toolu_2", "WebFetch"), + ) + res, err := Scan(dir, 30, nil) + if err != nil { + t.Fatal(err) + } + if got := res.CallCounts["WebFetch"]; got != 2 { + t.Errorf("WebFetch counted %d times, want 2 (ids deduplicated)", got) + } +} + +// TestScan_WindowsByTimestamp: a tool called only outside the window must show +// up as a candidate, which is the entire point of --days. +func TestScan_WindowsByTimestamp(t *testing.T) { + dir := t.TempDir() + writeTranscript(t, dir, "a.jsonl", + entry(time.Now().AddDate(0, 0, -90), "toolu_old", "NotebookEdit"), + entry(time.Now(), "toolu_new", "WebFetch"), + ) + res, err := Scan(dir, 30, nil) + if err != nil { + t.Fatal(err) + } + if contains(res.Called, "NotebookEdit") { + t.Error("NotebookEdit was called 90 days ago; outside a 30-day window it is not 'called'") + } + if !contains(res.Candidates, "NotebookEdit") { + t.Errorf("NotebookEdit should be a candidate: %v", res.Candidates) + } + if !contains(res.Called, "WebFetch") { + t.Errorf("WebFetch is inside the window: %v", res.Called) + } + if contains(res.Candidates, "WebFetch") { + t.Error("a tool called inside the window must never be a candidate") + } +} + +// TestScan_UnknownNamesAreNeverProposed is the safety property. An MCP tool or +// a built-in from a newer Claude Code release is not in the table, so it can +// never be proposed for removal however long it goes unused. +func TestScan_UnknownNamesAreNeverProposed(t *testing.T) { + dir := t.TempDir() + writeTranscript(t, dir, "a.jsonl", entry(time.Now(), "toolu_1", "Bash")) + res, err := Scan(dir, 30, nil) + if err != nil { + t.Fatal(err) + } + for _, c := range res.Candidates { + if !contains(KnownTools(), c) { + t.Errorf("candidate %q is not in the known-tool table", c) + } + } + // A tool that does the primary work is not even in the table, so an idle + // window can't propose it. + for _, never := range []string{"Read", "Write", "Edit", "Bash", "Grep", "Glob", "TodoWrite", "ExitPlanMode"} { + if contains(res.Candidates, never) { + t.Errorf("%q must never be a removal candidate", never) + } + } +} + +// TestScan_ImpliesWithholdsIndirectlyUsedTools: Agent drives SendMessage, so a +// transcript showing Agent must not propose removing SendMessage even though +// SendMessage never appears by name. +func TestScan_ImpliesWithholdsIndirectlyUsedTools(t *testing.T) { + dir := t.TempDir() + writeTranscript(t, dir, "a.jsonl", entry(time.Now(), "toolu_1", "Agent")) + res, err := Scan(dir, 30, nil) + if err != nil { + t.Fatal(err) + } + if contains(res.Candidates, "SendMessage") { + t.Error("Agent implies SendMessage; it must be withheld, not proposed") + } + if !contains(res.Kept, "SendMessage") { + t.Errorf("SendMessage should be reported as withheld: %v", res.Kept) + } +} + +func TestScan_KeepFlagWithholds(t *testing.T) { + dir := t.TempDir() + writeTranscript(t, dir, "a.jsonl", entry(time.Now(), "toolu_1", "Bash")) + res, err := Scan(dir, 30, []string{"NotebookEdit", " WebSearch "}) + if err != nil { + t.Fatal(err) + } + for _, kept := range []string{"NotebookEdit", "WebSearch"} { + if contains(res.Candidates, kept) { + t.Errorf("%q was passed to --keep; must not be proposed", kept) + } + } +} + +// TestScan_SkipsLinesWithoutToolUse verifies the prefilter does not change +// results — only cost. A transcript of pure text must yield no calls. +func TestScan_SkipsLinesWithoutToolUse(t *testing.T) { + dir := t.TempDir() + writeTranscript(t, dir, "a.jsonl", + `{"type":"user","message":{"role":"user","content":"hello"}}`, + `{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"hi"}]}}`, + ) + res, err := Scan(dir, 30, nil) + if err != nil { + t.Fatal(err) + } + if len(res.Called) != 0 { + t.Errorf("Called = %v, want none", res.Called) + } + if res.Lines != 0 { + t.Errorf("Lines = %d, want 0 (prefilter should reject all)", res.Lines) + } +} + +// TestScan_ToleratesMalformedLines: a truncated final line (a crashed session) +// must not abort the scan. +func TestScan_ToleratesMalformedLines(t *testing.T) { + dir := t.TempDir() + writeTranscript(t, dir, "a.jsonl", + entry(time.Now(), "toolu_1", "WebFetch"), + `{"type":"assistant","message":{"content":[{"type":"tool_use","id":"toolu_2"`, + ) + res, err := Scan(dir, 30, nil) + if err != nil { + t.Fatalf("a malformed line must not fail the scan: %v", err) + } + if res.CallCounts["WebFetch"] != 1 { + t.Errorf("valid line should still be counted: %+v", res.CallCounts) + } +} + +func TestScan_WalksNestedProjectDirs(t *testing.T) { + root := t.TempDir() + writeTranscript(t, filepath.Join(root, "proj-a"), "s1.jsonl", entry(time.Now(), "t1", "WebFetch")) + writeTranscript(t, filepath.Join(root, "proj-b"), "s2.jsonl", entry(time.Now(), "t2", "Monitor")) + res, err := Scan(root, 30, nil) + if err != nil { + t.Fatal(err) + } + if res.Files != 2 { + t.Errorf("Files = %d, want 2", res.Files) + } + if !contains(res.Called, "WebFetch") || !contains(res.Called, "Monitor") { + t.Errorf("Called = %v, want both", res.Called) + } +} + +func TestYAMLBlock(t *testing.T) { + r := &Result{Candidates: []string{"NotebookEdit", "WebSearch"}} + got := r.YAMLBlock() + if !strings.Contains(got, "remove: [NotebookEdit, WebSearch]") { + t.Errorf("block missing remove list:\n%s", got) + } + if !strings.Contains(got, "on_error: observe") { + t.Errorf("emitted block must default to observe (measure first):\n%s", got) + } + empty := (&Result{}).YAMLBlock() + if !strings.Contains(empty, "remove: []") { + t.Errorf("no candidates should render an empty list:\n%s", empty) + } +} diff --git a/authbridge/cmd/authbridge-envoy/plugins_toolprune.go b/authbridge/cmd/authbridge-envoy/plugins_toolprune.go new file mode 100644 index 000000000..e8ac589f3 --- /dev/null +++ b/authbridge/cmd/authbridge-envoy/plugins_toolprune.go @@ -0,0 +1,5 @@ +//go:build !exclude_plugin_toolprune + +package main + +import _ "github.com/rossoctl/cortex/authbridge/authlib/plugins/toolprune" diff --git a/authbridge/cmd/authbridge-proxy/demo.go b/authbridge/cmd/authbridge-proxy/demo.go index a05905d51..30724fccb 100644 --- a/authbridge/cmd/authbridge-proxy/demo.go +++ b/authbridge/cmd/authbridge-proxy/demo.go @@ -45,7 +45,23 @@ tls_bridge: generate_ca: true pipeline: outbound: - plugins: [inference-parser, mcp-parser, a2a-parser] + plugins: + - name: inference-parser + - name: mcp-parser + - name: a2a-parser + # tool-prune drops unused tool definitions from the outbound manifest. + # It ships inert: the remove list is empty, so it does nothing until you + # fill it in, and on_error: observe means even then it only measures -- + # counting what it *would* remove while the bytes on the wire stay + # untouched. Read the projection in abctl's plugin pane, then switch + # on_error to enforce once the numbers look right. Fill the list with: + # abctl tools scan --write + # Keep it last: it rewrites the request body, and body readers must + # precede the mutator so they see the original bytes. + - name: tool-prune + on_error: observe + config: + remove: [] ` } diff --git a/authbridge/cmd/authbridge-proxy/demo_test.go b/authbridge/cmd/authbridge-proxy/demo_test.go index bb1a6faf1..c6fc6e028 100644 --- a/authbridge/cmd/authbridge-proxy/demo_test.go +++ b/authbridge/cmd/authbridge-proxy/demo_test.go @@ -3,6 +3,7 @@ package main import ( "path/filepath" "slices" + "strings" "testing" "github.com/rossoctl/cortex/authbridge/authlib/config" @@ -72,8 +73,31 @@ func TestDemoConfig_WriteLoadsAndValidates(t *testing.T) { for i, p := range cfg.Pipeline.Outbound.Plugins { gotPlugins[i] = p.Name } - wantPlugins := []string{"inference-parser", "mcp-parser", "a2a-parser"} + // tool-prune must come last: it is the request-body mutator, and the + // pipeline refuses to build a chain where a body reader follows it. + wantPlugins := []string{"inference-parser", "mcp-parser", "a2a-parser", "tool-prune"} if !slices.Equal(gotPlugins, wantPlugins) { t.Errorf("outbound plugins = %v, want %v", gotPlugins, wantPlugins) } + + // tool-prune ships inert, and that is a property worth pinning: the demo + // must never silently start rewriting a user's traffic. Two independent + // guards — an empty remove list (nothing to do) and observe policy + // (measure only) — so a future edit has to defeat both to enable it. + var tp *config.PluginEntry + for i := range cfg.Pipeline.Outbound.Plugins { + if cfg.Pipeline.Outbound.Plugins[i].Name == "tool-prune" { + tp = &cfg.Pipeline.Outbound.Plugins[i] + } + } + if tp == nil { + t.Fatal("tool-prune entry not found") + } + if tp.OnError != "observe" { + t.Errorf("tool-prune on_error = %q, want observe so the demo only measures", tp.OnError) + } + if !strings.Contains(string(tp.Config), "\"remove\":[]") && + !strings.Contains(string(tp.Config), "\"remove\": []") { + t.Errorf("tool-prune must ship with an empty remove list, got %s", tp.Config) + } } diff --git a/authbridge/cmd/authbridge-proxy/plugins_toolprune.go b/authbridge/cmd/authbridge-proxy/plugins_toolprune.go new file mode 100644 index 000000000..e8ac589f3 --- /dev/null +++ b/authbridge/cmd/authbridge-proxy/plugins_toolprune.go @@ -0,0 +1,5 @@ +//go:build !exclude_plugin_toolprune + +package main + +import _ "github.com/rossoctl/cortex/authbridge/authlib/plugins/toolprune" diff --git a/authbridge/docs/framework-architecture.md b/authbridge/docs/framework-architecture.md index 44cfdd468..5c0d3ced8 100644 --- a/authbridge/docs/framework-architecture.md +++ b/authbridge/docs/framework-architecture.md @@ -586,18 +586,30 @@ Always sequential. No priority / mode / fire-and-forget semantics yet. This is t A plugin that declares `WritesRequestBody: true` may rewrite the request or response body. The framework owns the propagation to the wire; plugins only call `pctx.SetBody(newBytes)` / `pctx.SetResponseBody(newBytes)`. -**Capability model.** Three booleans on `PluginCapabilities`: +**Capability model.** Body access is declared per direction on `PluginCapabilities`: | Field | Meaning | Listener effect | |---|---|---| | `ReadsBody` | plugin reads `pctx.Body` / `pctx.ResponseBody` | buffers the body; plugin sees the bytes | -| `WritesRequestBody` | plugin may call `pctx.SetBody` / `pctx.SetResponseBody` | implies `ReadsBody`; propagates mutations | +| `WritesRequestBody` | plugin may call `pctx.SetBody` | implies `ReadsBody`; propagates request mutations | +| `WritesResponseBody` | plugin may call `pctx.SetResponseBody` | implies `ReadsBody`; propagates response mutations **and forces the buffered response path** | | `BodyAccess` (deprecated) | legacy alias for `ReadsBody` | folded by `Normalize()`, removed in a future release | +**Why the directions are separate.** `Pipeline.WritesResponseBody()` is the SSE +streaming predicate: both proxy listeners consult it to decide whether a +`text/event-stream` response may be relayed incrementally. It was previously one +undirected flag, which meant a plugin rewriting only the *request* body disabled +*response* streaming for bytes it never touched. The cost was latency and feel +rather than correctness — the buffered path restores the body verbatim — but a +long completion arriving in one lump after a silent wait is the first thing +anyone notices. Request bodies are never streamed (they arrive complete with a +`Content-Length` and are read end to end before dispatch), so a request-only +mutator now keeps incremental relay. + `pipeline.New` enforces two rules at build time: -1. **At most one `WritesRequestBody` plugin per pipeline.** Multiple mutators would have ambiguous ordering semantics; the error names both plugins so an operator debugging pod logs knows which two to reconcile. -2. **`WritesRequestBody` cannot precede a `ReadsBody`-only plugin.** A reader expects to see the original bytes; putting a mutator before it would silently feed the reader the post-rewrite content. +1. **At most one mutator per direction per pipeline.** Multiple mutators writing the same bytes would have ambiguous ordering semantics; the error names both plugins so an operator debugging pod logs knows which two to reconcile. A request mutator and a response mutator coexist fine. +2. **A mutator of either direction cannot precede a `ReadsBody`-only plugin.** A reader expects to see the original bytes; putting a mutator before it would silently feed the reader the post-rewrite content. **Mutation helpers.** `SetBody` / `SetResponseBody` replace the byte slice and flip an internal `bodyMutated` / `responseBodyMutated` flag that listeners read via `pctx.BodyMutated()` / `pctx.ResponseBodyMutated()`. They also auto-emit: diff --git a/authbridge/docs/plugin-catalog.md b/authbridge/docs/plugin-catalog.md index 450bbad7d..8ce5f5a73 100644 --- a/authbridge/docs/plugin-catalog.md +++ b/authbridge/docs/plugin-catalog.md @@ -36,6 +36,7 @@ AuthBridge pipeline YAML, not whether it is compiled into the binary | [`session-budget`](#session-budget) | Enforces per-session token, call, and duration budgets via Redis. | Alpha | Outbound | No | | [`token-broker`](#token-broker) | Exchanges incoming tokens against a configured IdP via a broker service. | Alpha | Outbound | No | | [`token-exchange`](#token-exchange) | RFC 8693 outbound token exchange per route. | Ready | Outbound | YES | +| [`tool-prune`](#tool-prune) | Removes unused tool definitions from inference requests. | Alpha | Outbound | No | ## `a2a-parser` @@ -225,3 +226,23 @@ ID, Okta, and any RFC 8693-compliant IdP. - `routes.rules` (list) — inline route entries (`host`, `target_audience`, `token_scopes`, `token_url`, `action`), combined with file-loaded routes. - `audience_from_host` (bool) — derive audience from host for unrouted requests (waypoint mode). Default `false`. - `resolve_placeholders` (bool) — resolve an inbound placeholder-prefixed bearer to its real token before exchange; unresolvable placeholders are denied. Default `false`. + +## `tool-prune` + +Removes unused tool definitions from the outbound inference manifest, so +the tokens for tools an agent never calls are not billed on every turn. +The manifest is assembled by the client, so the proxy is the only place to +trim it without changing every client. + +Requires `inference-parser` earlier in the chain, and must sit after any +body-reading plugin (it rewrites the request body). Declares +`WritesRequestBody` only, so response streaming is unaffected. + +- `remove` (`[]string`) — tool names to delete from the manifest. The complete verdict: no learning, no state, no storage. Names absent from a given request are ignored. +- `paths` (`[]string`) — request paths to act on, matched exactly or by suffix. Defaults to `/v1/chat/completions`, `/v1/completions`, `/v1/messages`. + +Generate the list from local transcripts with `abctl tools scan`, which +proposes only tools it recognises as Claude Code built-ins and never +proposes one it has seen called. See +[`tool-prune-plugin.md`](./tool-prune-plugin.md) for the measure-then-enforce +rollout, the metrics readout, and what the saving does and does not change. diff --git a/authbridge/docs/plugin-reference.md b/authbridge/docs/plugin-reference.md index f4a855339..d564db001 100644 --- a/authbridge/docs/plugin-reference.md +++ b/authbridge/docs/plugin-reference.md @@ -616,6 +616,45 @@ separator than against an escape convention. `Details`.** The session store has no auth on it; only safe-to-log data belongs in Invocations. +### 1b. Operator-facing counters (`MetricsProvider`) + +Invocations describe *this* request. For running totals an operator reads while +debugging — how many requests a plugin acted on, how many bytes it saved — +implement `pipeline.MetricsProvider`: + +```go +type Metric struct { + Name string `json:"name"` + Value float64 `json:"value"` + Unit string `json:"unit,omitempty"` // count | bytes | tokens | ratio + Note string `json:"note,omitempty"` // e.g. "estimate, n=1284" +} + +type MetricsProvider interface { Metrics() []Metric } +``` + +`describePipeline` calls `Metrics()` while serving `/v1/pipeline`, so it must be +safe for concurrent use with the request path and must not block — take a +mutex, copy, release. Returning `nil` is fine; abctl renders `(none)`. + +Rules worth honouring: + +- **Label anything derived.** A figure the plugin computed rather than counted + goes in with a `Note` naming its sample size. A derived number with no `Note` + reads as a measurement. +- **Report the sample alongside the conclusion.** Counters are per-process and + reset on restart *and on config hot-reload* (a reload rebuilds the plugin), so + a bare average is uninterpretable without the count behind it. +- **Don't route counters through `auth.Stats`.** That type is auth-shaped — + typed approval/denial enums and a custom `MarshalJSON` — and carrying + unrelated totals through it distorts its meaning. + +This is an optional interface, so it is **not** promoted through +`configuredPlugin`'s embedded `Plugin`. The wrapper forwards it explicitly, the +same way it forwards `Initializer` / `Shutdowner` / `Finisher` / `Readier`; a new +optional interface must be added there too or it will be invisible for every +plugin that has config. + ### 2. Named protocol extension (optional, for parsers) `MCP`, `A2A`, `Inference` are typed slots on `pipeline.Extensions`. @@ -699,26 +738,53 @@ before/after (never the raw body). ```go type PluginCapabilities struct { - ReadsBody bool // plugin reads pctx.Body / pctx.ResponseBody - WritesRequestBody bool // plugin may call pctx.SetBody / pctx.SetResponseBody + ReadsBody bool // plugin reads pctx.Body / pctx.ResponseBody + WritesRequestBody bool // plugin may call pctx.SetBody + WritesResponseBody bool // plugin may call pctx.SetResponseBody } ``` - `ReadsBody`: listener buffers the body; plugin sees bytes. - `WritesRequestBody`: implies `ReadsBody`. Listener propagates `pctx.SetBody` - rewrites to the upstream (and `pctx.SetResponseBody` to the - downstream client). + rewrites to the upstream. +- `WritesResponseBody`: implies `ReadsBody`. Listener propagates + `pctx.SetResponseBody` rewrites to the downstream client. + +**Declare the direction you actually write.** The two flags are not +interchangeable, and getting this wrong is silent: `WritesResponseBody` is the +SSE streaming predicate. A plugin that declares it forces every response on +that chain onto the buffered path, so a long completion arrives in one lump +after a silent wait instead of appearing incrementally. Declaring +`WritesRequestBody` costs nothing — requests arrive complete with a +`Content-Length` and are read end to end before dispatch, so rewriting one says +nothing about how the response may be relayed. + +| Plugin shape | Declares | Streams responses? | +|---|---|---| +| request-only mutator (`tool-prune`, `context-guru`) | `WritesRequestBody` | yes | +| response mutator | `WritesResponseBody` | no — buffered | +| both (`sparc`, `cpex`) | both | no — buffered | +| pure reader (parsers) | `ReadsBody` | yes | ### Build-time validation (enforced by `pipeline.New`) -- At most **one** `WritesRequestBody` plugin per pipeline. Two mutators in - the same direction would produce ambiguous ordering; `New` rejects - with an error naming both plugins. -- A `WritesRequestBody` plugin cannot precede a `ReadsBody`-only plugin. The - reader must see the original bytes. +- At most **one** mutator **per direction** per pipeline. Two request mutators + (or two response mutators) would produce ambiguous ordering; `New` rejects + with an error naming both plugins. One request mutator plus one response + mutator is fine — they never rewrite the same bytes. +- A mutator of **either** direction cannot precede a `ReadsBody`-only plugin. + The reader must see the original bytes. - Waypoint mode (ext_authz listener) cannot propagate body mutations — the ext_authz API has no body-mutation field. Do not combine - `WritesRequestBody: true` plugins with `mode: waypoint`. + body-mutating plugins with `mode: waypoint`. + +> **Declaring is a contract, not an enforcement.** `SetBody` flips +> `bodyMutated` unconditionally outside observe mode and the listeners gate +> purely on that flag, so a plugin that calls `SetBody` *without* declaring the +> capability still reaches the wire. This divergence is documented rather than +> closed, because adding the check silently would break out-of-tree plugins +> relying on today's behaviour. Do not read it as a way to keep response +> streaming — declare `WritesRequestBody`, which costs no streaming anyway. ### Mutation helpers diff --git a/authbridge/docs/plugin-tutorial.md b/authbridge/docs/plugin-tutorial.md index b6060cddd..d51ca5fe9 100644 --- a/authbridge/docs/plugin-tutorial.md +++ b/authbridge/docs/plugin-tutorial.md @@ -209,8 +209,15 @@ both stay nil even if you try to read them. ### Mutating the body If your plugin needs to **rewrite** the body — prompt-redaction, output -filtering, content transformation — declare `WritesRequestBody` and call -`pctx.SetBody` / `pctx.SetResponseBody`: +filtering, content transformation — declare the direction you write and call +the matching helper: `WritesRequestBody` for `pctx.SetBody`, +`WritesResponseBody` for `pctx.SetResponseBody`. + +Declare only what you actually write. `WritesResponseBody` is the SSE streaming +predicate, so claiming it when you only rewrite requests costs every caller on +that chain incremental relay — a long completion arrives in one lump after a +silent wait. `WritesRequestBody` costs nothing: requests arrive complete and are +read end to end before dispatch. ```go func (p *Redactor) Capabilities() pipeline.PluginCapabilities { @@ -232,10 +239,11 @@ for `SetResponseBody`) with a correct `Content-Length` and a cleared (never the raw body content). **Rules enforced by `pipeline.New`:** -- At most one `WritesRequestBody` plugin per pipeline. Two mutators = ambiguous - ordering → build fails at startup. -- A `WritesRequestBody` plugin must run **after** any `ReadsBody`-only plugin. - Readers see the original bytes; a mutator in front would silently +- At most one mutator **per direction** per pipeline. Two request mutators (or + two response mutators) = ambiguous ordering → build fails at startup. One of + each is fine; they never rewrite the same bytes. +- A mutator must run **after** any `ReadsBody`-only plugin, whichever direction + it writes. Readers see the original bytes; a mutator in front would silently feed them post-rewrite content. Don't assign `pctx.Body = newBytes` directly — the listener won't diff --git a/authbridge/docs/tool-prune-plugin.md b/authbridge/docs/tool-prune-plugin.md new file mode 100644 index 000000000..bc0ab2581 --- /dev/null +++ b/authbridge/docs/tool-prune-plugin.md @@ -0,0 +1,165 @@ +# `tool-prune` plugin + +Removes unused tool definitions from outbound inference requests. + +A Claude Code turn carries the full tool manifest on every request — tens of +thousands of tokens of JSON schema, billed each time, largely for tools the +agent will never call in a given deployment. The manifest is assembled by the +client, so the proxy is the only place to trim it without changing every client. + +The verdict is entirely configuration. `remove` names the tools to drop; there +is no learning, no state and no storage dependency. `abctl tools scan` proposes +a list, but the plugin only ever does what it was told. + +## Configuration + +```yaml +pipeline: + outbound: + plugins: + - inference-parser + - mcp-parser + - name: tool-prune + on_error: observe # measure only; switch to enforce when trusted + config: + remove: [NotebookEdit, ScheduleWakeup, TaskOutput] +``` + +| Field | Type | Default | Meaning | +|---|---|---|---| +| `remove` | `[]string` | `[]` | Tool names to delete. Names absent from a given request are ignored. | +| `paths` | `[]string` | `/v1/chat/completions`, `/v1/completions`, `/v1/messages` | Request paths to act on, matched exactly or by suffix. | + +**Placement matters.** `tool-prune` requires `inference-parser` earlier in the +chain, and because it rewrites the request body it must sit *after* every +body-reading plugin — readers have to see the original bytes. `pipeline.New` +enforces both and fails at startup rather than misbehaving quietly. + +It declares `WritesRequestBody` only, never `WritesResponseBody`, so responses +still stream incrementally. See +[`plugin-reference.md`](./plugin-reference.md#capability-fields). + +## Measure first, then enforce + +`on_error: observe` makes the plugin a projection: it computes exactly what it +would remove and counts it, while the bytes on the wire stay untouched. Nothing +about the plugin's code differs between the two modes — under observe, `SetBody` +is a no-op on bytes and leaves `BodyMutated()` false, which is how the plugin +knows which counter to increment. + +So the rollout is: add it in observe, read the projection, then flip one word. + +```sh +abctl tools scan --write ./cortex-ca/demo.yaml # fill in remove: +# read the projection in abctl's plugin pane, then change on_error to enforce +``` + +The config is hot-reloaded, so neither step needs a restart. Note that a reload +rebuilds the plugin and therefore **resets its counters** — the same as a +process restart. + +## Reading the metrics + +`abctl`'s plugin detail pane shows a `Metrics:` section (source: +`GET /v1/pipeline`): + +``` +Metrics: + requests seen 3 count + requests pruned 3 count + tools removed 6 count + bytes removed 825 bytes + bytes removed / request 275 bytes + tokens saved / request 343.75 tokens estimate, n=3 + removed: NotebookEdit 3 count + removed: ScheduleWakeup 3 count +``` + +In observe mode `requests projected` replaces `requests pruned`, so a +projection is never mistaken for a realised saving. + +Byte counts are exact. The token figure is an **estimate**, and labelled as one +with its sample size: rather than bundling a tokenizer or assuming a +bytes-per-token constant, the ratio is calibrated on your own traffic from the +response `usage` block. One approximation to state plainly — under `enforce`, +`PromptTokens` is already the post-pruning count, so the ratio is measured on +pruned requests. That is acceptable for a conversion factor, which is a property +of the tokenizer and content mix rather than of the pruning, but it is why the +number is an estimate. + +Counters are in-memory and per-process. That is the right trade for the +single-laptop case this targets and what keeps the plugin free of a storage +dependency; fleet aggregation belongs on the stats server later and would not +change the plugin. + +## Where the list comes from + +``` +abctl tools scan [--days 30] [--keep Name,Name] [--dir PATH] [--write CONFIG] +``` + +It reads `~/.claude/projects/**/*.jsonl`, deduplicates tool calls by their +unique `tool_use` block id (a transcript is rewritten on every resume, so raw +occurrences would inflate heavily-resumed sessions), and windows to the last +`--days`. Without `--write` it prints the YAML block; with `--write` it patches +the `remove:` list of the `tool-prune` entry in place, idempotently and without +reformatting the rest of the file. + +**The offered-set problem.** Transcripts record tools that were *called*, never +tools that were *offered*. This is structural, not a defect: a +configured-but-never-invoked tool leaves no trace. Two consequences: + +- The removal candidates are tools abctl knows Claude Code ships that you never + called — which is also where most of the wasted tokens sit. +- A tool name the scan has never heard of is **kept**. Removing a tool the model + needs is the harmful direction of failure; carrying a few extra definitions is + merely expensive. Drift in the known-tool table costs savings, never + correctness. + +A `--keep` flag and a small implies table cover tools whose use is indirect — +`Agent` implying `SendMessage`, say, which a transcript may never show being +called by name. At runtime the plugin also logs, once, any configured name +absent from the first manifest it sees, so a stale list surfaces as a warning +rather than a silent no-op. + +## Failure behaviour + +Every error path forwards the original bytes unmodified. A cost optimisation +must never be able to break a request, so the plugin fails open on a malformed +or truncated body, an unparseable manifest, a rewrite that does not shrink the +body, a rewrite that produces invalid JSON, an unexpected tool count afterwards, +and any panic. + +Two specifics worth knowing: + +- **Nothing else in the request changes.** Deletions are surgical: every byte + outside the removed array elements is preserved, including key order and + whitespace. +- **Removing every tool drops the keys.** An empty `tools: []` is not a safe + output — OpenAI rejects it, and rejects `tool_choice` without `tools` — so an + over-broad list removes both keys instead of emptying the array. + +## What the saving does and does not change + +`/cost` and anything derived from the API response `usage` block **do** move: +the server bills the request it received, so `input_tokens` and +`cache_read_input_tokens` genuinely drop. + +Claude Code's `/context` breakdown **does not**. It is a client-side pre-flight +view of what the CLI assembled, and it computes `Free space` itself; the pruning +happens downstream. This is the first place anyone looks, so it is worth stating +plainly: proxy-side pruning saves money but does not return context window. The +client still believes it sent the full manifest, so auto-compact triggers at the +same point. Recovering headroom needs client-side configuration +(`--allowedTools`, disabling unused MCP servers). AuthBridge's advantage is the +complement — it applies to every agent behind it with no per-client change, and +it measures. + +One further caveat on the list changing: a new `remove` list invalidates the +prompt-cache prefix once. That is inherent and bounded — the list is static, so +it happens on the change and then the prefix is stable again. + +## Build tag + +Compiled in by default; exclude with `-tags exclude_plugin_toolprune`. The +`authbridge-lite` image excludes it along with the other non-auth plugins. diff --git a/authbridge/install-demo.sh b/authbridge/install-demo.sh index ad31ebbcd..c76a09b91 100755 --- a/authbridge/install-demo.sh +++ b/authbridge/install-demo.sh @@ -196,6 +196,19 @@ else info "Cortex demo started (pid ${demo_pid}); couldn't confirm it's listening (install lsof or nc to verify). Logs: ${log}" fi info "" + +# tool-prune ships inert: present in the pipeline with an empty remove list and +# on_error: observe. Offer the scan that fills the list in. Only patch a config +# that already exists, so a first run never rewrites a file it just created +# behind the user's back -- print the command instead and let them look first. +demo_cfg="${ca_dir}/demo.yaml" +if [ -f "${demo_cfg}" ]; then + info " Measure tool-manifest waste (writes the remove: list, no restart needed):" + info " ${abctl_cmd} tools scan --write ${demo_cfg}" + info " Then read 'requests projected' in abctl's plugin pane before switching" + info " that entry's on_error to enforce." + info "" +fi info " Watch traffic: ${abctl_cmd} --endpoint http://localhost:${DEMO_SESSION_PORT}" info " Send traffic through it (e.g. Claude Code):" info " HTTPS_PROXY=http://localhost:${DEMO_FORWARD_PORT} \\" From 7a5e53205aa1f0295afe6d8c95a9430ec24e3c8d Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 14:04:07 -0400 Subject: [PATCH 07/28] fix(authbridge): Never prune a tool that tool_choice forces MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tool_choice can name a specific tool — Anthropic's {"type":"tool","name":"X"} or OpenAI's {"type":"function","function":{"name":"X"}}. A tool_choice naming a tool absent from the manifest is an invalid request, so pruning X while leaving the forced choice in place produced a body the provider rejects. tool_choice was only dropped when the remove list emptied the manifest entirely; the far likelier partial case was unhandled. The forced tool is now kept regardless of the configured list, and the rest of the list still applies, so the saving is preserved without constructing an invalid request. Also corrects the safety claim this contradicted. "A cost optimisation must never be able to break a request" overstated what the plugin can promise: its own failure paths do fail open, but whether a provider or gateway accepts a validly pruned manifest is outside what it can observe. The docs now say that plainly and point at on_error: observe as the way to establish it empirically — which is the whole reason measure-only mode exists. Signed-off-by: Hai Huang --- .../authlib/plugins/toolprune/plugin.go | 35 +++++++++++- .../authlib/plugins/toolprune/plugin_test.go | 57 +++++++++++++++++++ authbridge/docs/tool-prune-plugin.md | 27 ++++++--- docs/proposals/tool-prune.md | 11 +++- 4 files changed, 119 insertions(+), 11 deletions(-) diff --git a/authbridge/authlib/plugins/toolprune/plugin.go b/authbridge/authlib/plugins/toolprune/plugin.go index 89d4f1a1d..40ab9d6d7 100644 --- a/authbridge/authlib/plugins/toolprune/plugin.go +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -14,8 +14,14 @@ // // Safety is one-directional. Removing a tool the model needs is the harmful // failure; carrying a few extra definitions is not. So every error path fails -// open, forwarding the original bytes untouched: a cost optimisation must never -// be able to break a request. +// open, forwarding the original bytes untouched, and a tool named by a forced +// tool_choice is never removed — the manifest and tool_choice have to agree or +// the request is invalid. +// +// That is a promise about this plugin's own failure modes, not a claim that +// pruning is always safe: whether a provider or gateway accepts a validly +// pruned manifest is outside what the plugin can observe. on_error: observe +// exists to establish that empirically before any request changes. package toolprune import ( @@ -138,6 +144,24 @@ func toolNameAt(body []byte, i int) string { return gjson.GetBytes(body, fmt.Sprintf("tools.%d.function.name", i)).String() } +// forcedToolName returns the tool a forced tool_choice names, or "" when the +// request does not force one. Anthropic spells it tool_choice.name, OpenAI +// tool_choice.function.name; "auto" / "none" / "any" carry no name. +// +// This tool can never be removed: a tool_choice naming a tool absent from the +// manifest is an invalid request, so pruning it would turn a cost optimisation +// into a 400. +func forcedToolName(body []byte) string { + tc := gjson.GetBytes(body, "tool_choice") + if !tc.IsObject() { + return "" // "auto" / "none" / absent + } + if n := tc.Get("name"); n.Exists() { + return n.String() + } + return tc.Get("function.name").String() +} + // OnRequest prunes the manifest. Every failure path returns Continue with the // body untouched. func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action pipeline.Action) { @@ -189,6 +213,7 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action // Resolve indices from the raw bytes rather than from the parsed manifest: // inference-parser drops unnamed tools, so manifest position does not // reliably map back to array position. + forced := forcedToolName(body) var victims []int names := make([]string, 0, len(raw)) for i := range raw { @@ -196,6 +221,12 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action if name == "" { continue } + if name == forced { + // Removing the tool tool_choice forces would make the request + // invalid. Keep it and prune the rest. + slog.Debug("tool-prune: keeping tool forced by tool_choice", "tool", name) + continue + } if _, ok := p.remove[name]; ok { victims = append(victims, i) names = append(names, name) diff --git a/authbridge/authlib/plugins/toolprune/plugin_test.go b/authbridge/authlib/plugins/toolprune/plugin_test.go index 33e23d3cb..1ebf49d57 100644 --- a/authbridge/authlib/plugins/toolprune/plugin_test.go +++ b/authbridge/authlib/plugins/toolprune/plugin_test.go @@ -414,3 +414,60 @@ func findMetric(t *testing.T, ms []pipeline.Metric, name string) pipeline.Metric t.Fatalf("metric %q not found in %+v", name, ms) return pipeline.Metric{} } + +// TestPrune_NeverRemovesForcedToolChoice: a tool_choice that forces a specific +// tool must keep that tool, whichever dialect spells it. Removing it leaves a +// tool_choice naming a tool absent from the manifest, which providers reject — +// turning a cost optimisation into a 400, the one thing this plugin must never +// do. The rest of the remove list still applies. +func TestPrune_NeverRemovesForcedToolChoice(t *testing.T) { + cases := []struct { + name string + body string + }{ + { + name: "anthropic tool_choice.name", + body: `{"tools":[{"name":"Read"},{"name":"WebSearch"},{"name":"NotebookEdit"}],` + + `"tool_choice":{"type":"tool","name":"WebSearch"}}`, + }, + { + name: "openai tool_choice.function.name", + body: `{"tools":[{"type":"function","function":{"name":"Read"}},` + + `{"type":"function","function":{"name":"WebSearch"}},` + + `{"type":"function","function":{"name":"NotebookEdit"}}],` + + `"tool_choice":{"type":"function","function":{"name":"WebSearch"}}}`, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + // Both WebSearch (forced) and NotebookEdit are configured for removal. + p := configured(t, "WebSearch", "NotebookEdit") + pctx := inferenceCtx("/v1/messages", tc.body, "Read", "WebSearch", "NotebookEdit") + run(t, p, pctx) + + got := string(pctx.Body) + if !strings.Contains(got, "WebSearch") { + t.Errorf("forced tool was removed — request is now invalid:\n %s", got) + } + if strings.Contains(got, "NotebookEdit") { + t.Errorf("non-forced tool should still be pruned:\n %s", got) + } + }) + } +} + +// TestPrune_ToolChoiceAutoDoesNotBlockPruning: "auto" / "none" name no tool, so +// they must not be mistaken for a forced choice and suppress all pruning. +func TestPrune_ToolChoiceAutoDoesNotBlockPruning(t *testing.T) { + for _, choice := range []string{`"auto"`, `"none"`, `{"type":"auto"}`} { + t.Run(choice, func(t *testing.T) { + body := `{"tools":[{"name":"Read"},{"name":"NotebookEdit"}],"tool_choice":` + choice + `}` + p := configured(t, "NotebookEdit") + pctx := inferenceCtx("/v1/messages", body, "Read", "NotebookEdit") + run(t, p, pctx) + if strings.Contains(string(pctx.Body), "NotebookEdit") { + t.Errorf("tool_choice %s should not suppress pruning:\n %s", choice, pctx.Body) + } + }) + } +} diff --git a/authbridge/docs/tool-prune-plugin.md b/authbridge/docs/tool-prune-plugin.md index bc0ab2581..c2ca5e2ac 100644 --- a/authbridge/docs/tool-prune-plugin.md +++ b/authbridge/docs/tool-prune-plugin.md @@ -124,13 +124,26 @@ rather than a silent no-op. ## Failure behaviour -Every error path forwards the original bytes unmodified. A cost optimisation -must never be able to break a request, so the plugin fails open on a malformed -or truncated body, an unparseable manifest, a rewrite that does not shrink the -body, a rewrite that produces invalid JSON, an unexpected tool count afterwards, -and any panic. - -Two specifics worth knowing: +Every error path forwards the original bytes unmodified: the plugin fails open on +a malformed or truncated body, an unparseable manifest, a rewrite that does not +shrink the body, a rewrite that produces invalid JSON, an unexpected tool count +afterwards, and any panic. + +**What that does and does not promise.** It means the plugin's own failure modes +cannot break a request — a bug or a surprising input forwards the original bytes +rather than a damaged rewrite. It does **not** promise that a validly pruned +manifest is acceptable to every provider or gateway in front of one. Pruning +changes the request, so if a provider rejects a request for a reason the plugin +cannot see, `on_error: observe` is how you find out safely: it counts what it +would remove while sending the bytes untouched. + +Three specifics worth knowing: + +- **A forced `tool_choice` is never pruned.** `tool_choice: {"type":"tool", + "name":"X"}` (or OpenAI's `{"type":"function","function":{"name":"X"}}`) makes + `X` mandatory; a `tool_choice` naming a tool absent from the manifest is an + invalid request. `X` is kept even when the remove list names it, and the rest + of the list still applies. - **Nothing else in the request changes.** Deletions are surgical: every byte outside the removed array elements is preserved, including key order and diff --git a/docs/proposals/tool-prune.md b/docs/proposals/tool-prune.md index f5f80d284..58ae52c45 100644 --- a/docs/proposals/tool-prune.md +++ b/docs/proposals/tool-prune.md @@ -233,8 +233,15 @@ On each outbound request: Every byte outside the deleted array elements is unchanged. `gjson`/`sjson` are already in `authlib/go.mod` (currently indirect), so no new dependency. -Any error or panic fails open: the original body is forwarded unmodified. A -cost optimisation must never be able to break a request. +Any error or panic fails open: the original body is forwarded unmodified, so the +plugin's own failure modes cannot break a request. That is a narrower promise +than "pruning is always safe": pruning changes the request, and whether a +provider or gateway accepts a validly pruned manifest is outside what the plugin +can see. `on_error: observe` is how that is established safely. + +A forced `tool_choice` is the one case where the manifest and another field must +agree, so a tool named by `tool_choice` is never removed regardless of the +configured list. ### Configuration From 48476887f38198ea263b56e4d960668b035f33e0 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 14:06:09 -0400 Subject: [PATCH 08/28] feat(authbridge): Stamp a request id on session events MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Session events carried nothing tying a response back to its request, so abctl paired them positionally: a request row was bracketed with whatever response row followed it. The code said so — "a server-side correlation id would be the fix if that ever bites." It bit, and expensively. Claude Code issues its session-title request concurrently with the main one, both POSTs to the same host. Interleaved as req(title), req(main), resp(title, 400), the heuristic walks back from the 400 to the nearest unpaired request — the main one — and brackets them together. The result was a 400 rendered directly beneath the row where tool-prune reported rewriting a body, which reads unambiguously as the plugin having broken that request. It had not: every request it modified returned 200, and every 400 belonged to a title request carrying no tool manifest that the plugin skipped outright. Hours went into disproving a defect the display had invented. Context.RequestID() generates a short per-request id on first use, and all four listeners stamp it on both the request and response event. Generated lazily rather than as a constructor argument because there are ten Context construction sites and an eleventh required field would be a standing trap; Contexts are single-goroutine by contract, so the lazy write needs no synchronisation. abctl pairs on it exactly and keeps the adjacency heuristic only for events without one, so an older data plane still renders brackets. The regression test uses the real interleaving from the session store and fails against the heuristic alone. Signed-off-by: Hai Huang --- authbridge/authlib/listener/extproc/server.go | 6 ++ .../authlib/listener/forwardproxy/server.go | 3 + .../listener/forwardproxy/transparent.go | 1 + .../authlib/listener/reverseproxy/server.go | 4 ++ authbridge/authlib/pipeline/context.go | 4 ++ authbridge/authlib/pipeline/requestid.go | 48 +++++++++++++++ authbridge/authlib/pipeline/session.go | 7 +++ authbridge/cmd/abctl/tui/events_pane.go | 49 +++++++++++++-- authbridge/cmd/abctl/tui/events_pane_test.go | 59 +++++++++++++++++++ 9 files changed, 177 insertions(+), 4 deletions(-) create mode 100644 authbridge/authlib/pipeline/requestid.go diff --git a/authbridge/authlib/listener/extproc/server.go b/authbridge/authlib/listener/extproc/server.go index 3cfe5ed67..ed08d2c24 100644 --- a/authbridge/authlib/listener/extproc/server.go +++ b/authbridge/authlib/listener/extproc/server.go @@ -241,6 +241,7 @@ func (s *Server) recordInboundSession(pctx *pipeline.Context) { At: time.Now(), Direction: pipeline.Inbound, Phase: pipeline.SessionRequest, + RequestID: pctx.RequestID(), A2A: pipeline.SnapshotA2A(pctx.Extensions.A2A), Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest), Plugins: plugins, @@ -279,6 +280,7 @@ func (s *Server) recordInboundReject(pctx *pipeline.Context, action pipeline.Act At: time.Now(), Direction: pipeline.Inbound, Phase: pipeline.SessionDenied, + RequestID: pctx.RequestID(), Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest), Plugins: pipeline.SnapshotPlugins(pctx.Extensions.Custom), Identity: pipeline.SnapshotIdentity(pctx), @@ -332,6 +334,7 @@ func (s *Server) recordOutboundReject(pctx *pipeline.Context, action pipeline.Ac At: time.Now(), Direction: pipeline.Outbound, Phase: pipeline.SessionDenied, + RequestID: pctx.RequestID(), Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest), Plugins: pipeline.SnapshotPlugins(pctx.Extensions.Custom), Identity: pipeline.SnapshotIdentity(pctx), @@ -371,6 +374,7 @@ func (s *Server) recordInboundResponseSession(pctx *pipeline.Context) { At: time.Now(), Direction: pipeline.Inbound, Phase: pipeline.SessionResponse, + RequestID: pctx.RequestID(), A2A: pipeline.SnapshotA2A(pctx.Extensions.A2A), Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseResponse), Plugins: plugins, @@ -399,6 +403,7 @@ func (s *Server) recordOutboundResponseSession(pctx *pipeline.Context) { At: time.Now(), Direction: pipeline.Outbound, Phase: pipeline.SessionResponse, + RequestID: pctx.RequestID(), MCP: pipeline.SnapshotMCP(pctx.Extensions.MCP), Inference: pipeline.SnapshotInference(pctx.Extensions.Inference), Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseResponse), @@ -446,6 +451,7 @@ func (s *Server) recordOutboundSession(pctx *pipeline.Context) { At: time.Now(), Direction: pipeline.Outbound, Phase: pipeline.SessionRequest, + RequestID: pctx.RequestID(), MCP: pipeline.SnapshotMCP(pctx.Extensions.MCP), Inference: pipeline.SnapshotInference(pctx.Extensions.Inference), Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest), diff --git a/authbridge/authlib/listener/forwardproxy/server.go b/authbridge/authlib/listener/forwardproxy/server.go index 9d822be93..88a8ba1b7 100644 --- a/authbridge/authlib/listener/forwardproxy/server.go +++ b/authbridge/authlib/listener/forwardproxy/server.go @@ -305,6 +305,7 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, isBridge At: time.Now(), Direction: pipeline.Outbound, Phase: pipeline.SessionRequest, + RequestID: pctx.RequestID(), MCP: pipeline.SnapshotMCP(pctx.Extensions.MCP), Inference: pipeline.SnapshotInference(pctx.Extensions.Inference), Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest), @@ -570,6 +571,7 @@ func (s *Server) recordOutboundResponseEvent(pctx *pipeline.Context, statusCode At: time.Now(), Direction: pipeline.Outbound, Phase: pipeline.SessionResponse, + RequestID: pctx.RequestID(), MCP: pipeline.SnapshotMCP(pctx.Extensions.MCP), Inference: pipeline.SnapshotInference(pctx.Extensions.Inference), Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseResponse), @@ -882,6 +884,7 @@ func (s *Server) recordOutboundReject(pctx *pipeline.Context, action pipeline.Ac At: time.Now(), Direction: pipeline.Outbound, Phase: pipeline.SessionDenied, + RequestID: pctx.RequestID(), Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest), Host: pctx.Host, StatusCode: status, diff --git a/authbridge/authlib/listener/forwardproxy/transparent.go b/authbridge/authlib/listener/forwardproxy/transparent.go index 6941c836a..447a5478c 100644 --- a/authbridge/authlib/listener/forwardproxy/transparent.go +++ b/authbridge/authlib/listener/forwardproxy/transparent.go @@ -175,6 +175,7 @@ func (s *Server) recordTunnelOpened(pctx *pipeline.Context) { At: time.Now(), Direction: pipeline.Outbound, Phase: pipeline.SessionRequest, + RequestID: pctx.RequestID(), Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest), Plugins: plugins, Identity: pipeline.SnapshotIdentity(pctx), diff --git a/authbridge/authlib/listener/reverseproxy/server.go b/authbridge/authlib/listener/reverseproxy/server.go index 81b4d3988..62d89cb07 100644 --- a/authbridge/authlib/listener/reverseproxy/server.go +++ b/authbridge/authlib/listener/reverseproxy/server.go @@ -410,6 +410,7 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) { At: time.Now(), Direction: pipeline.Inbound, Phase: pipeline.SessionRequest, + RequestID: pctx.RequestID(), A2A: pipeline.SnapshotA2A(pctx.Extensions.A2A), Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest), Plugins: plugins, @@ -529,6 +530,7 @@ func (s *Server) modifyResponse(resp *http.Response) error { At: time.Now(), Direction: pipeline.Inbound, Phase: pipeline.SessionResponse, + RequestID: pctx.RequestID(), A2A: pipeline.SnapshotA2A(pctx.Extensions.A2A), Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseResponse), Plugins: plugins, @@ -583,6 +585,7 @@ func (s *Server) recordInboundReject(pctx *pipeline.Context, action pipeline.Act At: time.Now(), Direction: pipeline.Inbound, Phase: pipeline.SessionDenied, + RequestID: pctx.RequestID(), Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest), Host: pctx.Host, StatusCode: status, @@ -671,6 +674,7 @@ func (s *Server) recordInboundResponseEvent(pctx *pipeline.Context, statusCode i At: time.Now(), Direction: pipeline.Inbound, Phase: pipeline.SessionResponse, + RequestID: pctx.RequestID(), A2A: pipeline.SnapshotA2A(pctx.Extensions.A2A), Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseResponse), Plugins: plugins, diff --git a/authbridge/authlib/pipeline/context.go b/authbridge/authlib/pipeline/context.go index 4d93c9170..78182dd45 100644 --- a/authbridge/authlib/pipeline/context.go +++ b/authbridge/authlib/pipeline/context.go @@ -115,6 +115,10 @@ type Context struct { // compute SessionEvent.Duration without walking the event history. StartedAt time.Time + // requestID backs RequestID(), which generates it on first use. See + // requestid.go for why it is lazy rather than a constructor argument. + requestID string + Agent *AgentIdentity Identity Identity // nil before an auth plugin runs Session *SessionView // nil unless session tracking is enabled diff --git a/authbridge/authlib/pipeline/requestid.go b/authbridge/authlib/pipeline/requestid.go new file mode 100644 index 000000000..395dbd097 --- /dev/null +++ b/authbridge/authlib/pipeline/requestid.go @@ -0,0 +1,48 @@ +package pipeline + +import ( + "crypto/rand" + "encoding/hex" + "strconv" + "sync/atomic" +) + +// requestIDCounter is the fallback when crypto/rand is unavailable, so an id is +// always produced rather than an empty string that would silently disable +// pairing. +var requestIDCounter atomic.Uint64 + +// newRequestID returns a short, unique-per-process request identifier. +// +// Not a UUID on purpose: it exists to pair a request event with its response +// event in a session timeline, so it needs to be unique among in-flight +// requests and short enough to read in a terminal — not globally unique or +// cryptographically meaningful. +func newRequestID() string { + var b [6]byte + if _, err := rand.Read(b[:]); err != nil { + return "r" + strconv.FormatUint(requestIDCounter.Add(1), 36) + } + return hex.EncodeToString(b[:]) +} + +// RequestID returns a stable identifier for this request, generated on first +// use. Session events carry it so a consumer can pair a request event with its +// response event. +// +// Without it, pairing is positional — a UI matches a request row to whatever +// response row follows it — which silently misattributes whenever a client has +// more than one request in flight. That produced a real misdiagnosis: a plugin +// was blamed for a 400 that belonged to a concurrent request it never touched. +// +// Generated lazily rather than at Context construction so no listener can forget +// it; there are several construction sites and adding one more required field +// would be a standing trap. Contexts are single-goroutine by contract (plugins +// mutate Body, Headers and Extensions without locks), so the lazy write needs no +// synchronisation. +func (c *Context) RequestID() string { + if c.requestID == "" { + c.requestID = newRequestID() + } + return c.requestID +} diff --git a/authbridge/authlib/pipeline/session.go b/authbridge/authlib/pipeline/session.go index fa853ecbf..1ec1a94f5 100644 --- a/authbridge/authlib/pipeline/session.go +++ b/authbridge/authlib/pipeline/session.go @@ -83,6 +83,10 @@ type SessionEvent struct { At time.Time Direction Direction Phase SessionPhase + // RequestID pairs a request event with its response event. Without it a + // consumer can only pair positionally, which misattributes whenever a + // client has concurrent requests in flight. + RequestID string A2A *A2AExtension MCP *MCPExtension Inference *InferenceExtension @@ -212,6 +216,7 @@ type sessionEventWire struct { At time.Time `json:"at"` Direction Direction `json:"direction"` Phase SessionPhase `json:"phase"` + RequestID string `json:"requestId,omitempty"` A2A *A2AExtension `json:"a2a,omitempty"` MCP *MCPExtension `json:"mcp,omitempty"` Inference *InferenceExtension `json:"inference,omitempty"` @@ -232,6 +237,7 @@ func (e SessionEvent) MarshalJSON() ([]byte, error) { At: e.At, Direction: e.Direction, Phase: e.Phase, + RequestID: e.RequestID, A2A: e.A2A, MCP: e.MCP, Inference: e.Inference, @@ -260,6 +266,7 @@ func (e *SessionEvent) UnmarshalJSON(data []byte) error { At: w.At, Direction: w.Direction, Phase: w.Phase, + RequestID: w.RequestID, A2A: w.A2A, MCP: w.MCP, Inference: w.Inference, diff --git a/authbridge/cmd/abctl/tui/events_pane.go b/authbridge/cmd/abctl/tui/events_pane.go index 1fe21e487..556385288 100644 --- a/authbridge/cmd/abctl/tui/events_pane.go +++ b/authbridge/cmd/abctl/tui/events_pane.go @@ -477,21 +477,62 @@ func truncStr(s string, n int) string { // never gets a response) from stealing a later response that belongs to a // different method. // -// Closest-preceding adjacency is sufficient for current traffic, where a -// response follows its request. Concurrent same-host+method calls could in -// principle cross-pair, but this is a navigational cue, not a correctness -// guarantee; a server-side correlation id would be the fix if that ever bites. +// Pairing prefers SessionEvent.RequestID, which the proxy stamps on both the +// request and response event of the same exchange. That is exact, including +// under concurrency. +// +// The closest-preceding heuristic below remains for events with no RequestID — +// an older proxy, or a listener that has not been taught to stamp it. It matches +// on direction + host (port-normalized) + method, and it cross-pairs when a +// client has concurrent same-host+method calls in flight. That is not +// hypothetical: Claude Code fires its session-title request alongside the main +// one, and the heuristic drew a 400 from the title request under the main +// request's row, which read as the pipeline plugin on that row having caused it. // // IDs are keyed by event pointer so the render loop can look one up without // knowing the row index. They start at 1 and increment in first-seen row order // so adjacent exchanges get adjacent integers. func computeEventPairs(rows []eventRow) (map[*pipeline.SessionEvent]int, map[int]int) { partner := make(map[int]int) // row index → matched row index + + // Exact pass: pair by the proxy-stamped RequestID. Indexed by id so a + // response finds its request regardless of how much traffic interleaves + // between them. + reqByID := make(map[string]int) + for i := range rows { + e := rows[i].event + if e.RequestID == "" || e.Phase != pipeline.SessionRequest { + continue + } + if _, dup := reqByID[e.RequestID]; !dup { + reqByID[e.RequestID] = i + } + } + for j := range rows { + e := rows[j].event + if e.RequestID == "" || e.Phase != pipeline.SessionResponse { + continue + } + i, ok := reqByID[e.RequestID] + if !ok { + continue + } + if _, taken := partner[i]; taken { + continue + } + partner[i] = j + partner[j] = i + } + + // Heuristic pass: only for rows the exact pass could not place. for j := range rows { rj := rows[j].event if rj.Phase != pipeline.SessionResponse { continue } + if _, done := partner[j]; done { + continue // already paired exactly by RequestID + } for i := j - 1; i >= 0; i-- { if _, taken := partner[i]; taken { continue diff --git a/authbridge/cmd/abctl/tui/events_pane_test.go b/authbridge/cmd/abctl/tui/events_pane_test.go index c8d6b281b..34b857f57 100644 --- a/authbridge/cmd/abctl/tui/events_pane_test.go +++ b/authbridge/cmd/abctl/tui/events_pane_test.go @@ -822,3 +822,62 @@ func TestPlural(t *testing.T) { } } } + +// TestComputeEventPairs_RequestIDBeatsAdjacency reproduces a real +// misdiagnosis. Claude Code fires its session-title request concurrently with +// the main one; both are POSTs to the same host. Interleaved as +// req(main) req(title) resp(title,400) resp(main,200), the closest-preceding +// heuristic pairs req(title) with resp(title) — but pairs req(main) with +// resp(main) only by luck of ordering, and with a different interleaving it +// draws the title request's 400 under the main request's row. +// +// That is exactly what happened: a 400 belonging to a request tool-prune never +// touched was rendered beneath the row where tool-prune reported a body +// rewrite, which read as the plugin having broken the request. RequestID makes +// the pairing exact. +func TestComputeEventPairs_RequestIDBeatsAdjacency(t *testing.T) { + ev := func(phase pipeline.SessionPhase, id string, code int) *pipeline.SessionEvent { + return &pipeline.SessionEvent{ + Direction: pipeline.Outbound, + Phase: phase, + Host: "litellm.example", + RequestID: id, + StatusCode: code, + } + } + // The real interleaving observed in the session store: the title request + // is issued first, the main request second, and the title's 400 arrives + // before the main response. The heuristic then walks back from the 400 to + // the nearest unpaired request — the MAIN one — and brackets them together. + title := ev(pipeline.SessionRequest, "bbb", 0) + mainReq := ev(pipeline.SessionRequest, "aaa", 0) + titleResp := ev(pipeline.SessionResponse, "bbb", 400) + mainResp := ev(pipeline.SessionResponse, "aaa", 200) + rows := []eventRow{{event: title}, {event: mainReq}, {event: titleResp}, {event: mainResp}} + + _, partner := computeEventPairs(rows) + + if partner[0] != 2 { + t.Errorf("title request (row 0) paired with row %d, want 2 (its own 400)", partner[0]) + } + if partner[1] != 3 { + t.Errorf("main request (row 1) paired with row %d, want 3 (its own 200)", partner[1]) + } + // The specific failure this fixes: the main request owning the title's 400. + if partner[1] == 2 { + t.Error("main request paired with the title request's 400 — the misdiagnosis this fixes") + } +} + +// TestComputeEventPairs_FallsBackWithoutRequestID keeps the heuristic working +// for events from a proxy that does not stamp an id, so an older data plane +// still renders brackets. +func TestComputeEventPairs_FallsBackWithoutRequestID(t *testing.T) { + req := &pipeline.SessionEvent{Direction: pipeline.Outbound, Phase: pipeline.SessionRequest, Host: "h"} + resp := &pipeline.SessionEvent{Direction: pipeline.Outbound, Phase: pipeline.SessionResponse, Host: "h", StatusCode: 200} + rows := []eventRow{{event: req}, {event: resp}} + _, partner := computeEventPairs(rows) + if partner[0] != 1 || partner[1] != 0 { + t.Errorf("heuristic pairing broke for id-less events: partner=%v", partner) + } +} From 9340821d27053e3e116c0d29ecbfd23b647fa94b Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 14:06:09 -0400 Subject: [PATCH 09/28] docs(authbridge): Correct stale body-capability references MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The directional split left several places describing the old single flag, which is worse than no documentation because it reads as current: - framework-architecture.md still showed WritesRequestBody permitting pctx.SetResponseBody in five places, and listed a deprecated BodyAccess field that no longer exists on the struct at all. Updated to the two directional flags; the BodyAccess mentions that remain are changelog entries, accurate as history. - reverseproxy's streaming comment said WritesRequestBody is incompatible with streaming, directly beside the check that now reads WritesResponseBody — the exact inversion the split fixes. Also documents a known gap the split makes visible rather than introducing. Reader-ordering is validated in list order, which is request order; RunResponse iterates in reverse, so on the response pass the rule inverts and a reader needs to sit after a WritesResponseBody plugin. The two rules conflict for a both-direction mutator whenever a body reader is present, so no single ordering satisfies both. It does not bite in-tree because RunResponse skips StreamingResponders and every body-reading parser is one; a non-streaming reader (opa, ibac) before a response mutator would see rewritten bytes. Deliberately not enforced: the check would reject chains that validate today, and this change promised that no working configuration starts failing. Closing it needs direction-specific read capabilities, which is its own compatibility review. Signed-off-by: Hai Huang --- .../authlib/listener/reverseproxy/server.go | 2 +- authbridge/authlib/pipeline/pipeline.go | 20 ++++++++++++++++ authbridge/docs/framework-architecture.md | 23 +++++++++---------- authbridge/docs/plugin-reference.md | 11 +++++++++ 4 files changed, 43 insertions(+), 13 deletions(-) diff --git a/authbridge/authlib/listener/reverseproxy/server.go b/authbridge/authlib/listener/reverseproxy/server.go index 62d89cb07..a58f7519e 100644 --- a/authbridge/authlib/listener/reverseproxy/server.go +++ b/authbridge/authlib/listener/reverseproxy/server.go @@ -441,7 +441,7 @@ func (s *Server) modifyResponse(resp *http.Response) error { // called on this path — streaming-aware plugins finalize via // OnResponseFrame(last=true). // - // WritesRequestBody is incompatible with streaming (we can't rewrite a + // WritesResponseBody is incompatible with streaming (we can't rewrite a // body we've already started forwarding) — fall back to buffered // with a warning. if isEventStream(resp.Header.Get("Content-Type")) && diff --git a/authbridge/authlib/pipeline/pipeline.go b/authbridge/authlib/pipeline/pipeline.go index 496b41fd8..e77542fc6 100644 --- a/authbridge/authlib/pipeline/pipeline.go +++ b/authbridge/authlib/pipeline/pipeline.go @@ -597,6 +597,26 @@ func validateCapabilities(plugins []Plugin) error { } // Reader-ordering is triggered by either write flag: a reader placed // after any mutator would no longer see the original bytes. + // + // KNOWN GAP, response direction. This check is in list order, which is + // request order. RunResponse iterates in reverse, so on the response + // pass the rule inverts: a reader must appear AFTER a + // WritesResponseBody plugin to see original response bytes. The two + // rules therefore conflict for a plugin that writes both directions + // (sparc, cpex) whenever a body reader is in the chain — no single + // ordering satisfies both. + // + // It does not bite in-tree today because RunResponse skips + // StreamingResponders, and every body-reading parser (inference-, + // a2a-, mcp-parser) is one. A non-streaming reader (opa, ibac) placed + // before a response mutator would genuinely see rewritten bytes. + // + // Deliberately not enforced here: adding the reverse-order check would + // reject chains that validate today (e.g. [opa, sparc]), and the + // directional-capability change promised that no working configuration + // starts failing. Closing it needs direction-specific READ capabilities + // so the two passes can be validated independently, which is its own + // compatibility review. if caps.ReadsBody && firstMutator != "" && readerAfterMutator == "" { readerAfterMutator = plugin.Name() } diff --git a/authbridge/docs/framework-architecture.md b/authbridge/docs/framework-architecture.md index 5c0d3ced8..2f0889ad5 100644 --- a/authbridge/docs/framework-architecture.md +++ b/authbridge/docs/framework-architecture.md @@ -74,8 +74,8 @@ type PluginCapabilities struct { Reads []string // extension slot names this plugin reads Writes []string // extension slot names this plugin writes ReadsBody bool // plugin reads pctx.Body / pctx.ResponseBody - WritesRequestBody bool // plugin mutates body via pctx.SetBody / pctx.SetResponseBody - BodyAccess bool // deprecated: alias for ReadsBody (folded by Normalize) + WritesRequestBody bool // plugin mutates the request body via pctx.SetBody + WritesResponseBody bool // plugin mutates the response body via pctx.SetResponseBody } ``` @@ -85,9 +85,9 @@ Declared once per plugin instance. `pipeline.New` validates that every `Read` is plugin "guardrail" reads slot "mcp" but no earlier plugin writes it ``` -`ReadsBody: true` (or the legacy `BodyAccess` alias) on *any* plugin in a chain causes `Pipeline.NeedsBody()` to return true, which the **listener** uses to negotiate Envoy's `ProcessingMode` (BUFFERED vs HEADERS-only). Without this, the gRPC ext_proc server never asks for the body and parsers see `pctx.Body == nil`. +`ReadsBody: true` on *any* plugin in a chain causes `Pipeline.NeedsBody()` to return true, which the **listener** uses to negotiate Envoy's `ProcessingMode` (BUFFERED vs HEADERS-only). Without this, the gRPC ext_proc server never asks for the body and parsers see `pctx.Body == nil`. -`WritesRequestBody: true` declares that the plugin may rewrite the body via `pctx.SetBody` / `pctx.SetResponseBody`; the listener propagates the mutation to the wire. See §6, "Body mutation" for the full body-mutation contract (capability rules, ordering constraints, Content-Encoding policy). +`WritesRequestBody: true` declares that the plugin may rewrite the request body via `pctx.SetBody`; `WritesResponseBody: true` declares the response side via `pctx.SetResponseBody`. The listener propagates the mutation to the wire, and only `WritesResponseBody` forfeits incremental SSE relay. See §6, "Body mutation" for the full body-mutation contract (capability rules, ordering constraints, Content-Encoding policy). ### `OnRequest(ctx, pctx) Action` Called when a request is entering the pipeline. Plugins typically read request headers / body, mutate one or more extension slots, and return `Continue` or `Reject`. @@ -111,7 +111,7 @@ type Context struct { Host string // :authority / Host Path string // :path Headers http.Header - Body []byte // nil unless a plugin declared BodyAccess: true + Body []byte // nil unless a plugin declared ReadsBody: true StartedAt time.Time // listener wall-clock at request entry Agent *AgentIdentity // this workload's SPIFFE / Keycloak identity @@ -130,8 +130,8 @@ type Context struct { **Ownership rules:** - Plugins **read** any field they declared in `Capabilities.Reads`. - Plugins **write** fields they declared in `Capabilities.Writes`. By convention each extension slot has exactly one writer (the parser plugin). -- Plugins read `pctx.Body` / `pctx.ResponseBody` only if they declared `ReadsBody: true` (or the deprecated `BodyAccess: true`). -- Plugins mutate body content via `pctx.SetBody(newBytes)` / `pctx.SetResponseBody(newBytes)`, and only if they declared `WritesRequestBody: true`. Direct assignment (`pctx.Body = ...`) compiles but bypasses listener propagation and misses the Invocation + body-mutation event emission — see §6, "Body mutation." +- Plugins read `pctx.Body` / `pctx.ResponseBody` only if they declared `ReadsBody: true`. +- Plugins mutate body content via `pctx.SetBody(newBytes)` if they declared `WritesRequestBody: true`, or `pctx.SetResponseBody(newBytes)` if they declared `WritesResponseBody: true`. Direct assignment (`pctx.Body = ...`) compiles but bypasses listener propagation and misses the Invocation + body-mutation event emission — see §6, "Body mutation." - `Identity` is populated by whichever auth plugin ran (jwt-validation ships a `claimsIdentity` adapter around `validation.Claims`; a SAML / mTLS / custom plugin publishes its own adapter). The framework reads it through the `Identity` interface (`Subject()` / `ClientID()` / `Scopes()`) so no plugin-specific type leaks into `pipeline/`. - `Agent`, `Session` are populated by the listener before `Run`. Plugins treat them as read-only. - `ResponseBody` appears between `Run` and `RunResponse` — plugins must not read it in `OnRequest`. @@ -417,7 +417,7 @@ func (p *Pipeline) RunFinish(ctx context.Context, pctx *Context, outcome Outcome func (p *Pipeline) Start(ctx context.Context) error // invoke Init on Initializer plugins func (p *Pipeline) Stop(ctx context.Context) // invoke Shutdown on Shutdowner plugins func (p *Pipeline) Plugins() []Plugin // defensive copy -func (p *Pipeline) NeedsBody() bool // OR over all plugins' BodyAccess +func (p *Pipeline) NeedsBody() bool // OR over ReadsBody + both write flags ``` `New` validates capability wiring at startup: every `Read` must be satisfied by some earlier plugin's `Write`. `plugins.Build` additionally validates the cross-plugin relationship declarations — `Requires`, `RequiresAny`, `After`, `Claims` — before returning the pipeline to the listener. See [`plugin-reference.md` "Declaring plugin relationships"](./plugin-reference.md#declaring-plugin-relationships). @@ -584,7 +584,7 @@ Always sequential. No priority / mode / fire-and-forget semantics yet. This is t ### Body mutation -A plugin that declares `WritesRequestBody: true` may rewrite the request or response body. The framework owns the propagation to the wire; plugins only call `pctx.SetBody(newBytes)` / `pctx.SetResponseBody(newBytes)`. +A plugin declares the direction it rewrites: `WritesRequestBody: true` for the request body (`pctx.SetBody`), `WritesResponseBody: true` for the response body (`pctx.SetResponseBody`). The framework owns the propagation to the wire; plugins only call the helper. **Capability model.** Body access is declared per direction on `PluginCapabilities`: @@ -593,7 +593,6 @@ A plugin that declares `WritesRequestBody: true` may rewrite the request or resp | `ReadsBody` | plugin reads `pctx.Body` / `pctx.ResponseBody` | buffers the body; plugin sees the bytes | | `WritesRequestBody` | plugin may call `pctx.SetBody` | implies `ReadsBody`; propagates request mutations | | `WritesResponseBody` | plugin may call `pctx.SetResponseBody` | implies `ReadsBody`; propagates response mutations **and forces the buffered response path** | -| `BodyAccess` (deprecated) | legacy alias for `ReadsBody` | folded by `Normalize()`, removed in a future release | **Why the directions are separate.** `Pipeline.WritesResponseBody()` is the SSE streaming predicate: both proxy listeners consult it to decide whether a @@ -833,9 +832,9 @@ Breaking changes will be announced in `authbridge/CHANGELOG.md` (TBD) before a 1 **Package sources:** -- `pipeline.go` — `Pipeline` type, `New`, `Run`, `RunResponse`, `Start`, `Stop`, `Plugins`, `NeedsBody`, `WritesRequestBody`. +- `pipeline.go` — `Pipeline` type, `New`, `Run`, `RunResponse`, `Start`, `Stop`, `Plugins`, `NeedsBody`, `WritesRequestBody`, `WritesResponseBody`. - `holder.go` — `Holder`, the atomic slot listeners hold in place of a raw `*Pipeline`. -- `plugin.go` — `Plugin` interface, `PluginCapabilities` (with `ReadsBody` / `WritesRequestBody` / deprecated `BodyAccess` + `Normalize()`; chain-scoped relationship fields `Requires` / `RequiresAny` / `After` / `Claims`), `Configurable`, `Initializer`, `Shutdowner`, `Readier`, `Finisher`. +- `plugin.go` — `Plugin` interface, `PluginCapabilities` (with `ReadsBody` / `WritesRequestBody` / `WritesResponseBody` + `Normalize()`; chain-scoped relationship fields `Requires` / `RequiresAny` / `After` / `Claims`), `Configurable`, `Initializer`, `Shutdowner`, `Readier`, `Finisher`. - `outcome.go` — `Outcome` struct + `OutcomeAction` (allow / deny / error) for `Finisher` consumers; `Context.Outcome()` getter. - `action.go` — `Action`, `ActionType`, `Violation`, helper constructors (`Deny`, `DenyStatus`, `DenyWithDetails`, `Challenge`, `RateLimited`), `StatusFromCode`. - `context.go` — `Context`, `Direction`, `AgentIdentity`, the `pctx.Record` / `Allow` / `Skip` / `Observe` / `Modify` / `DenyAndRecord` helpers, and `pctx.SetBody` / `SetResponseBody` / `BodyMutated` / `ResponseBodyMutated` for body mutation. diff --git a/authbridge/docs/plugin-reference.md b/authbridge/docs/plugin-reference.md index d564db001..f35266fe9 100644 --- a/authbridge/docs/plugin-reference.md +++ b/authbridge/docs/plugin-reference.md @@ -778,6 +778,17 @@ nothing about how the response may be relayed. the ext_authz API has no body-mutation field. Do not combine body-mutating plugins with `mode: waypoint`. +> **Reader-ordering is validated in request order only.** `RunResponse` iterates +> the chain in reverse, so on the response pass the rule inverts — a reader needs +> to sit *after* a `WritesResponseBody` plugin to see original response bytes. +> The two rules conflict for a both-direction mutator whenever a body reader is +> present, so no single ordering satisfies both. In practice this is invisible +> in-tree: `RunResponse` skips `StreamingResponder`s and every body-reading +> parser is one. A non-streaming reader (`opa`, `ibac`) placed before a response +> mutator would see rewritten bytes. Not enforced, because the check would reject +> chains that validate today; closing it needs direction-specific *read* +> capabilities. + > **Declaring is a contract, not an enforcement.** `SetBody` flips > `bodyMutated` unconditionally outside observe mode and the listeners gate > purely on that flag, so a plugin that calls `SetBody` *without* declaring the From 4ccbaa2276c001c7c10af65bf4fd9be2b6900612 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 14:06:09 -0400 Subject: [PATCH 10/28] feat(authbridge): Surface the provider's error type on 4xx/5xx MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A rejected upstream request recorded {"kind":"backend_error","code":"400"} and nothing else, so the session timeline showed that something failed but never why. Debugging one meant reproducing it outside the proxy. DeriveError now reads the provider's own machine-readable classification out of the error body already buffered on that path: error.type, falling back to error.code. That is what an operator acts on — invalid_request_error means fix the request, rate_limit_error means back off, authentication_error means fix credentials. The human-readable error.message is deliberately excluded. Provider messages routinely quote the offending part of the request, and the session store is unauthenticated — the same reason body-mutation events carry only length and sha256. type and code are enum-like: bounded vocabularies chosen by the provider, carrying no request content. A test asserts a credential embedded in a provider message does not reach the event. Costs no new body reads: it uses what is already buffered, bounds the parse, and returns empty for anything that is not a JSON error document, so an HTML 502 or a truncated body still yields the bare event rather than noise. Signed-off-by: Hai Huang --- authbridge/authlib/pipeline/errorkind_test.go | 77 +++++++++++++++++++ authbridge/authlib/pipeline/snapshot.go | 46 ++++++++++- 2 files changed, 121 insertions(+), 2 deletions(-) create mode 100644 authbridge/authlib/pipeline/errorkind_test.go diff --git a/authbridge/authlib/pipeline/errorkind_test.go b/authbridge/authlib/pipeline/errorkind_test.go new file mode 100644 index 000000000..8e95994a9 --- /dev/null +++ b/authbridge/authlib/pipeline/errorkind_test.go @@ -0,0 +1,77 @@ +package pipeline + +import "testing" + +// TestUpstreamErrorKind: a bare "backend_error / 400" gives an operator nothing +// to act on. The provider's own classification does — and it must be the +// classification only, never the human message, which quotes request content +// into an unauthenticated store. +func TestUpstreamErrorKind(t *testing.T) { + tests := []struct { + name string + body string + want string + }{ + { + name: "anthropic error type", + body: `{"type":"error","error":{"type":"invalid_request_error","message":"tools.3: unexpected"}}`, + want: "invalid_request_error", + }, + { + name: "openai style falls back to code", + body: `{"error":{"message":"bad","code":"context_length_exceeded"}}`, + want: "context_length_exceeded", + }, + {"type preferred over code", `{"error":{"type":"rate_limit_error","code":"429"}}`, "rate_limit_error"}, + {"no error object", `{"ok":true}`, ""}, + {"malformed json", `{"error":{"type":`, ""}, + {"empty body", ``, ""}, + {"not json at all", `502 Bad Gateway`, ""}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := upstreamErrorKind([]byte(tc.body)); got != tc.want { + t.Errorf("upstreamErrorKind() = %q, want %q", got, tc.want) + } + }) + } +} + +// TestUpstreamErrorKind_NeverLeaksTheMessage is the privacy assertion: the +// provider's prose can quote the request, so it must never reach the event. +func TestUpstreamErrorKind_NeverLeaksTheMessage(t *testing.T) { + secret := "sk-live-abcdef123456" + body := `{"error":{"type":"authentication_error","message":"invalid key ` + secret + `"}}` + got := upstreamErrorKind([]byte(body)) + if got != "authentication_error" { + t.Fatalf("got %q, want the type", got) + } + if got == secret || len(got) > 64 { + t.Errorf("message content leaked into the event: %q", got) + } +} + +// TestDeriveError_PopulatesKindFrom4xxBody wires it to the event an operator +// actually reads. +func TestDeriveError_PopulatesKindFrom4xxBody(t *testing.T) { + pctx := &Context{ + StatusCode: 400, + ResponseBody: []byte(`{"type":"error","error":{"type":"invalid_request_error","message":"x"}}`), + } + e := DeriveError(pctx) + if e == nil { + t.Fatal("expected an error event for a 400") + } + if e.Kind != "backend_error" || e.Code != "400" { + t.Errorf("kind/code = %q/%q", e.Kind, e.Code) + } + if e.Message != "invalid_request_error" { + t.Errorf("Message = %q, want the provider's error type", e.Message) + } + // A 4xx with no parseable body must still produce the event, just without + // a classification — never an error swallowed for lack of a body. + bare := DeriveError(&Context{StatusCode: 503}) + if bare == nil || bare.Code != "503" || bare.Message != "" { + t.Errorf("bare 5xx = %+v, want backend_error/503 with empty message", bare) + } +} diff --git a/authbridge/authlib/pipeline/snapshot.go b/authbridge/authlib/pipeline/snapshot.go index 079cc36d0..693421422 100644 --- a/authbridge/authlib/pipeline/snapshot.go +++ b/authbridge/authlib/pipeline/snapshot.go @@ -1,6 +1,8 @@ package pipeline import ( + "github.com/tidwall/gjson" + "encoding/json" "log/slog" "strconv" @@ -129,9 +131,49 @@ func DeriveError(pctx *Context) *EventError { } if pctx.StatusCode >= 400 { return &EventError{ - Kind: "backend_error", - Code: strconv.Itoa(pctx.StatusCode), + Kind: "backend_error", + Code: strconv.Itoa(pctx.StatusCode), + Message: upstreamErrorKind(pctx.ResponseBody), } } return nil } + +// upstreamErrorKind extracts the provider's machine-readable error type from an +// error response body, or "" when there isn't one. +// +// A bare `backend_error / 400` tells an operator nothing about why, which turns +// every upstream rejection into a guessing exercise. The provider already +// classifies its own failures, and the classification is what an operator acts +// on: invalid_request_error means fix the request, rate_limit_error means back +// off, authentication_error means fix credentials. +// +// The human-readable error.message is deliberately NOT captured. Provider +// messages routinely quote the offending part of the request, and the session +// store is unauthenticated — the same reason body-mutation events carry only +// length and sha256. The type and code are enum-like: bounded vocabularies +// chosen by the provider, carrying no request content. +func upstreamErrorKind(body []byte) string { + if len(body) == 0 { + return "" + } + // Bound the parse: an error body is small, and a huge one here means this + // isn't an error document at all. + if len(body) > 64*1024 { + body = body[:64*1024] + } + if !gjson.ValidBytes(body) { + return "" + } + t := gjson.GetBytes(body, "error.type").String() + if t == "" { + t = gjson.GetBytes(body, "error.code").String() + } + if t == "" { + return "" + } + if len(t) > 64 { + t = t[:64] + } + return t +} From e9bd6144eb712e9424811db4976d5de434e1cc99 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 14:12:42 -0400 Subject: [PATCH 11/28] test(abctl): Pin event pairing against a real captured trace MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replays a 20-event interleaving captured from a live Claude Code session through the demo proxy. The adjacency heuristic mispaired 6 of its 15 responses: a three-way rotation across three concurrent litellm requests and a straight swap on a later pair. Ownership in the fixture is not guesswork. Each response event carries a duration measured from its own request's start, so subtracting it from the response timestamp identifies the true owning request independently of the RequestID the test exercises — which is how the mispairing was established in the first place. The assertion that matters is the last one: neither request tool-prune modified may own a 400. In the real trace both did on screen, and both actually returned 200 — the 400s belonged to concurrent requests the plugin never touched. That display artifact was read as the plugin breaking requests, so it is worth a test that fails loudly (12 assertions) if pairing ever regresses to adjacency. Signed-off-by: Hai Huang --- authbridge/cmd/abctl/tui/events_pane_test.go | 79 ++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/authbridge/cmd/abctl/tui/events_pane_test.go b/authbridge/cmd/abctl/tui/events_pane_test.go index 34b857f57..70d6092a3 100644 --- a/authbridge/cmd/abctl/tui/events_pane_test.go +++ b/authbridge/cmd/abctl/tui/events_pane_test.go @@ -881,3 +881,82 @@ func TestComputeEventPairs_FallsBackWithoutRequestID(t *testing.T) { t.Errorf("heuristic pairing broke for id-less events: partner=%v", partner) } } + +// TestComputeEventPairs_FieldTrace replays a real interleaving captured from a +// Claude Code session, where the adjacency heuristic mispaired 6 of 15 +// responses — a 3-way rotation (rows 10/11/12) and a straight swap (20/21). +// +// The mispairing was not cosmetic. It rendered a 400 beneath every row where +// tool-prune reported rewriting a body, when each of those 400s belonged to a +// different concurrent request and every request tool-prune touched returned +// 200. Ownership here is not guesswork: each response's duration is measured +// from its own request's start, so subtracting it identifies the true owner +// independently of the id being tested. +func TestComputeEventPairs_FieldTrace(t *testing.T) { + type spec struct { + id string // true owning request id + phase pipeline.SessionPhase + host string + code int + } + // Order is wall-clock order as observed; ids are the true owners. + trace := []spec{ + {"r07", pipeline.SessionRequest, "mcp.ete", 0}, + {"r08", pipeline.SessionRequest, "mcp.ete", 0}, + {"r08", pipeline.SessionResponse, "mcp.ete", 200}, + {"r09", pipeline.SessionRequest, "litellm", 0}, + {"r09", pipeline.SessionResponse, "litellm", 200}, + {"r10", pipeline.SessionRequest, "litellm", 0}, + {"r11", pipeline.SessionRequest, "litellm", 0}, // tool-prune modified this one + {"r10", pipeline.SessionResponse, "litellm", 400}, + {"r12", pipeline.SessionRequest, "litellm", 0}, + {"r11", pipeline.SessionResponse, "litellm", 200}, // the modify's real outcome + {"r12", pipeline.SessionResponse, "litellm", 400}, + {"r13", pipeline.SessionRequest, "litellm", 0}, + {"r14", pipeline.SessionRequest, "litellm", 0}, // tool-prune modified this one + {"r07", pipeline.SessionResponse, "mcp.ete", 200}, + {"r13", pipeline.SessionResponse, "litellm", 400}, + {"r15", pipeline.SessionRequest, "litellm", 0}, + {"r16", pipeline.SessionRequest, "mcp.ete", 0}, + {"r16", pipeline.SessionResponse, "mcp.ete", 400}, + {"r15", pipeline.SessionResponse, "litellm", 400}, + {"r14", pipeline.SessionResponse, "litellm", 200}, // the modify's real outcome + } + + rows := make([]eventRow, 0, len(trace)) + for _, s := range trace { + rows = append(rows, eventRow{event: &pipeline.SessionEvent{ + Direction: pipeline.Outbound, Phase: s.phase, + Host: s.host, RequestID: s.id, StatusCode: s.code, + }}) + } + + _, partner := computeEventPairs(rows) + + for i, s := range trace { + j, ok := partner[i] + if !ok { + if s.id == "r07" || s.phase == pipeline.SessionRequest { + // every request in this trace does get a response + t.Errorf("row %d (%s %s) unpaired", i, s.id, s.phase) + } + continue + } + if got := rows[j].event.RequestID; got != s.id { + t.Errorf("row %d (%s) paired with %s — pairing crossed requests", i, s.id, got) + } + } + + // The specific regression: no tool-prune-modified request may own a 400. + for _, modified := range []string{"r11", "r14"} { + for i, s := range trace { + if s.phase != pipeline.SessionRequest || s.id != modified { + continue + } + j := partner[i] + if code := rows[j].event.StatusCode; code != 200 { + t.Errorf("%s (tool-prune modified) paired with a %d; its real response was 200", modified, code) + } + } + } +} From dc582137f6c32acc93ea48732a2d27fad8d06b63 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 14:21:48 -0400 Subject: [PATCH 12/28] refactor(authbridge): Make the remove list the only switch for tool-prune MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Enabling tool-prune took two steps: fill the remove list, then hand-edit on_error from observe to enforce. The second step was friction without much safety, because the two guards were never independent — enforce with an empty remove list is already a no-op. The list was always the real gate; the policy was belt-and-braces on top of it. So the demo config now ships on_error: enforce with an empty list, and `abctl tools scan --write` is the single deliberate act that turns the plugin on. The emitted YAML block drops the on_error line entirely, since "" already normalizes to enforce and printing a policy line implies it is the switch. observe is not removed and could not usefully be: it is a framework-wide policy in pipeline/policy.go that every plugin gets, and the plugin implements nothing for it beyond choosing which counter to increment. What changes is its billing — from a mandatory rollout stage to a deliberate instrument, documented for the two occasions it earns its keep: sizing the saving before it affects traffic, and clearing the plugin of suspicion when requests start failing. The second is worth keeping sharp; during development a display bug attributed another request's 400 to this plugin, and "set observe and see if the failures persist" is the cheapest way to settle that question. The demo test now pins only the empty list, not the policy. Asserting the policy would pin a default that is meant to be edited. Signed-off-by: Hai Huang --- authbridge/cmd/abctl/toolscan/scan.go | 4 ++- authbridge/cmd/abctl/toolscan/scan_test.go | 6 ++-- authbridge/cmd/authbridge-proxy/demo.go | 17 +++++---- authbridge/cmd/authbridge-proxy/demo_test.go | 11 +++--- authbridge/docs/tool-prune-plugin.md | 37 +++++++++++++------- authbridge/install-demo.sh | 5 ++- 6 files changed, 49 insertions(+), 31 deletions(-) diff --git a/authbridge/cmd/abctl/toolscan/scan.go b/authbridge/cmd/abctl/toolscan/scan.go index 96d688bf3..be8f78c9f 100644 --- a/authbridge/cmd/abctl/toolscan/scan.go +++ b/authbridge/cmd/abctl/toolscan/scan.go @@ -151,8 +151,10 @@ func scanFile(path string, since time.Time, seenIDs map[string]struct{}, res *Re // pastes (or --write patches) into the tool-prune entry. func (r *Result) YAMLBlock() string { var b strings.Builder + // on_error is omitted: it defaults to enforce, and the empty remove list + // below is what gates the plugin. Set on_error: observe only when you want + // a projection instead of a saving. b.WriteString(" - name: tool-prune\n") - b.WriteString(" on_error: observe # measure only; switch to enforce when trusted\n") b.WriteString(" config:\n") if len(r.Candidates) == 0 { b.WriteString(" remove: []\n") diff --git a/authbridge/cmd/abctl/toolscan/scan_test.go b/authbridge/cmd/abctl/toolscan/scan_test.go index 08e790138..2f8971ca5 100644 --- a/authbridge/cmd/abctl/toolscan/scan_test.go +++ b/authbridge/cmd/abctl/toolscan/scan_test.go @@ -198,8 +198,10 @@ func TestYAMLBlock(t *testing.T) { if !strings.Contains(got, "remove: [NotebookEdit, WebSearch]") { t.Errorf("block missing remove list:\n%s", got) } - if !strings.Contains(got, "on_error: observe") { - t.Errorf("emitted block must default to observe (measure first):\n%s", got) + // on_error is intentionally absent: it defaults to enforce, and the remove + // list is the gate. Emitting a policy line would imply it is the switch. + if strings.Contains(got, "on_error") { + t.Errorf("block should not emit an on_error line:\n%s", got) } empty := (&Result{}).YAMLBlock() if !strings.Contains(empty, "remove: []") { diff --git a/authbridge/cmd/authbridge-proxy/demo.go b/authbridge/cmd/authbridge-proxy/demo.go index 30724fccb..1b43fbaa9 100644 --- a/authbridge/cmd/authbridge-proxy/demo.go +++ b/authbridge/cmd/authbridge-proxy/demo.go @@ -50,16 +50,21 @@ pipeline: - name: mcp-parser - name: a2a-parser # tool-prune drops unused tool definitions from the outbound manifest. - # It ships inert: the remove list is empty, so it does nothing until you - # fill it in, and on_error: observe means even then it only measures -- - # counting what it *would* remove while the bytes on the wire stay - # untouched. Read the projection in abctl's plugin pane, then switch - # on_error to enforce once the numbers look right. Fill the list with: + # The empty remove list is the off switch: with nothing named it does + # nothing at all. Fill it in and it takes effect immediately -- # abctl tools scan --write + # -- and the config is hot-reloaded, so no restart. + # + # Watch the Metrics section of abctl's plugin pane for what it saved. If + # you ever suspect the plugin of breaking a request, set + # on_error: observe here: it then counts what it *would* remove while + # leaving every byte on the wire untouched, which settles the question + # without unconfiguring anything. + # # Keep it last: it rewrites the request body, and body readers must # precede the mutator so they see the original bytes. - name: tool-prune - on_error: observe + on_error: enforce config: remove: [] ` diff --git a/authbridge/cmd/authbridge-proxy/demo_test.go b/authbridge/cmd/authbridge-proxy/demo_test.go index c6fc6e028..296540bc2 100644 --- a/authbridge/cmd/authbridge-proxy/demo_test.go +++ b/authbridge/cmd/authbridge-proxy/demo_test.go @@ -81,9 +81,11 @@ func TestDemoConfig_WriteLoadsAndValidates(t *testing.T) { } // tool-prune ships inert, and that is a property worth pinning: the demo - // must never silently start rewriting a user's traffic. Two independent - // guards — an empty remove list (nothing to do) and observe policy - // (measure only) — so a future edit has to defeat both to enable it. + // must never silently start rewriting a user's traffic. The empty remove + // list is the guard — with no tool named there is nothing to remove, whatever + // the policy — so filling the list is the single, deliberate act that + // enables it. Asserting the policy too would just pin a default that is + // meant to be edited. var tp *config.PluginEntry for i := range cfg.Pipeline.Outbound.Plugins { if cfg.Pipeline.Outbound.Plugins[i].Name == "tool-prune" { @@ -93,9 +95,6 @@ func TestDemoConfig_WriteLoadsAndValidates(t *testing.T) { if tp == nil { t.Fatal("tool-prune entry not found") } - if tp.OnError != "observe" { - t.Errorf("tool-prune on_error = %q, want observe so the demo only measures", tp.OnError) - } if !strings.Contains(string(tp.Config), "\"remove\":[]") && !strings.Contains(string(tp.Config), "\"remove\": []") { t.Errorf("tool-prune must ship with an empty remove list, got %s", tp.Config) diff --git a/authbridge/docs/tool-prune-plugin.md b/authbridge/docs/tool-prune-plugin.md index c2ca5e2ac..606a64f9e 100644 --- a/authbridge/docs/tool-prune-plugin.md +++ b/authbridge/docs/tool-prune-plugin.md @@ -39,24 +39,35 @@ It declares `WritesRequestBody` only, never `WritesResponseBody`, so responses still stream incrementally. See [`plugin-reference.md`](./plugin-reference.md#capability-fields). -## Measure first, then enforce +## Turning it on -`on_error: observe` makes the plugin a projection: it computes exactly what it -would remove and counts it, while the bytes on the wire stay untouched. Nothing -about the plugin's code differs between the two modes — under observe, `SetBody` -is a no-op on bytes and leaves `BodyMutated()` false, which is how the plugin -knows which counter to increment. - -So the rollout is: add it in observe, read the projection, then flip one word. +**The empty `remove` list is the off switch.** With no tool named the plugin does +nothing, whatever the policy, so filling the list is the single act that enables +it: ```sh -abctl tools scan --write ./cortex-ca/demo.yaml # fill in remove: -# read the projection in abctl's plugin pane, then change on_error to enforce +abctl tools scan --write ./cortex-ca/demo.yaml ``` -The config is hot-reloaded, so neither step needs a restart. Note that a reload -rebuilds the plugin and therefore **resets its counters** — the same as a -process restart. +The config is hot-reloaded, so no restart. A reload does rebuild the plugin and +therefore **resets its counters** — the same as a process restart. + +### Measure instead of enforce, when you want to + +`on_error: observe` turns the plugin into a projection: it computes exactly what +it would remove and counts it, while every byte on the wire stays untouched. +Nothing in the plugin differs between the modes — under observe `SetBody` is a +no-op on bytes and leaves `BodyMutated()` false, which is how it knows which +counter to increment. + +Two occasions worth it: + +- **Sizing the change** before it affects anything: read `bytes removed` and + `tokens saved / request`, decide, then remove the line. +- **Clearing the plugin of suspicion.** If requests start failing and you are + not sure whether this is the cause, set `observe` and watch: the bytes are then + provably unmodified, so a failure that persists is not this plugin. That is + faster than reasoning about it, and it costs no configuration. ## Reading the metrics diff --git a/authbridge/install-demo.sh b/authbridge/install-demo.sh index c76a09b91..7ca0871a9 100755 --- a/authbridge/install-demo.sh +++ b/authbridge/install-demo.sh @@ -203,10 +203,9 @@ info "" # behind the user's back -- print the command instead and let them look first. demo_cfg="${ca_dir}/demo.yaml" if [ -f "${demo_cfg}" ]; then - info " Measure tool-manifest waste (writes the remove: list, no restart needed):" + info " Cut tool-manifest waste (fills the remove: list; hot-reloaded, no restart):" info " ${abctl_cmd} tools scan --write ${demo_cfg}" - info " Then read 'requests projected' in abctl's plugin pane before switching" - info " that entry's on_error to enforce." + info " Then watch the Metrics section of tool-prune's pane in abctl." info "" fi info " Watch traffic: ${abctl_cmd} --endpoint http://localhost:${DEMO_SESSION_PORT}" From ae3f0a3562bb77b70624741087253eae48392977 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 15:14:16 -0400 Subject: [PATCH 13/28] fix(authbridge): Make tool-prune's silence diagnosable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three ways this could do nothing while looking correctly configured, all hit during testing, none of which said anything: 1. A query string defeated the path gate. The gate suffix-matched the raw request target, so /v1/messages?beta=true — a request Claude Code really makes — never matched and every such request passed through untouched. Now the query and any trailing slash are stripped first. context-guru shares the pattern and has the same latent bug; not changed here, but worth a follow-up. 2. A skip did not say what it saw. The invocation recorded "path_not_inference" with no path, so the timeline could not distinguish "the path did not match" from "there was no path". A CONNECT tunnel has no path, and conflating the two sent a real investigation looking for a routing problem when TLS simply was not being decrypted. Tunnels now report no_path_tunnelled, and a genuine mismatch records the offending path. 3. The TLS bridge silently not decrypting. If the client does not trust the bridge CA, every HTTPS request opens an opaque CONNECT tunnel: parsers and tool-prune correctly no-op because there is no plaintext, nothing errors, and the only symptom is silence. The forward proxy now counts tunnels against decrypted requests and warns once — after five tunnels with nothing bridged, so passthrough hosts and startup races do not cry wolf — naming the absolute trust-anchor path to point a client at. Engine.CAFile carries that path for diagnostics; absolute because --demo anchors the CA to its launch directory, so a relative path is only right for someone standing in that directory. Also replaces the bare os error from `abctl tools scan --write` on a missing config. "open ./cortex-ca/demo.yaml: no such file or directory" is complete and useless: it names a relative path that resolves against the wrong directory more often than the right one. It now reports the absolute path it looked at, explains that --demo anchors the config to its launch directory, and gives the command that finds the real one. The bridge-health warning is covered by unit tests over the threshold, the bridged>0 case, once-only firing and concurrent tunnels under -race. It is not verified against a live untrusting client, because the health server's port is hardcoded to :9091 and a second proxy cannot start alongside a running one. Signed-off-by: Hai Huang --- .../forwardproxy/bridgehealth_test.go | 109 ++++++++++++++++++ .../authlib/listener/forwardproxy/server.go | 57 +++++++++ .../authlib/plugins/toolprune/plugin.go | 33 +++++- .../authlib/plugins/toolprune/plugin_test.go | 72 ++++++++++++ authbridge/authlib/tlsbridge/engine.go | 6 + authbridge/cmd/abctl/toolscan/patch.go | 17 +++ authbridge/cmd/abctl/toolscan/patch_test.go | 20 ++++ authbridge/cmd/authbridge-proxy/main.go | 13 +++ 8 files changed, 326 insertions(+), 1 deletion(-) create mode 100644 authbridge/authlib/listener/forwardproxy/bridgehealth_test.go diff --git a/authbridge/authlib/listener/forwardproxy/bridgehealth_test.go b/authbridge/authlib/listener/forwardproxy/bridgehealth_test.go new file mode 100644 index 000000000..9c663d50d --- /dev/null +++ b/authbridge/authlib/listener/forwardproxy/bridgehealth_test.go @@ -0,0 +1,109 @@ +package forwardproxy + +import ( + "strings" + "sync" + "testing" + + "github.com/rossoctl/cortex/authbridge/authlib/tlsbridge" +) + +// TestNoteTunnel_WarnsOnlyWhenNothingIsEverDecrypted covers the failure that +// looks like a plugin bug: the bridge is on, the client does not trust its CA, +// so every HTTPS request opens an opaque tunnel and every body-reading plugin +// correctly does nothing. Nothing errors — the only symptom is silence. +func TestNoteTunnel_WarnsOnlyWhenNothingIsEverDecrypted(t *testing.T) { + tests := []struct { + name string + bridge *tlsbridge.Engine + tunnels int + bridged uint64 + wantWarns int + }{ + { + name: "bridge disabled: never warn, tunnels are the expected behaviour", + bridge: nil, + tunnels: 50, + wantWarns: 0, + }, + { + name: "below threshold: a few tunnels are normal (passthrough hosts, startup races)", + bridge: &tlsbridge.Engine{}, + tunnels: tunnelWarnThreshold - 1, + wantWarns: 0, + }, + { + name: "tunnels but something was decrypted: bridge is working", + bridge: &tlsbridge.Engine{}, + tunnels: 50, + bridged: 1, + wantWarns: 0, + }, + { + name: "many tunnels, nothing decrypted: warn", + bridge: &tlsbridge.Engine{}, + tunnels: tunnelWarnThreshold, + wantWarns: 1, + }, + { + name: "and only once, however much traffic follows", + bridge: &tlsbridge.Engine{}, + tunnels: 200, + wantWarns: 1, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + s := &Server{TLSBridge: tc.bridge} + s.bridgedRequests.Store(tc.bridged) + var warns int + // bridgeWarnOnce is the mechanism under test; count how many times + // the guarded block would run by observing the sync.Once directly. + for i := 0; i < tc.tunnels; i++ { + before := s.warnFired() + s.noteTunnel() + if !before && s.warnFired() { + warns++ + } + } + if warns != tc.wantWarns { + t.Errorf("warned %d times, want %d", warns, tc.wantWarns) + } + }) + } +} + +// TestCaFileHint_NamesTheAbsolutePath: a relative path in the fix hint is only +// correct for someone standing in the directory --demo was launched from, which +// is precisely how the trust anchor gets mismatched in the first place. +func TestCaFileHint_NamesTheAbsolutePath(t *testing.T) { + s := &Server{TLSBridge: &tlsbridge.Engine{CAFile: "/abs/cortex-ca/ca.crt"}} + if got := s.caFileHint(); got != "/abs/cortex-ca/ca.crt" { + t.Errorf("caFileHint() = %q", got) + } + // Degrade to a placeholder rather than an empty string, so the log line + // still reads as an instruction. + bare := &Server{TLSBridge: &tlsbridge.Engine{}} + if got := bare.caFileHint(); !strings.Contains(got, "ca.crt") { + t.Errorf("caFileHint() = %q, want something naming ca.crt", got) + } +} + +// TestNoteTunnel_ConcurrentIsRaceFree: tunnels open on many goroutines. +func TestNoteTunnel_ConcurrentIsRaceFree(t *testing.T) { + s := &Server{TLSBridge: &tlsbridge.Engine{}} + var wg sync.WaitGroup + for i := 0; i < 16; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 64; j++ { + s.noteTunnel() + } + }() + } + wg.Wait() + if got := s.tunnelsOpened.Load(); got != 16*64 { + t.Errorf("tunnelsOpened = %d, want %d", got, 16*64) + } +} diff --git a/authbridge/authlib/listener/forwardproxy/server.go b/authbridge/authlib/listener/forwardproxy/server.go index 88a8ba1b7..40e1dde6f 100644 --- a/authbridge/authlib/listener/forwardproxy/server.go +++ b/authbridge/authlib/listener/forwardproxy/server.go @@ -17,6 +17,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" "github.com/rossoctl/cortex/authbridge/authlib/listener/httpx" @@ -68,6 +69,17 @@ type Server struct { SkipHosts *skiphost.Matcher TLSBridge *tlsbridge.Engine // nil = disabled; set by caller after NewServer + + // Bridge-health counters. When the TLS bridge is enabled but the client + // does not trust its CA, every HTTPS request opens a CONNECT tunnel and + // nothing is ever decrypted: the pipeline sees opaque tunnels, every + // body-reading plugin no-ops, and the proxy looks configured but inert. + // Nothing errors, so the only symptom is silence. These count the two + // outcomes so the listener can say so out loud. + tunnelsOpened atomic.Uint64 + bridgedRequests atomic.Uint64 + bridgeWarnOnce sync.Once + bridgeWarned atomic.Bool } // MTLSOptions configures outbound mTLS for the forward proxy. When @@ -211,6 +223,9 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) { // they are origin-form (the caller sets r.URL.Scheme/Host) and must re-originate // via the dedicated upstream client, never the mesh-mTLS s.Client. func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, isBridge bool) { + if isBridge { + s.bridgedRequests.Add(1) + } pctx := &pipeline.Context{ Direction: pipeline.Outbound, Method: r.Method, @@ -916,6 +931,7 @@ const connectDialTimeout = 30 * time.Second // trust path. CONNECT targets are opaque externals (LiteMaaS, Bedrock, // GitHub API, etc.) where the agent's existing TLS is the right answer. func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { + s.noteTunnel() pctx := &pipeline.Context{ Direction: pipeline.Outbound, Method: r.Method, // always "CONNECT" here, but populated for parity with handleRequest @@ -1225,3 +1241,44 @@ func portOf(authority string) int { } return 443 } + +// tunnelWarnThreshold is how many tunnels may open with nothing decrypted +// before the listener speaks up. A handful is normal — passthrough hosts, a +// non-HTTPS CONNECT, the first request racing startup — so warning on the +// first one would cry wolf. By this many, with zero bridged requests, the +// client is not trusting the CA. +const tunnelWarnThreshold = 5 + +// noteTunnel counts a CONNECT tunnel and, once, warns if the bridge is enabled +// yet has never decrypted anything. +// +// This is the failure that looks like a bug in whatever plugin you are testing: +// tool-prune, the parsers and every body reader correctly do nothing, because +// there is no plaintext to act on. Naming the trust anchor turns a silent +// dead end into a one-line fix. +func (s *Server) noteTunnel() { + n := s.tunnelsOpened.Add(1) + if s.TLSBridge == nil || n < tunnelWarnThreshold || s.bridgedRequests.Load() > 0 { + return + } + s.bridgeWarnOnce.Do(func() { + s.bridgeWarned.Store(true) + slog.Warn("tls-bridge: enabled but nothing has been decrypted — every request is tunnelling through opaquely, so body-reading plugins (parsers, tool-prune) cannot act", + "tunnels_opened", n, + "bridged_requests", 0, + "likely_cause", "the client does not trust the bridge CA", + "fix", "point the client at the trust anchor, e.g. NODE_EXTRA_CA_CERTS="+s.caFileHint()) + }) +} + +func (s *Server) caFileHint() string { + if s.TLSBridge != nil && s.TLSBridge.CAFile != "" { + return s.TLSBridge.CAFile + } + return "/ca.crt" +} + +// warnFired reports whether the bridge-health warning has already been emitted. +// Exists for tests: sync.Once has no public "has it run" query, and asserting +// on log output would couple the test to the message text. +func (s *Server) warnFired() bool { return s.bridgeWarned.Load() } diff --git a/authbridge/authlib/plugins/toolprune/plugin.go b/authbridge/authlib/plugins/toolprune/plugin.go index 40ab9d6d7..888c55025 100644 --- a/authbridge/authlib/plugins/toolprune/plugin.go +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -126,7 +126,13 @@ func (p *ToolPrune) Configure(raw json.RawMessage) error { } // gated reports whether the request path is one the plugin acts on. +// +// The query string is stripped first. Providers accept query parameters on +// these endpoints — /v1/messages?beta=true is a real request Claude Code makes — +// and a suffix match against the raw target silently misses every one of them, +// which reads as the plugin doing nothing for no visible reason. func (p *ToolPrune) gated(path string) bool { + path = pathOnly(path) for _, s := range p.cfg.Paths { if path == s || strings.HasSuffix(path, s) { return true @@ -135,6 +141,18 @@ func (p *ToolPrune) gated(path string) bool { return false } +// pathOnly drops a query string and any trailing slash, so the configured +// suffixes match the endpoint rather than the exact request target. +func pathOnly(target string) string { + if i := strings.IndexAny(target, "?#"); i >= 0 { + target = target[:i] + } + if len(target) > 1 && strings.HasSuffix(target, "/") { + target = strings.TrimRight(target, "/") + } + return target +} + // toolNameAt extracts a tool's name from raw manifest element i, covering both // dialects: Anthropic puts it at tools.i.name, OpenAI at tools.i.function.name. func toolNameAt(body []byte, i int) string { @@ -178,7 +196,20 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action }() if !p.gated(pctx.Path) { - pctx.Record(pipeline.Invocation{Action: pipeline.ActionSkip, Reason: "path_not_inference"}) + // Distinguish "this is not an HTTP request at all" from "the path did + // not match". A CONNECT tunnel has no path, and reporting it as a path + // mismatch sends an operator hunting for a routing problem when the + // real answer is that TLS is not being decrypted — so the client does + // not trust the bridge CA and nothing downstream can see the request. + reason := "path_not_inference" + if pctx.Path == "" { + reason = "no_path_tunnelled" + } + pctx.Record(pipeline.Invocation{ + Action: pipeline.ActionSkip, + Reason: reason, + Path: pctx.Path, + }) return action } // inference-parser establishes that this is an inference call at all. Its diff --git a/authbridge/authlib/plugins/toolprune/plugin_test.go b/authbridge/authlib/plugins/toolprune/plugin_test.go index 1ebf49d57..097693518 100644 --- a/authbridge/authlib/plugins/toolprune/plugin_test.go +++ b/authbridge/authlib/plugins/toolprune/plugin_test.go @@ -471,3 +471,75 @@ func TestPrune_ToolChoiceAutoDoesNotBlockPruning(t *testing.T) { }) } } + +// TestPrune_PathGateIgnoresQueryString: providers accept query parameters on +// these endpoints, and Claude Code really does send /v1/messages?beta=true. A +// suffix match against the raw target misses every such request and the plugin +// silently does nothing — the least debuggable possible failure, because +// everything looks configured correctly. +func TestPrune_PathGateIgnoresQueryString(t *testing.T) { + body := `{"tools":[{"name":"Read"},{"name":"NotebookEdit"}]}` + for _, path := range []string{ + "/v1/messages", + "/v1/messages?beta=true", + "/v1/messages?beta=true&x=1", + "/v1/messages/", + "/v1/chat/completions?stream=false", + "https://host/v1/messages?beta=true", // absolute-form target via a proxy + } { + t.Run(path, func(t *testing.T) { + p := configured(t, "NotebookEdit") + pctx := inferenceCtx(path, body, "Read", "NotebookEdit") + run(t, p, pctx) + if !pctx.BodyMutated() { + t.Errorf("path %q was not treated as an inference endpoint", path) + } + }) + } +} + +// TestPrune_NonInferencePathsStillSkip guards the other direction: loosening the +// gate must not make it match everything. +func TestPrune_NonInferencePathsStillSkip(t *testing.T) { + body := `{"tools":[{"name":"NotebookEdit"}]}` + for _, path := range []string{"/mcp", "/v1/models", "/healthz", "/v1/messages/batches"} { + t.Run(path, func(t *testing.T) { + p := configured(t, "NotebookEdit") + pctx := inferenceCtx(path, body, "NotebookEdit") + run(t, p, pctx) + if pctx.BodyMutated() { + t.Errorf("path %q must not be pruned", path) + } + }) + } +} + +// TestPrune_TunnelSkipIsDistinguishable: a CONNECT tunnel has no path. Reporting +// that as a path mismatch sent a real investigation hunting for a routing +// problem when the actual cause was that TLS was never decrypted. The reason +// code has to say which. +func TestPrune_TunnelSkipIsDistinguishable(t *testing.T) { + p := configured(t, "NotebookEdit") + pctx := inferenceCtx("", `{"tools":[{"name":"NotebookEdit"}]}`, "NotebookEdit") + run(t, p, pctx) + + if pctx.Extensions.Invocations == nil || len(pctx.Extensions.Invocations.Inbound) == 0 { + t.Fatal("expected a skip invocation") + } + inv := pctx.Extensions.Invocations.Inbound[0] + if inv.Reason != "no_path_tunnelled" { + t.Errorf("reason = %q, want no_path_tunnelled so a tunnel is not mistaken for a routing problem", inv.Reason) + } +} + +// TestPrune_PathMismatchRecordsThePath: a skip that does not say what it saw +// cannot be diagnosed from the session timeline. +func TestPrune_PathMismatchRecordsThePath(t *testing.T) { + p := configured(t, "NotebookEdit") + pctx := inferenceCtx("/v1/models", `{"tools":[{"name":"NotebookEdit"}]}`, "NotebookEdit") + run(t, p, pctx) + inv := pctx.Extensions.Invocations.Inbound[0] + if inv.Reason != "path_not_inference" || inv.Path != "/v1/models" { + t.Errorf("inv = %+v, want path_not_inference with the offending path recorded", inv) + } +} diff --git a/authbridge/authlib/tlsbridge/engine.go b/authbridge/authlib/tlsbridge/engine.go index d90ebaa3b..e9f87ffa8 100644 --- a/authbridge/authlib/tlsbridge/engine.go +++ b/authbridge/authlib/tlsbridge/engine.go @@ -12,4 +12,10 @@ type Engine struct { Skip *SkipSet Upstream *http.Client CAPEM []byte + + // CAFile is the on-disk trust anchor clients must load. Diagnostics only: + // the bridge itself works from CAPEM. It exists so a listener that notices + // nothing is being decrypted can name the exact file to trust, which is + // the single most common cause of that state. + CAFile string } diff --git a/authbridge/cmd/abctl/toolscan/patch.go b/authbridge/cmd/abctl/toolscan/patch.go index 154d903ef..67f9ea380 100644 --- a/authbridge/cmd/abctl/toolscan/patch.go +++ b/authbridge/cmd/abctl/toolscan/patch.go @@ -1,8 +1,10 @@ package toolscan import ( + "errors" "fmt" "os" + "path/filepath" "regexp" "strings" ) @@ -24,6 +26,21 @@ var ( func PatchConfig(path string, candidates []string) (changed bool, err error) { orig, err := os.ReadFile(path) //nolint:gosec // operator-supplied config path if err != nil { + if errors.Is(err, os.ErrNotExist) { + // The bare os error ("open ./cortex-ca/demo.yaml: no such file or + // directory") is technically complete and practically useless: the + // demo anchors its config to the directory it was launched from, so + // a relative path resolves against the wrong place more often than + // the right one. Say where we looked and what to do about it. + abs, aerr := filepath.Abs(path) + if aerr != nil { + abs = path + } + return false, fmt.Errorf("no config at %s\n"+ + " authbridge-proxy --demo writes cortex-ca/demo.yaml into the directory it is started from,\n"+ + " so run this from there or pass an absolute path. To find it:\n"+ + " curl -s localhost:47602/config | grep ca_dir", abs) + } return false, err } lines := strings.Split(string(orig), "\n") diff --git a/authbridge/cmd/abctl/toolscan/patch_test.go b/authbridge/cmd/abctl/toolscan/patch_test.go index e04f3dd3e..fa2f5ce3c 100644 --- a/authbridge/cmd/abctl/toolscan/patch_test.go +++ b/authbridge/cmd/abctl/toolscan/patch_test.go @@ -173,3 +173,23 @@ func TestPatchConfig_DoesNotEscapeTheEntry(t *testing.T) { t.Errorf("tool-prune's list was not patched:\n%s", got) } } + +// TestPatchConfig_MissingFileExplainsWhere: the bare os error names a relative +// path and nothing else, which twice sent a real user hunting in the wrong +// directory — the demo anchors its config to wherever it was launched, so a +// relative path usually resolves somewhere unintended. +func TestPatchConfig_MissingFileExplainsWhere(t *testing.T) { + _, err := PatchConfig("./definitely-not-here/demo.yaml", []string{"NotebookEdit"}) + if err == nil { + t.Fatal("expected an error") + } + msg := err.Error() + for _, want := range []string{"no config at", "/definitely-not-here/demo.yaml", "--demo", "absolute path", "ca_dir"} { + if !strings.Contains(msg, want) { + t.Errorf("error should mention %q:\n%s", want, msg) + } + } + if strings.Contains(msg, "no such file or directory") { + t.Errorf("should replace the bare os error, not wrap it:\n%s", msg) + } +} diff --git a/authbridge/cmd/authbridge-proxy/main.go b/authbridge/cmd/authbridge-proxy/main.go index a2e85b819..3788602da 100644 --- a/authbridge/cmd/authbridge-proxy/main.go +++ b/authbridge/cmd/authbridge-proxy/main.go @@ -378,6 +378,7 @@ func main() { Skip: tlsbridge.NewSkipSet(), Upstream: up, CAPEM: src.CACertPEM(), + CAFile: caTrustPath(cfg.TLSBridge.CADir), } slog.Info("tls-bridge enabled", "ca_dir", cfg.TLSBridge.CADir) } @@ -557,3 +558,15 @@ func startTransparentProxy(fp *forwardproxy.Server, addr string) *net.TCPListene }() return ln } + +// caTrustPath returns the absolute path of the CA clients must trust. Absolute +// because --demo anchors the CA to its launch directory, so a relative path in +// a log line is only correct for someone standing in that same directory — +// which is exactly how the trust anchor gets mismatched. +func caTrustPath(caDir string) string { + p := filepath.Join(caDir, "ca.crt") + if abs, err := filepath.Abs(p); err == nil { + return abs + } + return p +} From f33d066ac753e3c20616dbb9a1bacc8550d259c0 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 15:22:36 -0400 Subject: [PATCH 14/28] fix(abctl): Refresh plugin metrics instead of freezing them at connect MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The plugin detail pane rendered "Metrics: (none)" no matter how much traffic a plugin had processed. /v1/pipeline reported the counters correctly; abctl fetched that view exactly once at startup, on the documented assumption that "the pipeline is static for the duration of a process so there's no periodic refresh." That assumption was true of the composition and false of the counters I attached to the same view. On a freshly started proxy the fetch happens before any traffic, every counter is zero, snapshot() returns nil, and omitempty drops the key — so the pane showed (none) permanently, which reads as "this plugin does nothing" rather than "this number is stale". The view is now refetched when it can be seen: immediately on opening the plugin detail pane, because that is exactly when someone wants current numbers, and on the existing 2s refresh tick while the detail or pipeline pane is open. Elsewhere it is left alone — the composition genuinely does not change, so polling it while nobody is looking at metrics would be overhead for nothing. A refreshed view is also re-rendered into an already-open detail pane, resolved by name and direction. Without that the pane keeps displaying the snapshot it was opened with, which was the actual bug: the fetch was happening on the tick, and the pane was ignoring the result. Signed-off-by: Hai Huang --- authbridge/cmd/abctl/tui/app.go | 23 +++++++++- authbridge/cmd/abctl/tui/keys.go | 4 +- .../cmd/abctl/tui/plugin_detail_pane.go | 18 ++++++++ .../cmd/abctl/tui/plugin_metrics_test.go | 46 +++++++++++++++++++ 4 files changed, 88 insertions(+), 3 deletions(-) diff --git a/authbridge/cmd/abctl/tui/app.go b/authbridge/cmd/abctl/tui/app.go index edfe1fa14..1aa91abed 100644 --- a/authbridge/cmd/abctl/tui/app.go +++ b/authbridge/cmd/abctl/tui/app.go @@ -440,8 +440,11 @@ func (m *model) Init() tea.Cmd { return m.initSessionView() } -// loadPipelineCmd fetches /v1/pipeline once at startup. The pipeline is -// static for the duration of a process so there's no periodic refresh. +// loadPipelineCmd fetches /v1/pipeline. The plugin composition is static for +// the life of a process, but the view also carries each plugin's live +// Metrics counters — so a single fetch at startup would freeze them at zero, +// which on a fresh proxy is every number a user ever sees. It is refetched +// when a metrics-bearing pane is open; see refreshTickMsg. func (m *model) loadPipelineCmd() tea.Cmd { return func() tea.Msg { pv, err := m.client.GetPipeline(m.ctx) @@ -583,11 +586,27 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { if m.pane == paneNamespaces || m.pane == panePods { return m, refreshTickCmd() } + // Refresh the pipeline view too while a pane that displays plugin + // Metrics is open, so counters tick rather than sitting at whatever + // they were when the session was first opened. Skipped elsewhere: + // the composition itself does not change, so polling it while nobody + // is looking at metrics would be pure overhead. + if m.pane == panePluginDetail || m.pane == panePipeline { + return m, tea.Batch(m.loadSessionsCmd(), m.loadPipelineCmd(), refreshTickCmd()) + } return m, tea.Batch(m.loadSessionsCmd(), refreshTickCmd()) case pipelineLoadedMsg: m.pipeline = (*apiclient.PipelineView)(msg) m.rebuildPipelineTable() + // Re-render an open plugin detail pane against the new view. Without + // this the pane keeps showing the snapshot it was opened with, so + // Metrics would still read (none) however long traffic ran. + if m.pane == panePluginDetail && m.detailPlugin != nil { + if p := m.livePipelinePlugin(m.detailPlugin); p != nil { + m.showPluginDetail(p) + } + } return m, nil case catalogLoadedMsg: diff --git a/authbridge/cmd/abctl/tui/keys.go b/authbridge/cmd/abctl/tui/keys.go index 72a7a0fae..f903a8c4a 100644 --- a/authbridge/cmd/abctl/tui/keys.go +++ b/authbridge/cmd/abctl/tui/keys.go @@ -274,7 +274,9 @@ func (m *model) handleKey(msg tea.KeyMsg) tea.Cmd { m.previousPane = panePipeline m.showPluginDetail(p) m.pane = panePluginDetail - return nil + // Fetch immediately rather than waiting for the next refresh tick: + // opening the pane is exactly when someone wants current counters. + return m.loadPipelineCmd() case paneCatalog: p := m.selectedCatalogEntry() if p == nil { diff --git a/authbridge/cmd/abctl/tui/plugin_detail_pane.go b/authbridge/cmd/abctl/tui/plugin_detail_pane.go index efd12ac2b..c775337ae 100644 --- a/authbridge/cmd/abctl/tui/plugin_detail_pane.go +++ b/authbridge/cmd/abctl/tui/plugin_detail_pane.go @@ -93,3 +93,21 @@ func (m *model) showPluginDetail(p *apiclient.PipelinePlugin) { m.detailVp.SetContent(b.String()) m.detailVp.GotoTop() } + +// livePipelinePlugin re-resolves a plugin against the current pipeline view by +// name and direction, so a refreshed view can be rendered into an already-open +// detail pane. Returns nil for a catalog entry (blank direction), which has no +// counterpart in the active chain. +func (m *model) livePipelinePlugin(want *apiclient.PipelinePlugin) *apiclient.PipelinePlugin { + if want == nil || m.pipeline == nil || want.Direction == "" { + return nil + } + for _, set := range [][]apiclient.PipelinePlugin{m.pipeline.Inbound, m.pipeline.Outbound} { + for i := range set { + if set[i].Name == want.Name && set[i].Direction == want.Direction { + return &set[i] + } + } + } + return nil +} diff --git a/authbridge/cmd/abctl/tui/plugin_metrics_test.go b/authbridge/cmd/abctl/tui/plugin_metrics_test.go index fef61710a..9e18a323c 100644 --- a/authbridge/cmd/abctl/tui/plugin_metrics_test.go +++ b/authbridge/cmd/abctl/tui/plugin_metrics_test.go @@ -75,3 +75,49 @@ func TestFormatPluginMetrics_EmptyIsEmptyString(t *testing.T) { t.Errorf("formatPluginMetrics(nil) = %q, want empty (pane renders (none))", got) } } + +// TestLivePipelinePlugin_ResolvesAgainstRefreshedView: plugin Metrics are live +// counters riding on a view that was originally fetched once at startup, on the +// documented assumption that the pipeline composition never changes. That +// assumption held for composition and broke for counters — an open detail pane +// kept rendering its opening snapshot, so Metrics read "(none)" forever on a +// proxy that had counted nothing yet at connect time. +func TestLivePipelinePlugin_ResolvesAgainstRefreshedView(t *testing.T) { + shown := &apiclient.PipelinePlugin{Name: "tool-prune", Direction: "outbound"} + + m := &model{pipeline: &apiclient.PipelineView{ + Outbound: []apiclient.PipelinePlugin{ + {Name: "inference-parser", Direction: "outbound"}, + {Name: "tool-prune", Direction: "outbound", Metrics: []apiclient.PluginMetric{ + {Name: "requests pruned", Value: 2, Unit: "count"}, + }}, + }, + }} + + got := m.livePipelinePlugin(shown) + if got == nil { + t.Fatal("tool-prune not resolved against the refreshed view") + } + if len(got.Metrics) != 1 || got.Metrics[0].Value != 2 { + t.Errorf("resolved plugin carries no fresh metrics: %+v", got.Metrics) + } +} + +// TestLivePipelinePlugin_CatalogEntryHasNoLiveCounterpart: a catalog entry is +// synthesised with a blank direction and is not in the active chain, so there is +// nothing to refresh it from. +func TestLivePipelinePlugin_CatalogEntryHasNoLiveCounterpart(t *testing.T) { + m := &model{pipeline: &apiclient.PipelineView{ + Outbound: []apiclient.PipelinePlugin{{Name: "tool-prune", Direction: "outbound"}}, + }} + if got := m.livePipelinePlugin(&apiclient.PipelinePlugin{Name: "tool-prune"}); got != nil { + t.Errorf("catalog entry (blank direction) should not resolve, got %+v", got) + } + if got := m.livePipelinePlugin(nil); got != nil { + t.Error("nil input should return nil") + } + // No view fetched yet. + if got := (&model{}).livePipelinePlugin(&apiclient.PipelinePlugin{Name: "x", Direction: "outbound"}); got != nil { + t.Error("nil pipeline should return nil") + } +} From 1c9bdfae5a26f4ae209cae39dc009b903ff4464c Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 15:38:13 -0400 Subject: [PATCH 15/28] feat(authbridge): Cost tool-prune's saving per prompt-cache tier MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit "tokens saved / request" was a single blended number, and it was the wrong shape: prompt tokens are not fungible. Anthropic charges 1.25x the input rate for a cache write and 0.1x for a cache read, so identical pruned bytes are worth more than 12x more on a cache miss than on a hit. Any dollar figure derived from one blended count is wrong by up to that factor, and the observed traffic alternates miss/hit turn by turn. litellm-budget-track already learned this — "Flat pricing would overstate cache-heavy traffic (Claude Code) by up to ~10x" — so this follows its convention rather than inventing a second one: - Three operator-configured rates with the same names and semantics: input_cost_per_token, cache_write_cost_per_token, cache_read_cost_per_token, the cache rates defaulting to the input rate. No output rate: pruning only shrinks the prompt, so charging output to it would be false. - No price is ever assumed. With no rates set the row reads "set input_cost_per_token to cost this" rather than showing a number. A confidently wrong dollar figure is worse than none, because it gets quoted in decisions. The saving is attributed to the tier it actually came out of. The tool manifest sits inside the cached prefix — Claude Code puts cache_control on the tool block — so a cache-miss request saves cache-write tokens and a hit saves cache-read tokens. Reported as separate rows, never summed, so the readout cannot be multiplied by a single rate. The assumption about manifest placement is stated at the attribution site, since it is the one thing a differently-shaped client would invalidate. Per-request byte savings reach OnFinish through pipeline.SetState, the documented cross-phase state API, so the saving is paired with the usage split of the same request. The bytes-to-tokens ratio stays calibrated on observed traffic, now over the summed per-tier prompt counts rather than the legacy aggregate. Tests cover attribution for miss / hit / no-cache, the absence of a blended row, the unpriced case naming the field that enables costing, and that the write-vs-read cost ratio really is ~12.5x at published ratios — which is the whole reason the tiers are separate. Signed-off-by: Hai Huang --- .../authlib/plugins/toolprune/metrics.go | 99 +++++++++---- .../authlib/plugins/toolprune/plugin.go | 107 +++++++++++-- .../authlib/plugins/toolprune/plugin_test.go | 140 ++++++++++++++---- authbridge/docs/plugin-catalog.md | 1 + authbridge/docs/tool-prune-plugin.md | 67 +++++++-- 5 files changed, 334 insertions(+), 80 deletions(-) diff --git a/authbridge/authlib/plugins/toolprune/metrics.go b/authbridge/authlib/plugins/toolprune/metrics.go index 62bd2f5d0..3309e83d0 100644 --- a/authbridge/authlib/plugins/toolprune/metrics.go +++ b/authbridge/authlib/plugins/toolprune/metrics.go @@ -24,10 +24,14 @@ type metrics struct { bytesRemoved uint64 - // Calibration sample for bytes -> tokens, gathered from response usage. - promptTokens uint64 - requestBytes uint64 - requestsWithUsage uint64 + // Estimated tokens saved, split by the prompt tier the saving came out of. + // Kept apart because providers price the tiers very differently: a blended + // total cannot be multiplied by any single rate without being wrong by up + // to ~12x on cache-heavy traffic. + savedInput float64 + savedCacheWrite float64 + savedCacheRead float64 + requestsCosted uint64 } func (m *metrics) seen() { @@ -64,18 +68,24 @@ func (m *metrics) record(names []string, bytesRemoved int) { } } -func (m *metrics) observeUsage(promptTokens, requestBytes int) { +func (m *metrics) observeSaving(tokens float64, t tier) { m.mu.Lock() - m.promptTokens += uint64(promptTokens) - m.requestBytes += uint64(requestBytes) - m.requestsWithUsage++ + switch t { + case tierCacheWrite: + m.savedCacheWrite += tokens + case tierCacheRead: + m.savedCacheRead += tokens + default: + m.savedInput += tokens + } + m.requestsCosted++ m.mu.Unlock() } // snapshot renders the counters as operator-facing metrics. Every derived row // carries the sample it was computed from, so a figure can never be read as // more certain than it is. -func (m *metrics) snapshot() []pipeline.Metric { +func (m *metrics) snapshot(cfg *config) []pipeline.Metric { m.mu.Lock() defer m.mu.Unlock() @@ -109,29 +119,64 @@ func (m *metrics) snapshot() []pipeline.Metric { acted := m.requestsPruned + m.requestsProjected if acted > 0 { - perReq := float64(m.bytesRemoved) / float64(acted) out = append(out, pipeline.Metric{ - Name: "bytes removed / request", Value: perReq, Unit: "bytes", + Name: "bytes removed / request", Value: float64(m.bytesRemoved) / float64(acted), Unit: "bytes", }) - // Calibrate bytes -> tokens on the operator's own traffic instead of - // bundling a tokenizer or assuming a constant. With no usage sample - // yet, report zero rather than dividing by zero. - if m.requestBytes > 0 && m.promptTokens > 0 { - ratio := float64(m.promptTokens) / float64(m.requestBytes) - out = append(out, pipeline.Metric{ - Name: "tokens saved / request", - Value: perReq * ratio, - Unit: "tokens", - Note: fmt.Sprintf("estimate, n=%d", m.requestsWithUsage), - }) - } else { + } + + // Tokens saved, per prompt tier. Deliberately not summed: the tiers are + // priced differently enough (Anthropic: cache write 1.25x input, cache read + // 0.1x) that one total invites a multiplication that is wrong by >12x. + note := "" + if m.requestsCosted > 0 { + note = fmt.Sprintf("estimate, n=%d", m.requestsCosted) + } + tiers := []struct { + name string + val float64 + rate float64 + }{ + {"tokens saved: cache write", m.savedCacheWrite, cfg.cacheWriteRate()}, + {"tokens saved: cache read", m.savedCacheRead, cfg.cacheReadRate()}, + {"tokens saved: input", m.savedInput, cfg.InputCostPerToken}, + } + var usd float64 + for _, t := range tiers { + if t.val <= 0 { + continue + } + out = append(out, pipeline.Metric{ + Name: t.name, Value: t.val, Unit: "tokens", Note: note, + }) + usd += t.val * t.rate + } + if m.requestsCosted == 0 && acted > 0 { + out = append(out, pipeline.Metric{ + Name: "tokens saved", Value: 0, Unit: "tokens", + Note: "no response usage seen yet", + }) + } + + // Cost only when rates are configured. Without them no price is assumed: + // a confidently wrong dollar figure is worse than none, because it gets + // quoted in decisions. + if cfg.priced() && usd > 0 { + out = append(out, pipeline.Metric{ + Name: "$ saved", Value: usd, Unit: "usd", Note: note, + }) + if m.requestsCosted > 0 { out = append(out, pipeline.Metric{ - Name: "tokens saved / request", - Value: 0, - Unit: "tokens", - Note: "no usage sample yet", + Name: "$ saved / request", + Value: usd / float64(m.requestsCosted), + Unit: "usd", + Note: note, }) } + } else if acted > 0 && !cfg.priced() { + out = append(out, pipeline.Metric{ + Name: "$ saved", Value: 0, Unit: "usd", + Note: "set input_cost_per_token to cost this", + }) } // Per-tool attribution, sorted by count then name so the readout is diff --git a/authbridge/authlib/plugins/toolprune/plugin.go b/authbridge/authlib/plugins/toolprune/plugin.go index 888c55025..f26aaf83f 100644 --- a/authbridge/authlib/plugins/toolprune/plugin.go +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -53,6 +53,39 @@ type config struct { // Paths are the request paths this plugin acts on, matched exactly or by // suffix. Defaults to the three inference endpoints. Paths []string `json:"paths" description:"Request paths to act on (exact or suffix match)."` + + // Per-token rates, for costing the saving. Field names and semantics match + // litellm-budget-track so an operator configures rates in one familiar + // shape. USD per token; all optional. With none set, no cost is reported — + // a price is never assumed. + // + // There is deliberately no output rate: pruning only ever shrinks the + // prompt, so attributing any output cost to it would be false. + InputCostPerToken float64 `json:"input_cost_per_token" description:"USD per uncached input token."` + CacheWriteCostPerToken float64 `json:"cache_write_cost_per_token" description:"USD per cache-write input token; defaults to input_cost_per_token."` + CacheReadCostPerToken float64 `json:"cache_read_cost_per_token" description:"USD per cache-read input token; defaults to input_cost_per_token."` +} + +// cacheWriteRate / cacheReadRate default to the uncached input rate, matching +// litellm-budget-track. Flat pricing would misstate cache-heavy traffic badly — +// Anthropic charges 1.25x input for a cache write and 0.1x for a read, so the +// same saved bytes differ by more than 12x depending on which tier they land in. +func (c *config) cacheWriteRate() float64 { + if c.CacheWriteCostPerToken > 0 { + return c.CacheWriteCostPerToken + } + return c.InputCostPerToken +} + +func (c *config) cacheReadRate() float64 { + if c.CacheReadCostPerToken > 0 { + return c.CacheReadCostPerToken + } + return c.InputCostPerToken +} + +func (c *config) priced() bool { + return c.InputCostPerToken > 0 || c.CacheWriteCostPerToken > 0 || c.CacheReadCostPerToken > 0 } func (c *config) applyDefaults() { @@ -312,6 +345,10 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action } removedBytes := len(body) - len(out) + // Carry the saving to OnFinish, where the response reveals which token tier + // it came out of. SetState keeps it private to this plugin, unlike + // Extensions.Custom which is shared. + pipeline.SetState(pctx, p.Name(), &requestState{bytesRemoved: removedBytes}) pctx.SetBody(out) // Under ErrorPolicyObserve, SetBody is a no-op on bytes and leaves // bodyMutated false — so this same code path measures without enforcing, @@ -328,19 +365,71 @@ func (p *ToolPrune) OnResponse(_ context.Context, _ *pipeline.Context) pipeline. return pipeline.Action{Type: pipeline.Continue} } -// OnFinish calibrates the bytes-to-tokens ratio on the operator's own traffic, -// rather than bundling a tokenizer or hardcoding a constant. inference-parser -// is a StreamingResponder, so RunResponse skips its OnResponse — OnFinish is -// the hook where response-derived usage is reliably available. +// requestState carries the per-request byte saving from OnRequest to OnFinish. +type requestState struct{ bytesRemoved int } + +// OnFinish converts the request's byte saving into tokens and attributes it to +// the token tier it actually came out of. +// +// Two things make a single "tokens saved" number wrong, which is why this is +// per-tier. First, the ratio: rather than bundling a tokenizer or assuming +// bytes-per-token, it is calibrated on this request — prompt tokens over request +// bytes, both post-pruning, so the two sides are consistent. Second, and larger: +// providers price prompt tiers very differently. Anthropic charges 1.25x the +// input rate for a cache write and 0.1x for a cache read, so identical saved +// bytes are worth more than 12x more on a cache miss than on a hit. Reporting +// one blended figure would hide a factor of twelve. +// +// The tool manifest sits inside the cached prefix — Claude Code puts +// cache_control on the tool block — so on a cache-miss request the saving comes +// out of cache writes, and on a hit out of cache reads. That is the assumption +// this attribution rests on; it is stated here because it is the one thing that +// would need revisiting for a client that lays out its prompt differently. func (p *ToolPrune) OnFinish(_ context.Context, pctx *pipeline.Context) { - if pctx.Extensions.Inference == nil { + st := pipeline.GetState[requestState](pctx, p.Name()) + if st == nil || st.bytesRemoved <= 0 { + return + } + inf := pctx.Extensions.Inference + if inf == nil || len(pctx.Body) == 0 { + return + } + promptTotal := inf.InputTokens + inf.CacheReadTokens + inf.CacheWriteTokens + if promptTotal <= 0 { + // Fall back to the aggregate when a provider reports only a total. + promptTotal = inf.PromptTokens + } + if promptTotal <= 0 { return } - prompt := pctx.Extensions.Inference.PromptTokens - if prompt <= 0 || len(pctx.Body) == 0 { + tokens := float64(st.bytesRemoved) * float64(promptTotal) / float64(len(pctx.Body)) + if tokens <= 0 { return } - p.m.observeUsage(prompt, len(pctx.Body)) + p.m.observeSaving(tokens, tierOf(inf)) +} + +// tier names which prompt token tier a request's saving came out of. +type tier int + +const ( + tierInput tier = iota + tierCacheWrite + tierCacheRead +) + +// tierOf picks the tier the pruned manifest belonged to. The manifest is in the +// cached prefix, so a write-dominant request wrote it and a read-dominant one +// read it; with no cache tokens reported at all it was plain input. +func tierOf(inf *pipeline.InferenceExtension) tier { + switch { + case inf.CacheWriteTokens > inf.CacheReadTokens && inf.CacheWriteTokens > 0: + return tierCacheWrite + case inf.CacheReadTokens > 0: + return tierCacheRead + default: + return tierInput + } } // noteDrift logs, once, any configured name absent from the first manifest the @@ -370,4 +459,4 @@ func (p *ToolPrune) noteDrift(observed []pipeline.InferenceTool) { } // Metrics implements pipeline.MetricsProvider. -func (p *ToolPrune) Metrics() []pipeline.Metric { return p.m.snapshot() } +func (p *ToolPrune) Metrics() []pipeline.Metric { return p.m.snapshot(&p.cfg) } diff --git a/authbridge/authlib/plugins/toolprune/plugin_test.go b/authbridge/authlib/plugins/toolprune/plugin_test.go index 097693518..1b13f580e 100644 --- a/authbridge/authlib/plugins/toolprune/plugin_test.go +++ b/authbridge/authlib/plugins/toolprune/plugin_test.go @@ -293,45 +293,129 @@ func TestPrune_EnforceCountsPruned(t *testing.T) { } } -// TestMetrics_NoUsageSampleReportsZeroNotNaN: the bytes-to-tokens ratio divides -// by a sample that starts empty. Report a zero-valued estimate with a note -// rather than NaN or a panic. -func TestMetrics_NoUsageSampleReportsZeroNotNaN(t *testing.T) { - p := configured(t, "NotebookEdit") - pctx := inferenceCtx("/v1/messages", anthropicBody, "Read", "NotebookEdit") +// finish drives OnFinish with a given per-tier usage split. +func finish(t *testing.T, p *ToolPrune, pctx *pipeline.Context, input, cacheRead, cacheWrite int) { + t.Helper() + pctx.Extensions.Inference.InputTokens = input + pctx.Extensions.Inference.CacheReadTokens = cacheRead + pctx.Extensions.Inference.CacheWriteTokens = cacheWrite + p.OnFinish(context.Background(), pctx) +} + +func pruneOnce(t *testing.T, p *ToolPrune) *pipeline.Context { + t.Helper() + pctx := inferenceCtx("/v1/messages", anthropicBody, "Read", "NotebookEdit", "Bash") run(t, p, pctx) + if !pctx.BodyMutated() { + t.Fatal("expected a prune") + } + return pctx +} - m := findMetric(t, p.Metrics(), "tokens saved / request") - if m.Value != 0 { - t.Errorf("value = %v, want 0 with no usage sample", m.Value) +// TestMetrics_NoUsageYetReportsZero: before any response usage is seen there is +// no ratio to convert bytes with, so report zero with the reason rather than a +// number or a NaN. +func TestMetrics_NoUsageYetReportsZero(t *testing.T) { + p := configured(t, "NotebookEdit") + pruneOnce(t, p) + m := findMetric(t, p.Metrics(), "tokens saved") + if m.Value != 0 || m.Note != "no response usage seen yet" { + t.Errorf("got %+v, want 0 with the missing-sample reason", m) } - if m.Note != "no usage sample yet" { - t.Errorf("note = %q, want the missing-sample caveat", m.Note) +} + +// TestMetrics_AttributesSavingToTheRightTier is the core of the design. The tool +// manifest lives in the cached prefix, so on a cache-miss request the saving +// comes out of cache writes and on a hit out of cache reads. Reporting one +// blended token count would hide which — and the tiers are priced up to 12x +// apart, so that distinction is the whole number. +func TestMetrics_AttributesSavingToTheRightTier(t *testing.T) { + tests := []struct { + name string + input, cacheRead, cacheWrite int + wantRow string + }{ + {"cache miss writes the prefix", 8881, 0, 24701, "tokens saved: cache write"}, + {"cache hit reads the prefix", 26, 24701, 8907, "tokens saved: cache read"}, + {"no caching at all", 40000, 0, 0, "tokens saved: input"}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + p := configured(t, "NotebookEdit") + pctx := pruneOnce(t, p) + finish(t, p, pctx, tc.input, tc.cacheRead, tc.cacheWrite) + + ms := p.Metrics() + got := findMetric(t, ms, tc.wantRow) + if got.Value <= 0 { + t.Errorf("%s = %v, want positive", tc.wantRow, got.Value) + } + if !strings.HasPrefix(got.Note, "estimate, n=") { + t.Errorf("note = %q, want it labelled an estimate with its sample", got.Note) + } + // No other tier may be credited, and there must be no blended total. + for _, m := range ms { + if m.Name == "tokens saved" { + t.Error("a blended 'tokens saved' row invites multiplying by one rate") + } + if strings.HasPrefix(m.Name, "tokens saved: ") && m.Name != tc.wantRow { + t.Errorf("saving also credited to %q", m.Name) + } + } + }) } } -// TestMetrics_TokenEstimateCalibratesOnObservedUsage: once OnFinish has seen a -// response usage block, the estimate is derived from the operator's own -// traffic and labelled with its sample size. -func TestMetrics_TokenEstimateCalibratesOnObservedUsage(t *testing.T) { +// TestMetrics_NoRatesMeansNoDollarFigure: a price is never assumed. An +// unconfigured plugin says so instead of inventing one. +func TestMetrics_NoRatesMeansNoDollarFigure(t *testing.T) { p := configured(t, "NotebookEdit") - pctx := inferenceCtx("/v1/messages", anthropicBody, "Read", "NotebookEdit") - run(t, p, pctx) + pctx := pruneOnce(t, p) + finish(t, p, pctx, 0, 0, 24701) + m := findMetric(t, p.Metrics(), "$ saved") + if m.Value != 0 { + t.Errorf("$ saved = %v with no rates configured, want 0", m.Value) + } + if !strings.Contains(m.Note, "input_cost_per_token") { + t.Errorf("note = %q, want it to name the field that enables costing", m.Note) + } +} - // 1 prompt token per 4 body bytes. - pctx.Extensions.Inference.PromptTokens = len(pctx.Body) / 4 - p.OnFinish(context.Background(), pctx) +// TestMetrics_TierRatesDifferBy12x pins the reason the tiers are separate. The +// same pruned bytes, priced as a cache write versus a cache read at Anthropic's +// published ratios, differ by more than an order of magnitude. A flat rate would +// be wrong by that factor. +func TestMetrics_TierRatesDifferBy12x(t *testing.T) { + const inputRate = 15.0 / 1e6 // USD per token + cfg := func(t *testing.T) *ToolPrune { + p := New() + raw := []byte(`{"remove":["NotebookEdit"],` + + `"input_cost_per_token":1.5e-05,` + + `"cache_write_cost_per_token":1.875e-05,` + // 1.25x input + `"cache_read_cost_per_token":1.5e-06}`) // 0.1x input + if err := p.Configure(raw); err != nil { + t.Fatal(err) + } + return p + } + _ = inputRate + + write := cfg(t) + finish(t, write, pruneOnce(t, write), 0, 0, 24701) + read := cfg(t) + finish(t, read, pruneOnce(t, read), 0, 24701, 0) - m := findMetric(t, p.Metrics(), "tokens saved / request") - if m.Value <= 0 { - t.Errorf("value = %v, want a positive estimate", m.Value) + w := findMetric(t, write.Metrics(), "$ saved").Value + r := findMetric(t, read.Metrics(), "$ saved").Value + if w <= 0 || r <= 0 { + t.Fatalf("expected both priced: write=%v read=%v", w, r) } - if !strings.HasPrefix(m.Note, "estimate, n=") { - t.Errorf("note = %q, want it labelled an estimate with its sample size", m.Note) + if ratio := w / r; ratio < 12 || ratio > 13 { + t.Errorf("cache-write / cache-read cost ratio = %.2f, want ~12.5 (1.25x vs 0.1x input)", ratio) } - perReq := findMetric(t, p.Metrics(), "bytes removed / request") - if want := perReq.Value / 4; m.Value < want*0.9 || m.Value > want*1.1 { - t.Errorf("estimate %v not within 10%% of calibrated %v", m.Value, want) + // And a per-request figure alongside the total. + if pr := findMetric(t, write.Metrics(), "$ saved / request"); pr.Value <= 0 { + t.Errorf("$ saved / request = %v, want positive", pr.Value) } } diff --git a/authbridge/docs/plugin-catalog.md b/authbridge/docs/plugin-catalog.md index 8ce5f5a73..ad2b5b146 100644 --- a/authbridge/docs/plugin-catalog.md +++ b/authbridge/docs/plugin-catalog.md @@ -240,6 +240,7 @@ body-reading plugin (it rewrites the request body). Declares - `remove` (`[]string`) — tool names to delete from the manifest. The complete verdict: no learning, no state, no storage. Names absent from a given request are ignored. - `paths` (`[]string`) — request paths to act on, matched exactly or by suffix. Defaults to `/v1/chat/completions`, `/v1/completions`, `/v1/messages`. +- `input_cost_per_token`, `cache_write_cost_per_token`, `cache_read_cost_per_token` (`float`) — USD per token, for costing the saving. Names match `litellm-budget-track`; cache rates default to the input rate. All optional — with none set, no dollar figure is reported rather than a price being assumed. No output rate: pruning only shrinks the prompt. Generate the list from local transcripts with `abctl tools scan`, which proposes only tools it recognises as Claude Code built-ins and never diff --git a/authbridge/docs/tool-prune-plugin.md b/authbridge/docs/tool-prune-plugin.md index 606a64f9e..e96f27017 100644 --- a/authbridge/docs/tool-prune-plugin.md +++ b/authbridge/docs/tool-prune-plugin.md @@ -76,27 +76,62 @@ Two occasions worth it: ``` Metrics: - requests seen 3 count - requests pruned 3 count - tools removed 6 count - bytes removed 825 bytes - bytes removed / request 275 bytes - tokens saved / request 343.75 tokens estimate, n=3 - removed: NotebookEdit 3 count - removed: ScheduleWakeup 3 count + requests seen 2 count + requests pruned 2 count + tools removed 22 count + bytes removed 57,136 bytes + bytes removed / request 28,568 bytes + tokens saved: cache write 13,044 tokens estimate, n=2 + tokens saved: cache read 13,064 tokens estimate, n=2 + $ saved 0.2642 usd estimate, n=2 + $ saved / request 0.1321 usd estimate, n=2 + removed: NotebookEdit 2 count ``` In observe mode `requests projected` replaces `requests pruned`, so a projection is never mistaken for a realised saving. -Byte counts are exact. The token figure is an **estimate**, and labelled as one -with its sample size: rather than bundling a tokenizer or assuming a -bytes-per-token constant, the ratio is calibrated on your own traffic from the -response `usage` block. One approximation to state plainly — under `enforce`, -`PromptTokens` is already the post-pruning count, so the ratio is measured on -pruned requests. That is acceptable for a conversion factor, which is a property -of the tokenizer and content mix rather than of the pruning, but it is why the -number is an estimate. +### Why tokens are reported per tier and never summed + +Byte counts are exact. Tokens are an estimate, and — more importantly — they +are **not fungible**. Providers price prompt tiers very differently: Anthropic +charges 1.25x the input rate for a cache write and 0.1x for a cache read, so the +same pruned bytes are worth more than **12x** more on a cache miss than on a +cache hit. + +A single "tokens saved" figure would invite multiplying by one rate, which is +wrong by that factor. So the saving is attributed to the tier it actually came +out of and reported separately. The tool manifest sits inside the cached prefix +(Claude Code puts `cache_control` on the tool block), so a cache-miss request +saves cache-*write* tokens and a hit saves cache-*read* tokens. Traffic that +alternates shows both rows, and the honest headline is a range rather than a +point. + +The bytes-to-tokens ratio is calibrated on your own traffic — prompt tokens over +request bytes for the same request, both post-pruning so the two sides agree — +rather than bundling a tokenizer or assuming a constant. + +### Costing it + +No price is assumed. Set any of these and the `$` rows appear; leave them and +the row says so instead of inventing a figure: + +| Field | Meaning | +|---|---| +| `input_cost_per_token` | USD per uncached input token | +| `cache_write_cost_per_token` | USD per cache-write token; defaults to the input rate | +| `cache_read_cost_per_token` | USD per cache-read token; defaults to the input rate | + +Field names and semantics match +[`litellm-budget-track`](./plugin-catalog.md#litellm-budget-track), so rates are +configured once in a familiar shape. There is deliberately no output rate: +pruning only ever shrinks the prompt, so attributing output cost to it would be +false. + +If your gateway reports authoritative per-request cost (LiteLLM's +`x-litellm-response-cost`), `litellm-budget-track` is the plugin that consumes +it; this one prices from rates because a saving is a counterfactual — the cost +of a request that was never sent. Counters are in-memory and per-process. That is the right trade for the single-laptop case this targets and what keeps the plugin free of a storage From cc90c02f21e3f79c02fa1cf67ed6095e80e86e3c Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 16:40:13 -0400 Subject: [PATCH 16/28] feat(authbridge): Price tool-prune's saving per model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rates are per model, not per deployment. Measured from one gateway's own cost headers: claude-opus-5 bills input at $3.80/Mtok, aws/claude-sonnet-5 at $1.52, aws/claude-haiku-4-5 at $0.76 — a 5x spread. A single flat rate misprices the saving by that factor depending on which model served the request, and Claude Code uses more than one (its session-title calls need not be the model doing the work). So `pricing` is a map keyed by model name, each entry carrying the same three tier rates, and each request is priced at its own model's rate with the dollars accumulated. Pricing at snapshot time from blended token totals cannot express this and has been replaced. The flat fields remain as a fallback for models absent from the table, so the simpler single-model config still works. A model with no entry and no fallback is counted, not guessed: a `requests unpriced` row names the models, so an incomplete table shows as a visible gap instead of a quietly understated total. Charging it at another model's rate would be wrong by up to 5x — worse than reporting nothing. Tokens are still reported for those requests; only the dollars are withheld. Model keys are matched case-insensitively, folded once at Configure rather than per request. Gateways vary in how they echo model names and a case mismatch would silently unprice the traffic, which is the same class of invisible failure as the earlier query-string path bug. The docs also record why rates are configured rather than read from the gateway. LiteLLM reports x-litellm-response-cost: 0 for streaming responses because the total is unknown when headers are sent, and Claude Code streams every /v1/messages — so the authoritative per-request cost is unavailable for exactly the traffic this plugin prunes. litellm-budget-track hits the same wall. Beyond that, a saving is a counterfactual: the cost of a request never sent can only be priced from rates, never measured. The method for deriving real rates from non-streaming probes is documented so an operator can obtain their own. Signed-off-by: Hai Huang --- .../authlib/plugins/toolprune/metrics.go | 84 ++++++++++----- .../authlib/plugins/toolprune/plugin.go | 95 +++++++++++++---- .../authlib/plugins/toolprune/plugin_test.go | 100 ++++++++++++++++++ authbridge/docs/plugin-catalog.md | 3 +- authbridge/docs/tool-prune-plugin.md | 81 ++++++++++---- 5 files changed, 295 insertions(+), 68 deletions(-) diff --git a/authbridge/authlib/plugins/toolprune/metrics.go b/authbridge/authlib/plugins/toolprune/metrics.go index 3309e83d0..8cb3c6c3f 100644 --- a/authbridge/authlib/plugins/toolprune/metrics.go +++ b/authbridge/authlib/plugins/toolprune/metrics.go @@ -3,6 +3,7 @@ package toolprune import ( "fmt" "sort" + "strings" "sync" "github.com/rossoctl/cortex/authbridge/authlib/pipeline" @@ -32,6 +33,18 @@ type metrics struct { savedCacheWrite float64 savedCacheRead float64 requestsCosted uint64 + + // Dollars are accumulated at request time, not derived at snapshot time, + // because the rate depends on which model served the request — a 5x spread + // across opus/sonnet/haiku on one observed gateway. Multiplying a blended + // token total by any single rate would be wrong by that factor. + usdSaved float64 + + // Requests whose model had no configured rate. Counted and named rather + // than charged at another model's rate, so an incomplete pricing table + // shows up as a gap instead of silently under-reporting the total. + unpriced uint64 + unpricedModels map[string]uint64 } func (m *metrics) seen() { @@ -68,7 +81,7 @@ func (m *metrics) record(names []string, bytesRemoved int) { } } -func (m *metrics) observeSaving(tokens float64, t tier) { +func (m *metrics) observeSaving(tokens float64, t tier, usd float64, priced bool, model string) { m.mu.Lock() switch t { case tierCacheWrite: @@ -79,6 +92,18 @@ func (m *metrics) observeSaving(tokens float64, t tier) { m.savedInput += tokens } m.requestsCosted++ + if priced { + m.usdSaved += usd + } else { + m.unpriced++ + if m.unpricedModels == nil { + m.unpricedModels = make(map[string]uint64) + } + if model == "" { + model = "(unknown)" + } + m.unpricedModels[model]++ + } m.mu.Unlock() } @@ -131,24 +156,17 @@ func (m *metrics) snapshot(cfg *config) []pipeline.Metric { if m.requestsCosted > 0 { note = fmt.Sprintf("estimate, n=%d", m.requestsCosted) } - tiers := []struct { + for _, t := range []struct { name string val float64 - rate float64 }{ - {"tokens saved: cache write", m.savedCacheWrite, cfg.cacheWriteRate()}, - {"tokens saved: cache read", m.savedCacheRead, cfg.cacheReadRate()}, - {"tokens saved: input", m.savedInput, cfg.InputCostPerToken}, - } - var usd float64 - for _, t := range tiers { - if t.val <= 0 { - continue + {"tokens saved: cache write", m.savedCacheWrite}, + {"tokens saved: cache read", m.savedCacheRead}, + {"tokens saved: input", m.savedInput}, + } { + if t.val > 0 { + out = append(out, pipeline.Metric{Name: t.name, Value: t.val, Unit: "tokens", Note: note}) } - out = append(out, pipeline.Metric{ - Name: t.name, Value: t.val, Unit: "tokens", Note: note, - }) - usd += t.val * t.rate } if m.requestsCosted == 0 && acted > 0 { out = append(out, pipeline.Metric{ @@ -157,25 +175,37 @@ func (m *metrics) snapshot(cfg *config) []pipeline.Metric { }) } - // Cost only when rates are configured. Without them no price is assumed: - // a confidently wrong dollar figure is worse than none, because it gets - // quoted in decisions. - if cfg.priced() && usd > 0 { - out = append(out, pipeline.Metric{ - Name: "$ saved", Value: usd, Unit: "usd", Note: note, - }) - if m.requestsCosted > 0 { + // Dollars, accumulated per request at that request's model rate. + switch { + case m.usdSaved > 0: + out = append(out, pipeline.Metric{Name: "$ saved", Value: m.usdSaved, Unit: "usd", Note: note}) + if priced := m.requestsCosted - m.unpriced; priced > 0 { out = append(out, pipeline.Metric{ Name: "$ saved / request", - Value: usd / float64(m.requestsCosted), + Value: m.usdSaved / float64(priced), Unit: "usd", - Note: note, + Note: fmt.Sprintf("estimate, n=%d", priced), }) } - } else if acted > 0 && !cfg.priced() { + case acted > 0 && !cfg.priced(): out = append(out, pipeline.Metric{ Name: "$ saved", Value: 0, Unit: "usd", - Note: "set input_cost_per_token to cost this", + Note: "set pricing..input_cost_per_token to cost this", + }) + } + + // An incomplete pricing table is a gap in the dollar total, so name it. + if m.unpriced > 0 { + models := make([]string, 0, len(m.unpricedModels)) + for k := range m.unpricedModels { + models = append(models, k) + } + sort.Strings(models) + out = append(out, pipeline.Metric{ + Name: "requests unpriced", + Value: float64(m.unpriced), + Unit: "count", + Note: "no rate for: " + strings.Join(models, ", "), }) } diff --git a/authbridge/authlib/plugins/toolprune/plugin.go b/authbridge/authlib/plugins/toolprune/plugin.go index f26aaf83f..0621df716 100644 --- a/authbridge/authlib/plugins/toolprune/plugin.go +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -54,44 +54,95 @@ type config struct { // suffix. Defaults to the three inference endpoints. Paths []string `json:"paths" description:"Request paths to act on (exact or suffix match)."` - // Per-token rates, for costing the saving. Field names and semantics match - // litellm-budget-track so an operator configures rates in one familiar - // shape. USD per token; all optional. With none set, no cost is reported — - // a price is never assumed. + // Pricing gives per-token rates per model. Rates are per model because they + // differ enormously: on one observed gateway claude-opus-5 bills input at + // $3.80/Mtok, sonnet at $1.52 and haiku at $0.76 — a 5x spread, so one flat + // rate misprices by that factor depending on which model served the request. + // Keys match the model name the parser records + // (pctx.Extensions.Inference.Model), matched case-insensitively. + Pricing map[string]modelRates `json:"pricing" description:"Per-token rates keyed by model name."` + + // pricing is Pricing with keys lower-cased; built by applyDefaults. + pricing map[string]modelRates `json:"-"` + + // The flat fields are the fallback for models absent from Pricing. Names and + // semantics match litellm-budget-track. All optional; with nothing set no + // cost is reported rather than a price being assumed. // // There is deliberately no output rate: pruning only ever shrinks the - // prompt, so attributing any output cost to it would be false. + // prompt, so attributing output cost to it would be false. + InputCostPerToken float64 `json:"input_cost_per_token" description:"Fallback USD per uncached input token, for models absent from pricing."` + CacheWriteCostPerToken float64 `json:"cache_write_cost_per_token" description:"Fallback USD per cache-write token; defaults to input_cost_per_token."` + CacheReadCostPerToken float64 `json:"cache_read_cost_per_token" description:"Fallback USD per cache-read token; defaults to input_cost_per_token."` +} + +// modelRates is one model's prompt-tier pricing. Cache rates fall back to the +// input rate, matching litellm-budget-track — though on Anthropic-family models +// that fallback is poor (a real cache read is 0.1x input), so set them when known. +type modelRates struct { InputCostPerToken float64 `json:"input_cost_per_token" description:"USD per uncached input token."` - CacheWriteCostPerToken float64 `json:"cache_write_cost_per_token" description:"USD per cache-write input token; defaults to input_cost_per_token."` - CacheReadCostPerToken float64 `json:"cache_read_cost_per_token" description:"USD per cache-read input token; defaults to input_cost_per_token."` + CacheWriteCostPerToken float64 `json:"cache_write_cost_per_token" description:"USD per cache-write token; defaults to input_cost_per_token."` + CacheReadCostPerToken float64 `json:"cache_read_cost_per_token" description:"USD per cache-read token; defaults to input_cost_per_token."` } -// cacheWriteRate / cacheReadRate default to the uncached input rate, matching -// litellm-budget-track. Flat pricing would misstate cache-heavy traffic badly — -// Anthropic charges 1.25x input for a cache write and 0.1x for a read, so the -// same saved bytes differ by more than 12x depending on which tier they land in. -func (c *config) cacheWriteRate() float64 { - if c.CacheWriteCostPerToken > 0 { - return c.CacheWriteCostPerToken +func (r modelRates) rateFor(t tier) float64 { + switch t { + case tierCacheWrite: + if r.CacheWriteCostPerToken > 0 { + return r.CacheWriteCostPerToken + } + case tierCacheRead: + if r.CacheReadCostPerToken > 0 { + return r.CacheReadCostPerToken + } } - return c.InputCostPerToken + return r.InputCostPerToken } -func (c *config) cacheReadRate() float64 { - if c.CacheReadCostPerToken > 0 { - return c.CacheReadCostPerToken +func (r modelRates) set() bool { + return r.InputCostPerToken > 0 || r.CacheWriteCostPerToken > 0 || r.CacheReadCostPerToken > 0 +} + +// ratesFor resolves rates for a model: its own entry when present, else the flat +// fallback. Reports false when neither is configured, so the caller counts the +// request as unpriced rather than charging it at another model's rate — which on +// a 5x spread would be worse than reporting nothing. +func (c *config) ratesFor(model string) (modelRates, bool) { + if r, ok := c.pricing[strings.ToLower(model)]; ok && r.set() { + return r, true + } + fallback := modelRates{ + InputCostPerToken: c.InputCostPerToken, + CacheWriteCostPerToken: c.CacheWriteCostPerToken, + CacheReadCostPerToken: c.CacheReadCostPerToken, } - return c.InputCostPerToken + return fallback, fallback.set() } +// priced reports whether any pricing is configured at all. func (c *config) priced() bool { - return c.InputCostPerToken > 0 || c.CacheWriteCostPerToken > 0 || c.CacheReadCostPerToken > 0 + if c.InputCostPerToken > 0 || c.CacheWriteCostPerToken > 0 || c.CacheReadCostPerToken > 0 { + return true + } + for _, r := range c.Pricing { + if r.set() { + return true + } + } + return false } func (c *config) applyDefaults() { if len(c.Paths) == 0 { c.Paths = append([]string(nil), defaultPaths...) } + // Fold model keys to lower case once, so lookup is case-insensitive + // without allocating per request. Gateways vary in how they echo model + // names, and a case mismatch would silently unprice the traffic. + c.pricing = make(map[string]modelRates, len(c.Pricing)) + for k, v := range c.Pricing { + c.pricing[strings.ToLower(k)] = v + } } // ToolPrune is the plugin. Counters live in metrics, guarded by its own mutex; @@ -406,7 +457,9 @@ func (p *ToolPrune) OnFinish(_ context.Context, pctx *pipeline.Context) { if tokens <= 0 { return } - p.m.observeSaving(tokens, tierOf(inf)) + t := tierOf(inf) + rates, priced := p.cfg.ratesFor(inf.Model) + p.m.observeSaving(tokens, t, tokens*rates.rateFor(t), priced, inf.Model) } // tier names which prompt token tier a request's saving came out of. diff --git a/authbridge/authlib/plugins/toolprune/plugin_test.go b/authbridge/authlib/plugins/toolprune/plugin_test.go index 1b13f580e..d8e08eda9 100644 --- a/authbridge/authlib/plugins/toolprune/plugin_test.go +++ b/authbridge/authlib/plugins/toolprune/plugin_test.go @@ -627,3 +627,103 @@ func TestPrune_PathMismatchRecordsThePath(t *testing.T) { t.Errorf("inv = %+v, want path_not_inference with the offending path recorded", inv) } } + +// configuredJSON builds a plugin from raw config JSON. +func configuredJSON(t *testing.T, raw string) *ToolPrune { + t.Helper() + p := New() + if err := p.Configure(json.RawMessage(raw)); err != nil { + t.Fatalf("Configure: %v", err) + } + return p +} + +// pruneWithModel runs one prune and finishes it as the named model, with a +// cache-write split (the cache-miss shape). +func pruneWithModel(t *testing.T, p *ToolPrune, model string) { + t.Helper() + pctx := inferenceCtx("/v1/messages", anthropicBody, "Read", "NotebookEdit", "Bash") + run(t, p, pctx) + pctx.Extensions.Inference.Model = model + pctx.Extensions.Inference.CacheWriteTokens = 24701 + p.OnFinish(context.Background(), pctx) +} + +const perModelCfg = `{"remove":["NotebookEdit"],"pricing":{ + "claude-opus-5": {"input_cost_per_token":3.8e-06,"cache_write_cost_per_token":4.75e-06,"cache_read_cost_per_token":3.8e-07}, + "aws/claude-sonnet-5": {"input_cost_per_token":1.52e-06,"cache_write_cost_per_token":1.9e-06,"cache_read_cost_per_token":1.52e-07}, + "aws/claude-haiku-4-5":{"input_cost_per_token":7.6e-07,"cache_write_cost_per_token":9.5e-07,"cache_read_cost_per_token":7.6e-08}}}` + +// TestPricing_PerModelRatesDiffer is why pricing is keyed by model. On one +// observed gateway opus bills input at $3.80/Mtok, sonnet $1.52 and haiku $0.76 +// — a 5x spread. Charging every request at one rate would misstate the saving by +// that factor depending on which model happened to serve it. +func TestPricing_PerModelRatesDiffer(t *testing.T) { + usd := map[string]float64{} + for _, model := range []string{"claude-opus-5", "aws/claude-sonnet-5", "aws/claude-haiku-4-5"} { + p := configuredJSON(t, perModelCfg) + pruneWithModel(t, p, model) + usd[model] = findMetric(t, p.Metrics(), "$ saved").Value + if usd[model] <= 0 { + t.Fatalf("%s: no cost reported", model) + } + } + // Same saved bytes, same tier — cost must track the model's rate ratios. + if r := usd["claude-opus-5"] / usd["aws/claude-sonnet-5"]; r < 2.4 || r > 2.6 { + t.Errorf("opus/sonnet cost ratio = %.2f, want ~2.5", r) + } + if r := usd["claude-opus-5"] / usd["aws/claude-haiku-4-5"]; r < 4.9 || r > 5.1 { + t.Errorf("opus/haiku cost ratio = %.2f, want ~5.0", r) + } +} + +// TestPricing_UnknownModelIsCountedNotGuessed: charging an unpriced model at +// another model's rate would be wrong by up to 5x, so it is reported as a gap. +func TestPricing_UnknownModelIsCountedNotGuessed(t *testing.T) { + p := configuredJSON(t, perModelCfg) + pruneWithModel(t, p, "gcp/gemini-3-pro-preview") + + ms := p.Metrics() + gap := findMetric(t, ms, "requests unpriced") + if gap.Value != 1 { + t.Errorf("requests unpriced = %v, want 1", gap.Value) + } + if !strings.Contains(gap.Note, "gcp/gemini-3-pro-preview") { + t.Errorf("note should name the unpriced model, got %q", gap.Note) + } + // Tokens are still counted — only the dollars are withheld. + if findMetric(t, ms, "tokens saved: cache write").Value <= 0 { + t.Error("token saving should still be reported for an unpriced model") + } + for _, m := range ms { + if m.Name == "$ saved" && m.Value != 0 { + t.Errorf("$ saved = %v for an unpriced model, want no charge", m.Value) + } + } +} + +// TestPricing_FlatRatesActAsFallback keeps the simpler single-model config +// working: a model absent from the table is priced at the flat rates when set. +func TestPricing_FlatRatesActAsFallback(t *testing.T) { + p := configuredJSON(t, `{"remove":["NotebookEdit"],"input_cost_per_token":3.8e-06, + "pricing":{"aws/claude-haiku-4-5":{"input_cost_per_token":7.6e-07}}}`) + pruneWithModel(t, p, "some-other-model") + if findMetric(t, p.Metrics(), "$ saved").Value <= 0 { + t.Error("a model absent from pricing should fall back to the flat rates") + } + for _, m := range p.Metrics() { + if m.Name == "requests unpriced" { + t.Error("should not be counted unpriced when a fallback rate exists") + } + } +} + +// TestPricing_ModelMatchIsCaseInsensitive: gateways vary in how they echo model +// names, and a case mismatch would silently unprice the traffic. +func TestPricing_ModelMatchIsCaseInsensitive(t *testing.T) { + p := configuredJSON(t, `{"remove":["NotebookEdit"],"pricing":{"Claude-Opus-5":{"input_cost_per_token":3.8e-06}}}`) + pruneWithModel(t, p, "claude-opus-5") + if findMetric(t, p.Metrics(), "$ saved").Value <= 0 { + t.Error("model lookup should be case-insensitive") + } +} diff --git a/authbridge/docs/plugin-catalog.md b/authbridge/docs/plugin-catalog.md index ad2b5b146..d59dbbe65 100644 --- a/authbridge/docs/plugin-catalog.md +++ b/authbridge/docs/plugin-catalog.md @@ -240,7 +240,8 @@ body-reading plugin (it rewrites the request body). Declares - `remove` (`[]string`) — tool names to delete from the manifest. The complete verdict: no learning, no state, no storage. Names absent from a given request are ignored. - `paths` (`[]string`) — request paths to act on, matched exactly or by suffix. Defaults to `/v1/chat/completions`, `/v1/completions`, `/v1/messages`. -- `input_cost_per_token`, `cache_write_cost_per_token`, `cache_read_cost_per_token` (`float`) — USD per token, for costing the saving. Names match `litellm-budget-track`; cache rates default to the input rate. All optional — with none set, no dollar figure is reported rather than a price being assumed. No output rate: pruning only shrinks the prompt. +- `pricing` (`map[model]rates`) — per-token rates keyed by model name, each with `input_cost_per_token`, `cache_write_cost_per_token`, `cache_read_cost_per_token` (cache rates fall back to that model's input rate). Per model because rates differ up to 5x across opus/sonnet/haiku, so one flat rate misprices by that factor. Keys matched case-insensitively. +- `input_cost_per_token`, `cache_write_cost_per_token`, `cache_read_cost_per_token` (`float`) — optional flat fallback for models absent from `pricing`. With no pricing at all, no dollar figure is reported rather than a price being assumed; a model with no rate is counted in a `requests unpriced` row instead of being charged at another model's rate. No output rate: pruning only shrinks the prompt. Generate the list from local transcripts with `abctl tools scan`, which proposes only tools it recognises as Claude Code built-ins and never diff --git a/authbridge/docs/tool-prune-plugin.md b/authbridge/docs/tool-prune-plugin.md index e96f27017..b15bb9a03 100644 --- a/authbridge/docs/tool-prune-plugin.md +++ b/authbridge/docs/tool-prune-plugin.md @@ -113,25 +113,68 @@ rather than bundling a tokenizer or assuming a constant. ### Costing it -No price is assumed. Set any of these and the `$` rows appear; leave them and -the row says so instead of inventing a figure: - -| Field | Meaning | -|---|---| -| `input_cost_per_token` | USD per uncached input token | -| `cache_write_cost_per_token` | USD per cache-write token; defaults to the input rate | -| `cache_read_cost_per_token` | USD per cache-read token; defaults to the input rate | - -Field names and semantics match -[`litellm-budget-track`](./plugin-catalog.md#litellm-budget-track), so rates are -configured once in a familiar shape. There is deliberately no output rate: -pruning only ever shrinks the prompt, so attributing output cost to it would be -false. - -If your gateway reports authoritative per-request cost (LiteLLM's -`x-litellm-response-cost`), `litellm-budget-track` is the plugin that consumes -it; this one prices from rates because a saving is a counterfactual — the cost -of a request that was never sent. +No price is assumed. Rates are keyed **per model**, because they differ far more +than the tiers do — on one observed gateway: + +| model | input | vs opus | +|---|---|---| +| `claude-opus-5` | $3.80/Mtok | 1.0x | +| `aws/claude-sonnet-5` | $1.52/Mtok | 0.4x | +| `aws/claude-haiku-4-5` | $0.76/Mtok | 0.2x | + +A single flat rate misprices by up to 5x depending on which model served the +request, so each request is priced at its own model's rate and the dollars are +accumulated — never a blended token total multiplied by one number. + +```yaml +- name: tool-prune + config: + remove: [CronCreate, NotebookEdit] + pricing: + claude-opus-5: + input_cost_per_token: 0.0000038 + cache_write_cost_per_token: 0.00000475 + cache_read_cost_per_token: 0.00000038 + aws/claude-sonnet-5: + input_cost_per_token: 0.00000152 + cache_write_cost_per_token: 0.0000019 + cache_read_cost_per_token: 0.000000152 + # optional fallback for models absent from the table + input_cost_per_token: 0.0000038 +``` + +Model keys match what the parser records (`Extensions.Inference.Model`) and are +matched case-insensitively, since gateways vary in how they echo the name and a +case mismatch would silently unprice the traffic. + +A model with no entry and no fallback is **counted, not guessed**: the readout +grows a `requests unpriced` row naming the models, so an incomplete table shows +as a visible gap rather than a quietly understated total. Tokens are still +reported for those requests — only the dollars are withheld. + +Field names within each entry match +[`litellm-budget-track`](./plugin-catalog.md#litellm-budget-track). Cache rates +fall back to that model's input rate, though on Anthropic-family models that +fallback is poor — a real cache read is 0.1x input — so set them when known. +There is deliberately no output rate: pruning only shrinks the prompt. + +**Deriving your own rates.** If your gateway reports cost on non-streaming +responses (LiteLLM's `x-litellm-response-cost`), send two non-streaming requests +of different prompt length and difference them: `rate = Δcost / Δinput_tokens`. +Repeat with a `cache_control` block sent twice to get the write and read rates. +This is exact and specific to your deployment — it is how the numbers in the +table above were obtained, and they came out 4x below list because that gateway +bills negotiated rates. + +Why rates rather than the gateway's own number: LiteLLM reports +`x-litellm-response-cost: 0` for **streaming** responses, because the total is +not known when the headers are sent — and Claude Code streams every +`/v1/messages`. So the authoritative per-request cost is unavailable for exactly +the traffic this plugin prunes. `litellm-budget-track` hits the same wall and +falls back to configured rates for streams. + +A saving is also a counterfactual — the cost of a request that was never sent — +so even with a cost header it could only ever be priced from rates, not measured. Counters are in-memory and per-process. That is the right trade for the single-laptop case this targets and what keeps the plugin free of a storage From 7269213b5882c2088c69f1c94303674e09f01c62 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 16:51:10 -0400 Subject: [PATCH 17/28] refactor(abctl): Drop the request/response bracket glyphs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The PHASE column prefixed box-drawing corners (┌/│/└) to visually connect a request row to its response. They can only be correct for exchanges that NEST, and concurrent requests do not nest — they cross: A starts, B starts, A ends, B ends. A tree has no notation for partial overlap, so computeSpanGlyphs sorted the spans a row participates in by width and called the widest "outer" and narrowest "inner", which for crossing spans made both rows claim to contain each other. The observed output was "┌│ req" opening one exchange and "└│ resp" closing another, with no consistent reading. That was not a tuning problem. The glyphs encoded containment, the data has overlap, and the mismatch is structural — so they are removed rather than patched. Nothing replaces them: the # column already pairs exchanges exactly, by the proxy-stamped RequestID, which is what the brackets were a lossy approximation of. Reading "6 with 6, 7 with 7" off one column is both correct under concurrency and simpler than a bracket that is only correct when it happens not to overlap. Removes spanGlyph, spanLevels, prefix() and computeSpanGlyphs along with their tests. computeEventPairs keeps returning the partner map — it is what assigns a shared # to a paired request and response — so pairing is untouched and its tests, including the field-trace regression, still hold. Signed-off-by: Hai Huang --- authbridge/cmd/abctl/tui/events_pane.go | 135 ++----------------- authbridge/cmd/abctl/tui/events_pane_test.go | 125 +---------------- 2 files changed, 19 insertions(+), 241 deletions(-) diff --git a/authbridge/cmd/abctl/tui/events_pane.go b/authbridge/cmd/abctl/tui/events_pane.go index 556385288..9e24b23e1 100644 --- a/authbridge/cmd/abctl/tui/events_pane.go +++ b/authbridge/cmd/abctl/tui/events_pane.go @@ -3,7 +3,6 @@ package tui import ( "fmt" "net" - "sort" "strconv" "strings" @@ -90,17 +89,15 @@ func (m *model) rebuildEventsTable() { // shape as a plaintext call. eventRows := buildEventRows(events) - // Pair request rows with their response rows. ids drives the # column - // (one integer repeated across a request/response exchange); partner - // drives the PHASE-column span glyphs (┌/│/└) that visually bracket each - // exchange even when other events interleave between request and response. - ids, partner := computeEventPairs(eventRows) - glyphs := computeSpanGlyphs(partner, len(eventRows)) + // Pair request rows with their response rows. ids drives the # column: one + // integer repeated across a request/response exchange, which is how an + // exchange is read off the timeline. + ids, _ := computeEventPairs(eventRows) rows := make([]table.Row, 0, len(eventRows)) m.visibleRows = m.visibleRows[:0] m.hiddenInactive = 0 - for i, er := range eventRows { + for _, er := range eventRows { ev := er.event if m.filter != "" && !matchEventRow(er, m.filter) { continue @@ -121,15 +118,15 @@ func (m *model) rebuildEventsTable() { if id, ok := ids[ev]; ok { idCell = strconv.Itoa(id) } - // Prefix PHASE with the span glyph for this row's exchange. A request - // paired with a later response renders ┌; the response renders └; - // events nested between them render │ (with a second level when an - // inner exchange sits inside an outer one, e.g. inference calls inside - // an a2a message/stream). Unpaired rows get no prefix. + // PHASE carries no bracket glyphs. They were box-drawing corners + // (┌/│/└) meant to visually connect a request to its response, and they + // could only ever be correct for exchanges that NEST. Concurrent + // requests cross instead: A starts, B starts, A ends, B ends — for + // which a tree has no notation, so both rows claimed to contain each + // other and the output was actively misleading. The # column pairs + // exchanges exactly (by the proxy-stamped RequestID), which is what the + // glyphs were a lossy approximation of. phaseCell := shortPhase(ev.Phase) - if p := glyphs[i].prefix(); p != "" { - phaseCell = p + " " + phaseCell - } rows = append(rows, table.Row{ idCell, ev.At.Format("15:04:05.00"), @@ -571,112 +568,6 @@ func computeEventPairs(rows []eventRow) (map[*pipeline.SessionEvent]int, map[int return ids, partner } -// spanGlyph names which corner / side of a (request, response) exchange a row -// sits at, for the tree-style bracket in the PHASE column. rune (not byte) -// because the box-drawing characters are multi-byte in UTF-8. -type spanGlyph rune - -const ( - glyphNone spanGlyph = 0 - glyphStart spanGlyph = '┌' // request row that pairs with a later response - glyphMiddle spanGlyph = '│' // row between a paired request and its response - glyphEnd spanGlyph = '└' // response row paired with an earlier request -) - -// spanLevels holds the box-drawing glyphs for up to two nested exchanges on a -// single row. outer is the widest exchange containing the row; inner is the -// next-widest. Deeper nesting is dropped — operators only need the broad -// shape, and the PHASE column has a finite width budget. -type spanLevels struct { - outer spanGlyph - inner spanGlyph -} - -// prefix returns the concatenated rune string for the PHASE-column prefix: -// e.g. "│┌" when the row is inside an outer exchange and opens an inner one; -// "└" alone when only an outer endpoint applies; "" when the row is in no -// exchange span. -func (s spanLevels) prefix() string { - switch { - case s.outer == glyphNone: - return "" - case s.inner == glyphNone: - return string(rune(s.outer)) - default: - return string([]rune{rune(s.outer), rune(s.inner)}) - } -} - -// computeSpanGlyphs assigns each row up to two tree glyphs (outer + inner) -// from its position relative to all (request, response) exchange spans. The -// two widest spans containing the row are surfaced; deeper nesting is dropped -// so the PHASE column doesn't blow its width budget. -// -// pairs is the bidirectional map from computeEventPairs: pairs[i]=j AND -// pairs[j]=i for any matched pair (i, j). Unpaired rows are absent. n is the -// total row count. -func computeSpanGlyphs(pairs map[int]int, n int) []spanLevels { - out := make([]spanLevels, n) - if len(pairs) == 0 { - return out - } - // Collect each pair (a, b) with a < b once; the resp→req mirror entries - // are skipped. - type span struct{ a, b int } - spans := make([]span, 0, len(pairs)/2) - for a, b := range pairs { - if a < b { - spans = append(spans, span{a, b}) - } - } - - glyphAt := func(s span, i int) spanGlyph { - switch { - case i == s.a: - return glyphStart - case i == s.b: - return glyphEnd - case s.a < i && i < s.b: - return glyphMiddle - } - return glyphNone - } - - for i := range n { - // Find every span this row participates in (endpoint or strictly - // inside). - var participating []span - for _, s := range spans { - if s.a <= i && i <= s.b { - participating = append(participating, s) - } - } - if len(participating) == 0 { - continue - } - // Sort by width descending — widest first, narrowest last. Stable so - // equal-width spans keep declaration order (deterministic tests). - sort.SliceStable(participating, func(p, q int) bool { - return (participating[p].b - participating[p].a) > - (participating[q].b - participating[q].a) - }) - // outer = the widest containing span (the broadest context). inner = - // the NARROWEST containing span — the row's own tightest exchange — - // NOT the second-widest. A row that is an endpoint of a deeply-nested - // pair must still show its ┌/└ corner so its request and response - // connect visually; picking the second-widest would let an - // intermediate enclosing span's middle bar mask it. Example: a - // tools/list pair nested inside both an a2a message/stream span and a - // long-lived $transport/stream span would otherwise render "││" on - // both rows instead of "│┌" / "│└". - out[i].outer = glyphAt(participating[0], i) - if len(participating) > 1 { - out[i].inner = glyphAt(participating[len(participating)-1], i) - } - } - return out -} - // matchEventRow does a case-insensitive substring match across every string // field the operator might reasonably search for — the event's host/method, // the fields of every plugin invocation on it, and its protocol extensions. diff --git a/authbridge/cmd/abctl/tui/events_pane_test.go b/authbridge/cmd/abctl/tui/events_pane_test.go index 70d6092a3..99d068761 100644 --- a/authbridge/cmd/abctl/tui/events_pane_test.go +++ b/authbridge/cmd/abctl/tui/events_pane_test.go @@ -655,121 +655,11 @@ func TestHostOnly(t *testing.T) { } } -// TestSpanLevels_Prefix locks the PHASE-column prefix: empty levels render as -// empty string; one level renders one glyph; two levels render two glyphs. -func TestSpanLevels_Prefix(t *testing.T) { - cases := []struct { - name string - s spanLevels - want string - }{ - {"none", spanLevels{}, ""}, - {"outer only — start", spanLevels{outer: glyphStart}, "┌"}, - {"outer only — middle", spanLevels{outer: glyphMiddle}, "│"}, - {"outer only — end", spanLevels{outer: glyphEnd}, "└"}, - {"both — outer middle, inner start", spanLevels{outer: glyphMiddle, inner: glyphStart}, "│┌"}, - {"both — outer middle, inner end", spanLevels{outer: glyphMiddle, inner: glyphEnd}, "│└"}, - } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - if got := tc.s.prefix(); got != tc.want { - t.Errorf("prefix() = %q, want %q", got, tc.want) - } - }) - } -} - -// TestComputeSpanGlyphs covers per-row tree-glyph assignment for the PHASE -// column. Up to two levels of (request, response) nesting are surfaced — the -// widest containing span as outer, the next-widest as inner, deeper dropped. -func TestComputeSpanGlyphs(t *testing.T) { - none := spanLevels{} - outer := func(g spanGlyph) spanLevels { return spanLevels{outer: g} } - both := func(o, i spanGlyph) spanLevels { return spanLevels{outer: o, inner: i} } - - cases := []struct { - name string - pairs map[int]int - n int - want []spanLevels - }{ - {"no pairs", nil, 3, []spanLevels{none, none, none}}, - { - name: "adjacent pair", - pairs: map[int]int{0: 1, 1: 0}, - n: 2, - want: []spanLevels{outer(glyphStart), outer(glyphEnd)}, - }, - { - name: "one row in between", - pairs: map[int]int{0: 2, 2: 0}, - n: 3, - want: []spanLevels{outer(glyphStart), outer(glyphMiddle), outer(glyphEnd)}, - }, - { - // The real shape: an outer a2a exchange (0,5) bracketing two inner - // inference exchanges (1,2) and (3,4). - name: "nested exchanges (a2a containing two inference calls)", - pairs: map[int]int{ - 0: 5, 5: 0, - 1: 2, 2: 1, - 3: 4, 4: 3, - }, - n: 6, - want: []spanLevels{ - outer(glyphStart), - both(glyphMiddle, glyphStart), - both(glyphMiddle, glyphEnd), - both(glyphMiddle, glyphStart), - both(glyphMiddle, glyphEnd), - outer(glyphEnd), - }, - }, - { - // The #52 case: a pair (2,3) nested THREE deep — inside a middle - // span (1,4) inside an outer span (0,5). The innermost pair's - // endpoints must still show their ┌/└ corners (so its req/resp - // connect) rather than the middle span's bar masking them. inner = - // the row's narrowest containing span, not the second-widest. - name: "triple-nested innermost pair keeps its corners", - pairs: map[int]int{ - 0: 5, 5: 0, - 1: 4, 4: 1, - 2: 3, 3: 2, - }, - n: 6, - want: []spanLevels{ - outer(glyphStart), // 0: outer starts - both(glyphMiddle, glyphStart), // 1: outer mid, middle-span starts - both(glyphMiddle, glyphStart), // 2: outer mid, innermost STARTS (was masked to middle) - both(glyphMiddle, glyphEnd), // 3: outer mid, innermost ENDS (was masked to middle) - both(glyphMiddle, glyphEnd), // 4: outer mid, middle-span ends - outer(glyphEnd), // 5: outer ends - }, - }, - } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - got := computeSpanGlyphs(tc.pairs, tc.n) - if len(got) != len(tc.want) { - t.Fatalf("len = %d, want %d", len(got), len(tc.want)) - } - for i := range tc.want { - if got[i] != tc.want[i] { - t.Errorf("row %d: got {outer=%q inner=%q}, want {outer=%q inner=%q}", - i, string(rune(got[i].outer)), string(rune(got[i].inner)), - string(rune(tc.want[i].outer)), string(rune(tc.want[i].inner))) - } - } - }) - } -} - -// TestComputeEventPairs_NestedExchangeGlyphs is the end-to-end #23 shape: an +// TestComputeEventPairs_NestedExchanges is the end-to-end #23 shape: an // inbound a2a message/stream request, two outbound inference exchanges during // processing, then the a2a response. The a2a request/response must pair and -// bracket (┌ … └) with the inference exchanges nested (│┌ … │└) inside. -func TestComputeEventPairs_NestedExchangeGlyphs(t *testing.T) { +// exchange, with the inference exchanges falling inside its window. +func TestComputeEventPairs_NestedExchanges(t *testing.T) { a2aReq := pipeline.SessionEvent{Direction: pipeline.Inbound, Phase: pipeline.SessionRequest, Host: "claude-agent", A2A: &pipeline.A2AExtension{Method: "message/stream"}} infReq1 := pipeline.SessionEvent{Direction: pipeline.Outbound, Phase: pipeline.SessionRequest, @@ -795,12 +685,9 @@ func TestComputeEventPairs_NestedExchangeGlyphs(t *testing.T) { t.Errorf("a2a req/resp should share #, got %d vs %d", ids[&events[0]], ids[&events[5]]) } - glyphs := computeSpanGlyphs(partner, len(rows)) - want := []string{"┌", "│┌", "│└", "│┌", "│└", "└"} - for i, w := range want { - if got := glyphs[i].prefix(); got != w { - t.Errorf("row %d prefix = %q, want %q", i, got, w) - } + // The inner inference exchanges pair with each other, not across. + if partner[1] != 2 || partner[3] != 4 { + t.Errorf("inner exchanges should pair 1↔2 and 3↔4, got partner=%v", partner) } } From 74b8cbac4871fca3421b4719a9db466dc67367f7 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 17:03:35 -0400 Subject: [PATCH 18/28] docs: Add a laptop quickstart for cutting Claude Code token cost MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The existing quickstart installs the demo, which is the wrong foundation for someone who wants to actually run this: `--demo` regenerates cortex-ca/demo.yaml from a built-in template at startup — before it binds ports, so even a start that fails on a port clash discards edits — which means a prune list and pricing written there do not survive a restart. This is the real path instead: AUTHBRIDGE_INSTALL_ONLY=1 for the binaries, a config under ~/.cortex that nothing overwrites, `abctl tools scan --write` to fill the prune list, and the HTTPS_PROXY / NODE_EXTRA_CA_CERTS invocation. Every step was run verbatim before committing: the config parses, the CA generates, the scan writes 15 tools, the hot reload lands, and a request through the proxy arrives upstream with the configured tools removed. Pricing is shown as shape only, with zeroed placeholders. Rates are deployment-specific — a shared gateway commonly bills well below list — so the doc points at the derivation method rather than shipping numbers that would be wrong for most readers, and notes that the gateway's own per-request cost cannot substitute because LiteLLM reports 0 for streaming responses, which is all of Claude Code's traffic. Also states the two things people check first and misread: /cost drops but /context does not (client-side, computed before the request leaves), and an empty Metrics pane with every event marked `tunnel` means the CA is not trusted rather than the plugin being broken. Signed-off-by: Hai Huang --- README.md | 7 ++ authbridge/docs/laptop-token-savings.md | 130 ++++++++++++++++++++++++ 2 files changed, 137 insertions(+) create mode 100644 authbridge/docs/laptop-token-savings.md diff --git a/README.md b/README.md index 27136e11c..e1f1e3b67 100644 --- a/README.md +++ b/README.md @@ -37,6 +37,13 @@ Watch an AI agent's traffic — its model, tool, and agent-to-agent calls — de Its calls stream into `abctl`, decrypted and parsed. +## Cut Claude Code token cost on your laptop + +Already using Claude Code? Cortex can strip the tool definitions your agent never +calls out of every request — typically 20–25% of the prompt you pay for on each +turn. Four steps, about two minutes: +**[Cut Claude Code token cost](./authbridge/docs/laptop-token-savings.md)**. + ## Running on Kubernetes In a cluster, Cortex sidecars are injected automatically by the [operator](https://github.com/rossoctl/operator), with Keycloak + SPIFFE/SPIRE for identity and token exchange. Start with the end-to-end **[Weather Agent walkthrough](./authbridge/demos/weather-agent/demo-ui.md)** (or the [`abctl` version](./authbridge/demos/weather-agent/demo-with-abctl.md)); see the [demos index](./authbridge/demos/README.md) and the [architecture reference](./authbridge/README.md) for all modes and details. diff --git a/authbridge/docs/laptop-token-savings.md b/authbridge/docs/laptop-token-savings.md new file mode 100644 index 000000000..513fea16b --- /dev/null +++ b/authbridge/docs/laptop-token-savings.md @@ -0,0 +1,130 @@ +# Cut Claude Code token cost on your laptop + +Cortex runs as a local proxy in front of Claude Code and strips tool +definitions your agent never calls out of every request. Claude Code sends the +full tool manifest on every turn — tens of thousands of tokens of JSON schema, +billed each time — and the manifest is built by the client, so the proxy is the +only place to trim it without changing every client. + +Four steps, about two minutes. + +## 1. Install the binaries + +```sh +AUTHBRIDGE_INSTALL_ONLY=1 \ + curl -fsSL https://raw.githubusercontent.com/rossoctl/cortex/main/authbridge/install-demo.sh | sh +``` + +Puts `authbridge-proxy` and `abctl` in `~/.local/bin`. `INSTALL_ONLY` skips the +demo — you want a config that persists, which the next step writes. + +## 2. Write a config + +```sh +mkdir -p ~/.cortex +cat > ~/.cortex/config.yaml <<'YAML' +mode: proxy-sidecar +listener: + roles: [forward] + forward_proxy_addr: 127.0.0.1:47600 + session_api_addr: 127.0.0.1:47601 +stats: + address: 127.0.0.1:47602 +tls_bridge: + mode: enabled + ca_dir: "CA_DIR_PLACEHOLDER" + generate_ca: true +pipeline: + outbound: + plugins: + - name: inference-parser + # tool-prune must stay last: it rewrites the request body, and body + # readers have to precede it to see the original bytes. + - name: tool-prune + config: + remove: [] +YAML +sed -i.bak "s|CA_DIR_PLACEHOLDER|$HOME/.cortex/ca|" ~/.cortex/config.yaml && rm ~/.cortex/config.yaml.bak +``` + +Keep this outside any `cortex-ca/` directory. `authbridge-proxy --demo` +regenerates `cortex-ca/demo.yaml` from a built-in template on startup — before +it binds ports, so even a start that fails on a port clash discards your edits. +Running with `--config` avoids that entirely. + +## 3. Fill in the prune list and start + +```sh +authbridge-proxy --config ~/.cortex/config.yaml & +abctl tools scan --write ~/.cortex/config.yaml +``` + +`tools scan` reads your own `~/.claude/projects/*.jsonl` transcripts and proposes +the built-in tools you have not called in 30 days. It only ever proposes tools it +recognises, and never one it has seen you call — removing a tool the model needs +is the harmful direction of failure, so drift costs savings rather than +correctness. The config is hot-reloaded; no restart. + +## 4. Point Claude Code at it + +```sh +HTTPS_PROXY=http://localhost:47600 \ + NODE_EXTRA_CA_CERTS="$HOME/.cortex/ca/ca.crt" \ + CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1 \ + claude +``` + +Then watch what it saved: + +```sh +abctl --endpoint http://localhost:47601 +``` + +Plugin pane → `tool-prune` → `Metrics`. Expect `bytes removed / request` around +25–30 KB and a token saving of roughly 20–25% of your prompt. + +Stop it with `pkill -f 'authbridge-proxy --config'`. + +## Seeing the saving in money + +Token savings are reported per prompt-cache tier, never as one blended number: +providers charge ~1.25x the input rate for a cache write and ~0.1x for a cache +read, so the same saved bytes differ by more than 12x depending on cache state. +Rates also differ per model — often 5x across opus / sonnet / haiku. + +To get dollars, add your gateway's rates: + +```yaml + config: + remove: [...] + pricing: + : + input_cost_per_token: 0.0 + cache_write_cost_per_token: 0.0 + cache_read_cost_per_token: 0.0 +``` + +Rates are **deployment-specific** — a shared gateway often bills well below list +— so ask whoever runs yours, or derive them: send two non-streaming requests of +different prompt length and difference the reported cost, +`rate = Δcost / Δinput_tokens`. See +[`tool-prune-plugin.md`](./tool-prune-plugin.md#costing-it) for the full method +and why the gateway's own per-request cost cannot be used (it reports `0` for +streaming responses, which is all of Claude Code's traffic). + +With no rates set, the token rows still appear and the dollar row says so rather +than inventing a price. + +## What this does and does not change + +`/cost` and anything from the API response `usage` block **do** drop — the server +bills the request it received. + +`/context` **does not**. It is computed client-side before the request leaves, and +the pruning happens downstream. So this saves money, not context window; +auto-compact still triggers at the same point. Recovering headroom needs +client-side settings (`--allowedTools`, disabling unused MCP servers). + +If the Metrics pane stays empty and every event shows `tunnel`, Claude Code is not +trusting the bridge CA — check `NODE_EXTRA_CA_CERTS` points at the absolute path +above. The proxy also warns about this in its log after a few requests. From 581ce7e9af66a14d4aa5e500732dc96e0afee30d Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 17:05:06 -0400 Subject: [PATCH 19/28] docs: Remove one gateway's negotiated rates from the plugin docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-model pricing documentation used real measured numbers from the gateway they were derived on — $3.80/Mtok input for opus, $1.52 for sonnet, $0.76 for haiku, along with the observation that this is ~25% of Anthropic list. Those figures are one organisation's negotiated pricing, and this is a public repository; publishing them discloses a commercial discount that is not ours to disclose. Replaced with the ratios, which are what the argument actually needs: the input rate spans roughly 5x across the Claude family, so a flat rate misprices by that factor. Config examples now carry zeroed placeholders and point at the derivation method, and the "4x below list" aside becomes a general warning not to assume list pricing. Test fixtures move to synthetic round rates (1e-05 / 4e-06 / 2e-06) that preserve the ratios the assertions check — the tests verify a 5x model spread and a 12.5x cache-tier spread, neither of which needs a real price. Signed-off-by: Hai Huang --- .../authlib/plugins/toolprune/plugin.go | 6 +-- .../authlib/plugins/toolprune/plugin_test.go | 31 ++++++++------- authbridge/docs/tool-prune-plugin.md | 38 +++++++++---------- 3 files changed, 35 insertions(+), 40 deletions(-) diff --git a/authbridge/authlib/plugins/toolprune/plugin.go b/authbridge/authlib/plugins/toolprune/plugin.go index 0621df716..ea2906124 100644 --- a/authbridge/authlib/plugins/toolprune/plugin.go +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -55,9 +55,9 @@ type config struct { Paths []string `json:"paths" description:"Request paths to act on (exact or suffix match)."` // Pricing gives per-token rates per model. Rates are per model because they - // differ enormously: on one observed gateway claude-opus-5 bills input at - // $3.80/Mtok, sonnet at $1.52 and haiku at $0.76 — a 5x spread, so one flat - // rate misprices by that factor depending on which model served the request. + // differ enormously: across the Claude family the input rate spans roughly + // 5x (opus 1.0x, sonnet ~0.4x, haiku ~0.2x), so one flat rate misprices by + // that factor depending on which model served the request. // Keys match the model name the parser records // (pctx.Extensions.Inference.Model), matched case-insensitively. Pricing map[string]modelRates `json:"pricing" description:"Per-token rates keyed by model name."` diff --git a/authbridge/authlib/plugins/toolprune/plugin_test.go b/authbridge/authlib/plugins/toolprune/plugin_test.go index d8e08eda9..408501735 100644 --- a/authbridge/authlib/plugins/toolprune/plugin_test.go +++ b/authbridge/authlib/plugins/toolprune/plugin_test.go @@ -386,19 +386,17 @@ func TestMetrics_NoRatesMeansNoDollarFigure(t *testing.T) { // published ratios, differ by more than an order of magnitude. A flat rate would // be wrong by that factor. func TestMetrics_TierRatesDifferBy12x(t *testing.T) { - const inputRate = 15.0 / 1e6 // USD per token cfg := func(t *testing.T) *ToolPrune { p := New() raw := []byte(`{"remove":["NotebookEdit"],` + - `"input_cost_per_token":1.5e-05,` + - `"cache_write_cost_per_token":1.875e-05,` + // 1.25x input - `"cache_read_cost_per_token":1.5e-06}`) // 0.1x input + `"input_cost_per_token":1e-05,` + + `"cache_write_cost_per_token":1.25e-05,` + // 1.25x input + `"cache_read_cost_per_token":1e-06}`) // 0.1x input if err := p.Configure(raw); err != nil { t.Fatal(err) } return p } - _ = inputRate write := cfg(t) finish(t, write, pruneOnce(t, write), 0, 0, 24701) @@ -650,14 +648,15 @@ func pruneWithModel(t *testing.T, p *ToolPrune, model string) { } const perModelCfg = `{"remove":["NotebookEdit"],"pricing":{ - "claude-opus-5": {"input_cost_per_token":3.8e-06,"cache_write_cost_per_token":4.75e-06,"cache_read_cost_per_token":3.8e-07}, - "aws/claude-sonnet-5": {"input_cost_per_token":1.52e-06,"cache_write_cost_per_token":1.9e-06,"cache_read_cost_per_token":1.52e-07}, - "aws/claude-haiku-4-5":{"input_cost_per_token":7.6e-07,"cache_write_cost_per_token":9.5e-07,"cache_read_cost_per_token":7.6e-08}}}` - -// TestPricing_PerModelRatesDiffer is why pricing is keyed by model. On one -// observed gateway opus bills input at $3.80/Mtok, sonnet $1.52 and haiku $0.76 -// — a 5x spread. Charging every request at one rate would misstate the saving by -// that factor depending on which model happened to serve it. + "claude-opus-5": {"input_cost_per_token":1e-05,"cache_write_cost_per_token":1.25e-05,"cache_read_cost_per_token":1e-06}, + "aws/claude-sonnet-5": {"input_cost_per_token":4e-06,"cache_write_cost_per_token":5e-06,"cache_read_cost_per_token":4e-07}, + "aws/claude-haiku-4-5":{"input_cost_per_token":2e-06,"cache_write_cost_per_token":2.5e-06,"cache_read_cost_per_token":2e-07}}}` + +// TestPricing_PerModelRatesDiffer is why pricing is keyed by model. Across the +// Claude family the input rate spans roughly 5x (opus 1.0x, sonnet ~0.4x, haiku +// ~0.2x). Charging every request at one rate would misstate the saving by that +// factor depending on which model happened to serve it. The rates below are +// synthetic, chosen to reproduce those ratios exactly. func TestPricing_PerModelRatesDiffer(t *testing.T) { usd := map[string]float64{} for _, model := range []string{"claude-opus-5", "aws/claude-sonnet-5", "aws/claude-haiku-4-5"} { @@ -705,8 +704,8 @@ func TestPricing_UnknownModelIsCountedNotGuessed(t *testing.T) { // TestPricing_FlatRatesActAsFallback keeps the simpler single-model config // working: a model absent from the table is priced at the flat rates when set. func TestPricing_FlatRatesActAsFallback(t *testing.T) { - p := configuredJSON(t, `{"remove":["NotebookEdit"],"input_cost_per_token":3.8e-06, - "pricing":{"aws/claude-haiku-4-5":{"input_cost_per_token":7.6e-07}}}`) + p := configuredJSON(t, `{"remove":["NotebookEdit"],"input_cost_per_token":1e-05, + "pricing":{"aws/claude-haiku-4-5":{"input_cost_per_token":2e-06}}}`) pruneWithModel(t, p, "some-other-model") if findMetric(t, p.Metrics(), "$ saved").Value <= 0 { t.Error("a model absent from pricing should fall back to the flat rates") @@ -721,7 +720,7 @@ func TestPricing_FlatRatesActAsFallback(t *testing.T) { // TestPricing_ModelMatchIsCaseInsensitive: gateways vary in how they echo model // names, and a case mismatch would silently unprice the traffic. func TestPricing_ModelMatchIsCaseInsensitive(t *testing.T) { - p := configuredJSON(t, `{"remove":["NotebookEdit"],"pricing":{"Claude-Opus-5":{"input_cost_per_token":3.8e-06}}}`) + p := configuredJSON(t, `{"remove":["NotebookEdit"],"pricing":{"Claude-Opus-5":{"input_cost_per_token":1e-05}}}`) pruneWithModel(t, p, "claude-opus-5") if findMetric(t, p.Metrics(), "$ saved").Value <= 0 { t.Error("model lookup should be case-insensitive") diff --git a/authbridge/docs/tool-prune-plugin.md b/authbridge/docs/tool-prune-plugin.md index b15bb9a03..066abc42c 100644 --- a/authbridge/docs/tool-prune-plugin.md +++ b/authbridge/docs/tool-prune-plugin.md @@ -114,33 +114,29 @@ rather than bundling a tokenizer or assuming a constant. ### Costing it No price is assumed. Rates are keyed **per model**, because they differ far more -than the tiers do — on one observed gateway: - -| model | input | vs opus | -|---|---|---| -| `claude-opus-5` | $3.80/Mtok | 1.0x | -| `aws/claude-sonnet-5` | $1.52/Mtok | 0.4x | -| `aws/claude-haiku-4-5` | $0.76/Mtok | 0.2x | - -A single flat rate misprices by up to 5x depending on which model served the -request, so each request is priced at its own model's rate and the dollars are -accumulated — never a blended token total multiplied by one number. +than the tiers do. On one gateway the input rate varied 5x across the +Claude family — opus at 1.0x, sonnet at ~0.4x, haiku at ~0.2x — so a single flat +rate misprices the saving by that factor depending on which model served the +request. Each request is therefore priced at its own model's rate and the dollars +accumulated, never a blended token total multiplied by one number. ```yaml - name: tool-prune config: remove: [CronCreate, NotebookEdit] pricing: + # Your gateway's rates, in USD per token. The values below are + # placeholders — see "Deriving your own rates" below. claude-opus-5: - input_cost_per_token: 0.0000038 - cache_write_cost_per_token: 0.00000475 - cache_read_cost_per_token: 0.00000038 + input_cost_per_token: 0.0 + cache_write_cost_per_token: 0.0 + cache_read_cost_per_token: 0.0 aws/claude-sonnet-5: - input_cost_per_token: 0.00000152 - cache_write_cost_per_token: 0.0000019 - cache_read_cost_per_token: 0.000000152 + input_cost_per_token: 0.0 + cache_write_cost_per_token: 0.0 + cache_read_cost_per_token: 0.0 # optional fallback for models absent from the table - input_cost_per_token: 0.0000038 + input_cost_per_token: 0.0 ``` Model keys match what the parser records (`Extensions.Inference.Model`) and are @@ -162,9 +158,9 @@ There is deliberately no output rate: pruning only shrinks the prompt. responses (LiteLLM's `x-litellm-response-cost`), send two non-streaming requests of different prompt length and difference them: `rate = Δcost / Δinput_tokens`. Repeat with a `cache_control` block sent twice to get the write and read rates. -This is exact and specific to your deployment — it is how the numbers in the -table above were obtained, and they came out 4x below list because that gateway -bills negotiated rates. +This is exact and specific to your deployment. Do not assume list pricing: a +shared or enterprise gateway commonly bills at negotiated rates well below it, +and using list would overstate the saving by whatever that discount is. Why rates rather than the gateway's own number: LiteLLM reports `x-litellm-response-cost: 0` for **streaming** responses, because the total is From 48d78f55fcb99c8b4079ef32c4211ad0da184664 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 17:12:20 -0400 Subject: [PATCH 20/28] feat(authbridge): Ship default per-model rates so cost works unconfigured MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tool-prune reported token savings but no dollars until an operator supplied rates, and an optional setup step that stands between someone and the number they came for mostly does not get taken. The plugin now carries a rate table for the Claude models on the rossoctl LiteLLM gateway, measured from its own x-litellm-response-cost headers, so `$ saved` and `$ saved / request` appear with no configuration at all. Resolution is most-specific-first: an explicit pricing entry for the model, then the flat fallback, then the built-in table. Config always wins outright, so an operator on a different gateway corrects a model without deleting anything. Defaults are a starting point, not a fact about anyone's account — they are gateway-specific, that gateway bills below vendor list, and nothing refreshes them. So provenance travels with the figure: any dollar row derived from the table carries "default rates — set pricing. to use yours", and stops saying so once the model is configured. A model in neither the table nor the config is still counted in `requests unpriced` rather than charged at another model's rate, since the 5x spread across this family makes a wrong rate worse than no figure. The laptop quickstart drops its pricing block entirely as a result: install, write a config, scan, point Claude Code at it, and the saving shows in dollars. Verified by running the documented flow with no pricing configured: a pruned request reports $0.128918 saved, noted as default rates. Signed-off-by: Hai Huang --- .../authlib/plugins/toolprune/metrics.go | 35 ++++++---- .../authlib/plugins/toolprune/plugin.go | 42 ++++++++---- .../authlib/plugins/toolprune/plugin_test.go | 66 ++++++++++++++++--- .../authlib/plugins/toolprune/pricing.go | 51 ++++++++++++++ authbridge/docs/laptop-token-savings.md | 34 +++------- authbridge/docs/plugin-catalog.md | 4 +- authbridge/docs/tool-prune-plugin.md | 52 ++++++++++----- 7 files changed, 208 insertions(+), 76 deletions(-) create mode 100644 authbridge/authlib/plugins/toolprune/pricing.go diff --git a/authbridge/authlib/plugins/toolprune/metrics.go b/authbridge/authlib/plugins/toolprune/metrics.go index 8cb3c6c3f..67e0c3b55 100644 --- a/authbridge/authlib/plugins/toolprune/metrics.go +++ b/authbridge/authlib/plugins/toolprune/metrics.go @@ -45,6 +45,12 @@ type metrics struct { // shows up as a gap instead of silently under-reporting the total. unpriced uint64 unpricedModels map[string]uint64 + + // usedDefaultRates records that at least one request was priced from the + // built-in table rather than operator config, so the readout can say so. + // A dollar figure that silently mixes measured and assumed rates invites + // being quoted as though it were measured. + usedDefaultRates bool } func (m *metrics) seen() { @@ -81,7 +87,7 @@ func (m *metrics) record(names []string, bytesRemoved int) { } } -func (m *metrics) observeSaving(tokens float64, t tier, usd float64, priced bool, model string) { +func (m *metrics) observeSaving(tokens float64, t tier, usd float64, src rateSource, model string) { m.mu.Lock() switch t { case tierCacheWrite: @@ -92,8 +98,11 @@ func (m *metrics) observeSaving(tokens float64, t tier, usd float64, priced bool m.savedInput += tokens } m.requestsCosted++ - if priced { + if src != rateNone { m.usdSaved += usd + if src == rateDefault { + m.usedDefaultRates = true + } } else { m.unpriced++ if m.unpricedModels == nil { @@ -176,22 +185,26 @@ func (m *metrics) snapshot(cfg *config) []pipeline.Metric { } // Dollars, accumulated per request at that request's model rate. - switch { - case m.usdSaved > 0: - out = append(out, pipeline.Metric{Name: "$ saved", Value: m.usdSaved, Unit: "usd", Note: note}) + if m.usdSaved > 0 { + costNote := note + if m.usedDefaultRates { + // Provenance travels with the number. Built-in rates are + // gateway-specific and not refreshed, so a figure derived from them + // must not read as one measured on this account. + costNote = "default rates — set pricing. to use yours" + if note != "" { + costNote = note + "; " + costNote + } + } + out = append(out, pipeline.Metric{Name: "$ saved", Value: m.usdSaved, Unit: "usd", Note: costNote}) if priced := m.requestsCosted - m.unpriced; priced > 0 { out = append(out, pipeline.Metric{ Name: "$ saved / request", Value: m.usdSaved / float64(priced), Unit: "usd", - Note: fmt.Sprintf("estimate, n=%d", priced), + Note: costNote, }) } - case acted > 0 && !cfg.priced(): - out = append(out, pipeline.Metric{ - Name: "$ saved", Value: 0, Unit: "usd", - Note: "set pricing..input_cost_per_token to cost this", - }) } // An incomplete pricing table is a gap in the dollar total, so name it. diff --git a/authbridge/authlib/plugins/toolprune/plugin.go b/authbridge/authlib/plugins/toolprune/plugin.go index ea2906124..eae4aa176 100644 --- a/authbridge/authlib/plugins/toolprune/plugin.go +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -103,24 +103,42 @@ func (r modelRates) set() bool { return r.InputCostPerToken > 0 || r.CacheWriteCostPerToken > 0 || r.CacheReadCostPerToken > 0 } -// ratesFor resolves rates for a model: its own entry when present, else the flat -// fallback. Reports false when neither is configured, so the caller counts the -// request as unpriced rather than charging it at another model's rate — which on -// a 5x spread would be worse than reporting nothing. -func (c *config) ratesFor(model string) (modelRates, bool) { - if r, ok := c.pricing[strings.ToLower(model)]; ok && r.set() { - return r, true +// rateSource names where a request's rates came from, so a reported figure can +// carry its own provenance instead of looking equally authoritative either way. +type rateSource int + +const ( + rateNone rateSource = iota // no rates for this model + rateConfigured // operator-supplied, for this model or via the flat fallback + rateDefault // built-in table; see pricing.go +) + +// ratesFor resolves rates for a model, most specific first: an explicit pricing +// entry, then the flat fallback, then the built-in defaults. Explicit config +// always wins so an operator on a different gateway can correct the defaults +// per model without deleting anything. +func (c *config) ratesFor(model string) (modelRates, rateSource) { + key := strings.ToLower(model) + if r, ok := c.pricing[key]; ok && r.set() { + return r, rateConfigured } fallback := modelRates{ InputCostPerToken: c.InputCostPerToken, CacheWriteCostPerToken: c.CacheWriteCostPerToken, CacheReadCostPerToken: c.CacheReadCostPerToken, } - return fallback, fallback.set() + if fallback.set() { + return fallback, rateConfigured + } + if r, ok := defaultPricing[key]; ok { + return r, rateDefault + } + return modelRates{}, rateNone } -// priced reports whether any pricing is configured at all. -func (c *config) priced() bool { +// configuredPricing reports whether the operator supplied any rates of their +// own, as opposed to relying on the built-in defaults. +func (c *config) configuredPricing() bool { if c.InputCostPerToken > 0 || c.CacheWriteCostPerToken > 0 || c.CacheReadCostPerToken > 0 { return true } @@ -458,8 +476,8 @@ func (p *ToolPrune) OnFinish(_ context.Context, pctx *pipeline.Context) { return } t := tierOf(inf) - rates, priced := p.cfg.ratesFor(inf.Model) - p.m.observeSaving(tokens, t, tokens*rates.rateFor(t), priced, inf.Model) + rates, src := p.cfg.ratesFor(inf.Model) + p.m.observeSaving(tokens, t, tokens*rates.rateFor(t), src, inf.Model) } // tier names which prompt token tier a request's saving came out of. diff --git a/authbridge/authlib/plugins/toolprune/plugin_test.go b/authbridge/authlib/plugins/toolprune/plugin_test.go index 408501735..238893944 100644 --- a/authbridge/authlib/plugins/toolprune/plugin_test.go +++ b/authbridge/authlib/plugins/toolprune/plugin_test.go @@ -366,18 +366,64 @@ func TestMetrics_AttributesSavingToTheRightTier(t *testing.T) { } } -// TestMetrics_NoRatesMeansNoDollarFigure: a price is never assumed. An -// unconfigured plugin says so instead of inventing one. -func TestMetrics_NoRatesMeansNoDollarFigure(t *testing.T) { - p := configured(t, "NotebookEdit") - pctx := pruneOnce(t, p) - finish(t, p, pctx, 0, 0, 24701) +// TestPricing_DefaultsPriceKnownModelsWithoutConfig: the built-in table exists +// so a dollar figure appears with no configuration at all — the difference +// between a number an operator sees and one they never get around to enabling. +func TestPricing_DefaultsPriceKnownModelsWithoutConfig(t *testing.T) { + p := configured(t, "NotebookEdit") // no pricing configured whatsoever + pruneWithModel(t, p, "claude-opus-5") + m := findMetric(t, p.Metrics(), "$ saved") - if m.Value != 0 { - t.Errorf("$ saved = %v with no rates configured, want 0", m.Value) + if m.Value <= 0 { + t.Errorf("$ saved = %v, want a figure from the built-in rates", m.Value) + } + // Provenance must travel with the number: built-in rates are + // gateway-specific and never refreshed, so this must not read as measured. + if !strings.Contains(m.Note, "default rates") { + t.Errorf("note = %q, want it to disclose that default rates were used", m.Note) + } + if !strings.Contains(m.Note, "pricing.") { + t.Errorf("note = %q, want it to name how to override", m.Note) + } +} + +// TestPricing_ConfigOverridesDefaults: an operator on a different gateway must +// be able to correct a model without the built-in value leaking through, and the +// note must stop claiming defaults were used. +func TestPricing_ConfigOverridesDefaults(t *testing.T) { + base := configured(t, "NotebookEdit") + pruneWithModel(t, base, "claude-opus-5") + fromDefault := findMetric(t, base.Metrics(), "$ saved").Value + + // Ten times the built-in input rate. + over := configuredJSON(t, `{"remove":["NotebookEdit"], + "pricing":{"claude-opus-5":{"input_cost_per_token":3.8e-05,"cache_write_cost_per_token":4.75e-05}}}`) + pruneWithModel(t, over, "claude-opus-5") + m := findMetric(t, over.Metrics(), "$ saved") + + if ratio := m.Value / fromDefault; ratio < 9.5 || ratio > 10.5 { + t.Errorf("configured/default cost ratio = %.2f, want ~10 — config must win outright", ratio) } - if !strings.Contains(m.Note, "input_cost_per_token") { - t.Errorf("note = %q, want it to name the field that enables costing", m.Note) + if strings.Contains(m.Note, "default rates") { + t.Errorf("note = %q, must not claim defaults when the operator configured the model", m.Note) + } +} + +// TestPricing_UnknownModelStillUnpriced: the defaults cover a known set, not +// everything. A model in neither the table nor the config is counted, not +// charged at some other model's rate. +func TestPricing_UnknownModelStillUnpriced(t *testing.T) { + p := configured(t, "NotebookEdit") + pruneWithModel(t, p, "gcp/gemini-3-pro-preview") + + gap := findMetric(t, p.Metrics(), "requests unpriced") + if gap.Value != 1 || !strings.Contains(gap.Note, "gemini") { + t.Errorf("unpriced row = %+v, want 1 naming the model", gap) + } + for _, m := range p.Metrics() { + if m.Name == "$ saved" { + t.Errorf("$ saved = %v for a model with no rate anywhere, want no row", m.Value) + } } } diff --git a/authbridge/authlib/plugins/toolprune/pricing.go b/authbridge/authlib/plugins/toolprune/pricing.go new file mode 100644 index 000000000..e4b5e4a78 --- /dev/null +++ b/authbridge/authlib/plugins/toolprune/pricing.go @@ -0,0 +1,51 @@ +package toolprune + +// defaultPricing holds per-token rates for the models seen on the rossoctl +// LiteLLM gateway, measured from its own x-litellm-response-cost headers: +// send two non-streaming requests of differing prompt length and difference +// them, rate = Δcost / Δinput_tokens; the cache tiers were obtained the same +// way with a cache_control block sent twice. +// +// These exist so `$ saved` works with no configuration, which is the difference +// between a number an operator sees and one they never get around to enabling. +// They are a starting point, not a fact about your account: +// +// - Rates are gateway-specific. This gateway bills well below Anthropic list; +// a deployment talking straight to the vendor pays more, so these would +// understate its saving. +// - Rates change. Nothing here refreshes them. +// +// A figure derived from these is therefore labelled as coming from default +// rates wherever it is reported, so it cannot be mistaken for one measured on +// the operator's own account. Any `pricing` entry in config overrides the +// matching model outright. +// +// Keys are lower-case; lookup folds the observed model name the same way. +var defaultPricing = map[string]modelRates{ + // input 1.00x / cache write 1.25x / cache read 0.10x + "claude-opus-5": { + InputCostPerToken: 0.0000038, + CacheWriteCostPerToken: 0.00000475, + CacheReadCostPerToken: 0.00000038, + }, + "aws/claude-opus-5": { + InputCostPerToken: 0.0000038, + CacheWriteCostPerToken: 0.00000475, + CacheReadCostPerToken: 0.00000038, + }, + "aws/claude-sonnet-5": { + InputCostPerToken: 0.00000152, + CacheWriteCostPerToken: 0.0000019, + CacheReadCostPerToken: 0.000000152, + }, + "aws/claude-haiku-4-5": { + InputCostPerToken: 0.00000076, + CacheWriteCostPerToken: 0.00000095, + CacheReadCostPerToken: 0.000000076, + }, + "claude-haiku-4-5-20251001": { + InputCostPerToken: 0.00000076, + CacheWriteCostPerToken: 0.00000095, + CacheReadCostPerToken: 0.000000076, + }, +} diff --git a/authbridge/docs/laptop-token-savings.md b/authbridge/docs/laptop-token-savings.md index 513fea16b..f68129977 100644 --- a/authbridge/docs/laptop-token-savings.md +++ b/authbridge/docs/laptop-token-savings.md @@ -87,33 +87,19 @@ Stop it with `pkill -f 'authbridge-proxy --config'`. ## Seeing the saving in money +`$ saved` and `$ saved / request` appear with no extra configuration — the plugin +ships rates for the Claude models on the rossoctl gateway. The figure is labelled +`default rates` to be clear it comes from a built-in table rather than your own +account. + Token savings are reported per prompt-cache tier, never as one blended number: providers charge ~1.25x the input rate for a cache write and ~0.1x for a cache -read, so the same saved bytes differ by more than 12x depending on cache state. -Rates also differ per model — often 5x across opus / sonnet / haiku. - -To get dollars, add your gateway's rates: - -```yaml - config: - remove: [...] - pricing: - : - input_cost_per_token: 0.0 - cache_write_cost_per_token: 0.0 - cache_read_cost_per_token: 0.0 -``` - -Rates are **deployment-specific** — a shared gateway often bills well below list -— so ask whoever runs yours, or derive them: send two non-streaming requests of -different prompt length and difference the reported cost, -`rate = Δcost / Δinput_tokens`. See -[`tool-prune-plugin.md`](./tool-prune-plugin.md#costing-it) for the full method -and why the gateway's own per-request cost cannot be used (it reports `0` for -streaming responses, which is all of Claude Code's traffic). +read, so identical saved bytes differ by more than 12x depending on cache state. -With no rates set, the token rows still appear and the dollar row says so rather -than inventing a price. +If you are on a different gateway, or the rates have moved, override them per +model — see +[`tool-prune-plugin.md`](./tool-prune-plugin.md#costing-it), which also has the +method for measuring your own from the gateway's cost headers. ## What this does and does not change diff --git a/authbridge/docs/plugin-catalog.md b/authbridge/docs/plugin-catalog.md index d59dbbe65..62c5ee8db 100644 --- a/authbridge/docs/plugin-catalog.md +++ b/authbridge/docs/plugin-catalog.md @@ -240,8 +240,8 @@ body-reading plugin (it rewrites the request body). Declares - `remove` (`[]string`) — tool names to delete from the manifest. The complete verdict: no learning, no state, no storage. Names absent from a given request are ignored. - `paths` (`[]string`) — request paths to act on, matched exactly or by suffix. Defaults to `/v1/chat/completions`, `/v1/completions`, `/v1/messages`. -- `pricing` (`map[model]rates`) — per-token rates keyed by model name, each with `input_cost_per_token`, `cache_write_cost_per_token`, `cache_read_cost_per_token` (cache rates fall back to that model's input rate). Per model because rates differ up to 5x across opus/sonnet/haiku, so one flat rate misprices by that factor. Keys matched case-insensitively. -- `input_cost_per_token`, `cache_write_cost_per_token`, `cache_read_cost_per_token` (`float`) — optional flat fallback for models absent from `pricing`. With no pricing at all, no dollar figure is reported rather than a price being assumed; a model with no rate is counted in a `requests unpriced` row instead of being charged at another model's rate. No output rate: pruning only shrinks the prompt. +- `pricing` (`map[model]rates`) — per-token rates keyed by model name, each with `input_cost_per_token`, `cache_write_cost_per_token`, `cache_read_cost_per_token`. **Optional**: a built-in table covers the Claude models on the rossoctl gateway, so `$ saved` works unconfigured; any entry here overrides the built-in for that model. Per model because rates differ ~5x across opus/sonnet/haiku. Keys matched case-insensitively. +- `input_cost_per_token`, `cache_write_cost_per_token`, `cache_read_cost_per_token` (`float`) — optional flat fallback for models absent from `pricing`. A figure from built-in rates is labelled as such; a model in neither the table nor config is counted in a `requests unpriced` row instead of charged at another model's rate. No output rate: pruning only shrinks the prompt. Generate the list from local transcripts with `abctl tools scan`, which proposes only tools it recognises as Claude Code built-ins and never diff --git a/authbridge/docs/tool-prune-plugin.md b/authbridge/docs/tool-prune-plugin.md index 066abc42c..174daf80f 100644 --- a/authbridge/docs/tool-prune-plugin.md +++ b/authbridge/docs/tool-prune-plugin.md @@ -113,32 +113,50 @@ rather than bundling a tokenizer or assuming a constant. ### Costing it -No price is assumed. Rates are keyed **per model**, because they differ far more -than the tiers do. On one gateway the input rate varied 5x across the -Claude family — opus at 1.0x, sonnet at ~0.4x, haiku at ~0.2x — so a single flat -rate misprices the saving by that factor depending on which model served the -request. Each request is therefore priced at its own model's rate and the dollars -accumulated, never a blended token total multiplied by one number. +**Dollars work out of the box.** The plugin ships a rate table measured from the +rossoctl LiteLLM gateway, so `$ saved` appears with no configuration: + +| model | input | cache write (1.25x) | cache read (0.10x) | +|---|---|---|---| +| `claude-opus-5`, `aws/claude-opus-5` | $3.80/Mtok | $4.75/Mtok | $0.38/Mtok | +| `aws/claude-sonnet-5` | $1.52/Mtok | $1.90/Mtok | $0.152/Mtok | +| `aws/claude-haiku-4-5`, `claude-haiku-4-5-20251001` | $0.76/Mtok | $0.95/Mtok | $0.076/Mtok | + +Rates are keyed **per model** because they differ far more than the tiers do — +5x across this family — so a single flat rate would misprice the saving by that +factor depending on which model served the request. Each request is priced at its +own model's rate and the dollars accumulated, never a blended token total +multiplied by one number. + +Any figure derived from these carries `default rates — set pricing. to use +yours` in its note, because they are a starting point rather than a fact about +your account: they are specific to that gateway (which bills below vendor list), +and nothing refreshes them when they change. A model in neither the table nor +your config is reported in a `requests unpriced` row rather than charged at +another model's rate. + +To use your own, add a `pricing` entry — it overrides the built-in value for that +model outright: ```yaml - name: tool-prune config: remove: [CronCreate, NotebookEdit] pricing: - # Your gateway's rates, in USD per token. The values below are - # placeholders — see "Deriving your own rates" below. claude-opus-5: - input_cost_per_token: 0.0 - cache_write_cost_per_token: 0.0 - cache_read_cost_per_token: 0.0 - aws/claude-sonnet-5: - input_cost_per_token: 0.0 - cache_write_cost_per_token: 0.0 - cache_read_cost_per_token: 0.0 - # optional fallback for models absent from the table - input_cost_per_token: 0.0 + input_cost_per_token: 0.0000038 + cache_write_cost_per_token: 0.00000475 + cache_read_cost_per_token: 0.00000038 + # optional flat fallback for models absent from the table above + input_cost_per_token: 0.0000038 ``` +**Deriving your own rates.** If your gateway reports cost on non-streaming +responses (LiteLLM's `x-litellm-response-cost`), send two non-streaming requests +of different prompt length and difference them: `rate = Δcost / Δinput_tokens`. +Repeat with a `cache_control` block sent twice for the write and read rates. This +is how the table above was obtained, and it is exact for your deployment. + Model keys match what the parser records (`Extensions.Inference.Model`) and are matched case-insensitively, since gateways vary in how they echo the name and a case mismatch would silently unprice the traffic. From 834c51abe685af8f8fbebc228df9603ce6b6e74d Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 17:29:17 -0400 Subject: [PATCH 21/28] feat: Show per-request token and cost saving in abctl's TOKENS column MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The saving was only visible as a pane aggregate, which hides the thing worth seeing: cache-miss turns save an order of magnitude more than cache-hit turns, so an average describes neither. The events column now reads "33,604 −24.7k $0.117" — the request's own total, what tool-prune removed from it, and what that was worth. The two halves of the calculation necessarily live on different events. The byte saving is known when the request is rewritten; the tier it came out of, and the ratio converting bytes to tokens, only from the response. Emitting the finished figure from OnFinish is not an option: the listener defers RunFinish to the return of serveOutbound, so it runs after the response event is already recorded. So the plugin publishes what it knows at request time — bytes removed, post-prune body size, model, and the resolved per-tier rates — under "tool-prune/event". Carrying the rates rather than a dollar amount means a consumer needs no knowledge of the built-in default table, and abctl pairs request to response on the proxy-stamped RequestID (exact, including under the concurrency that made the old bracket glyphs unreadable) to finish it. Cost uses the tier the request actually used, so the ~12.5x spread between a cache write and a cache read lands on the right row. A model with no rate anywhere shows the token saving with no dollar figure rather than one priced at another model's rate. Sub-cent amounts format to four decimals: a per-request saving is often fractions of a cent, where %.2f would round every row to "0.00". Verified against a proxy alternating cache miss and hit on the same pruned request: $0.290 and $0.023, the 12.5x visible row to row. Signed-off-by: Hai Huang --- authbridge/authlib/plugins/toolprune/event.go | 45 +++++++ .../authlib/plugins/toolprune/plugin.go | 21 +++ .../authlib/plugins/toolprune/pricing.go | 11 ++ authbridge/cmd/abctl/tui/events_pane.go | 42 +++++- authbridge/cmd/abctl/tui/prune_saving.go | 118 +++++++++++++++++ authbridge/cmd/abctl/tui/prune_saving_test.go | 124 ++++++++++++++++++ authbridge/docs/tool-prune-plugin.md | 23 ++++ 7 files changed, 380 insertions(+), 4 deletions(-) create mode 100644 authbridge/authlib/plugins/toolprune/event.go create mode 100644 authbridge/cmd/abctl/tui/prune_saving.go create mode 100644 authbridge/cmd/abctl/tui/prune_saving_test.go diff --git a/authbridge/authlib/plugins/toolprune/event.go b/authbridge/authlib/plugins/toolprune/event.go new file mode 100644 index 000000000..19310efb7 --- /dev/null +++ b/authbridge/authlib/plugins/toolprune/event.go @@ -0,0 +1,45 @@ +package toolprune + +import "github.com/rossoctl/cortex/authbridge/authlib/pipeline" + +// pruneEvent is the per-request record published under "tool-prune/event", so a +// consumer can show what this one request saved instead of only an aggregate. +// +// It deliberately carries the applicable rates rather than a finished dollar +// figure. The dollar amount depends on which prompt-cache tier the saving came +// out of, and that is only known from the response — so the request-side event +// supplies the inputs and the consumer, which can pair request to response by +// RequestID, does the last step. Carrying the rates also means a consumer needs +// no knowledge of the built-in default table. +// +// No body content: the session store is unauthenticated, so this holds counts, +// tool names the operator themselves configured, and rates. +type pruneEvent struct { + ToolsRemoved []string `json:"toolsRemoved,omitempty"` + BytesRemoved int `json:"bytesRemoved"` + BodyBytesAfter int `json:"bodyBytesAfter"` + Model string `json:"model,omitempty"` + + // Rates are USD per token for this request's model, already resolved + // through config → flat fallback → built-in defaults. + RateInput float64 `json:"rateInput,omitempty"` + RateCacheWrite float64 `json:"rateCacheWrite,omitempty"` + RateCacheRead float64 `json:"rateCacheRead,omitempty"` + RateSource string `json:"rateSource,omitempty"` // configured | default | none +} + +func (p *ToolPrune) publish(pctx *pipeline.Context, ev pruneEvent) { + if pctx.Extensions.Custom == nil { + pctx.Extensions.Custom = map[string]any{} + } + pctx.Extensions.Custom[p.Name()+pipeline.PluginEventSuffix] = ev +} + +// inferenceModel returns the model the parser recorded, or "" when no parser has +// run — in which case rate lookup falls through to the flat fallback. +func inferenceModel(pctx *pipeline.Context) string { + if pctx.Extensions.Inference == nil { + return "" + } + return pctx.Extensions.Inference.Model +} diff --git a/authbridge/authlib/plugins/toolprune/plugin.go b/authbridge/authlib/plugins/toolprune/plugin.go index eae4aa176..270840050 100644 --- a/authbridge/authlib/plugins/toolprune/plugin.go +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -414,6 +414,27 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action } removedBytes := len(body) - len(out) + // Publish the per-request saving so a UI can show it on the row rather than + // only in an aggregate pane. Emitted here, in OnRequest, because the + // listener records the response session event before the deferred + // RunFinish, so anything published from OnFinish arrives too late to appear. + // + // Everything except the token tier is known now: inference-parser runs + // earlier in the chain and has already set the model, so the applicable + // rates resolve here. The consumer pairs this with the response event + // (matching on RequestID) to get the prompt token total and which tier the + // saving came out of, and finishes the arithmetic. + rates, src := p.cfg.ratesFor(inferenceModel(pctx)) + p.publish(pctx, pruneEvent{ + ToolsRemoved: names, + BytesRemoved: removedBytes, + BodyBytesAfter: len(out), + Model: inferenceModel(pctx), + RateInput: rates.InputCostPerToken, + RateCacheWrite: rates.rateFor(tierCacheWrite), + RateCacheRead: rates.rateFor(tierCacheRead), + RateSource: src.String(), + }) // Carry the saving to OnFinish, where the response reveals which token tier // it came out of. SetState keeps it private to this plugin, unlike // Extensions.Custom which is shared. diff --git a/authbridge/authlib/plugins/toolprune/pricing.go b/authbridge/authlib/plugins/toolprune/pricing.go index e4b5e4a78..4b8c2ae8c 100644 --- a/authbridge/authlib/plugins/toolprune/pricing.go +++ b/authbridge/authlib/plugins/toolprune/pricing.go @@ -49,3 +49,14 @@ var defaultPricing = map[string]modelRates{ CacheReadCostPerToken: 0.000000076, }, } + +func (s rateSource) String() string { + switch s { + case rateConfigured: + return "configured" + case rateDefault: + return "default" + default: + return "none" + } +} diff --git a/authbridge/cmd/abctl/tui/events_pane.go b/authbridge/cmd/abctl/tui/events_pane.go index 9e24b23e1..95b4560dd 100644 --- a/authbridge/cmd/abctl/tui/events_pane.go +++ b/authbridge/cmd/abctl/tui/events_pane.go @@ -27,7 +27,9 @@ func newEventsTable() table.Model { {Title: "METHOD", Width: 22}, {Title: "STATUS", Width: 7}, {Title: "DURATION", Width: 10}, - {Title: "TOKENS", Width: 8}, + // Wide enough for "33,650 −10.6k $0.1289": the request total, what + // tool-prune removed from it, and what that was worth. + {Title: "TOKENS / SAVED", Width: 24}, {Title: "HOST", Width: 20}, }), table.WithFocused(true), @@ -92,12 +94,12 @@ func (m *model) rebuildEventsTable() { // Pair request rows with their response rows. ids drives the # column: one // integer repeated across a request/response exchange, which is how an // exchange is read off the timeline. - ids, _ := computeEventPairs(eventRows) + ids, partner := computeEventPairs(eventRows) rows := make([]table.Row, 0, len(eventRows)) m.visibleRows = m.visibleRows[:0] m.hiddenInactive = 0 - for _, er := range eventRows { + for i, er := range eventRows { ev := er.event if m.filter != "" && !matchEventRow(er, m.filter) { continue @@ -137,7 +139,7 @@ func (m *model) rebuildEventsTable() { eventMethod(*ev), statusCell(*ev), durationCell(*ev), - tokensCell(*ev), + m.tokensCellWithSaving(eventRows, partner, i, ev), truncStr(ev.Host, 20), }) m.visibleRows = append(m.visibleRows, er) @@ -744,3 +746,35 @@ func truncateScopes(scopes []string, n int) string { } return strings.Join(scopes[:n], ", ") + fmt.Sprintf(" +%d more", len(scopes)-n) } + +// tokensCellWithSaving renders the TOKENS cell. For a response row that pairs +// with a request tool-prune rewrote, it appends the tokens removed and their +// cost — the saving belongs on the row it happened to, not only in an aggregate +// pane where cache-miss and cache-hit turns average into a number that describes +// neither. +// +// Returns the plain total when the row is not such a response, so every other +// event type renders exactly as before. +func (m *model) tokensCellWithSaving(rows []eventRow, partner map[int]int, i int, ev *pipeline.SessionEvent) string { + base := tokensCell(*ev) + if base == "" { + return "" + } + j, ok := partner[i] + if !ok || j < 0 || j >= len(rows) { + return base + } + req := rows[j].event + if req == nil || req.Phase != pipeline.SessionRequest { + return base + } + ps, ok := decodePruneSaving(req) + if !ok { + return base + } + tokens, usd, ok := savedTokensAndCost(ps, ev.Inference) + if !ok { + return base + } + return formatSavedCell(ev.Inference.TotalTokens, tokens, usd, ps.RateSource) +} diff --git a/authbridge/cmd/abctl/tui/prune_saving.go b/authbridge/cmd/abctl/tui/prune_saving.go new file mode 100644 index 000000000..df7388989 --- /dev/null +++ b/authbridge/cmd/abctl/tui/prune_saving.go @@ -0,0 +1,118 @@ +package tui + +import ( + "encoding/json" + "fmt" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" +) + +// pruneSaving is the tool-prune per-request event as published on the request +// event under "tool-prune". Mirrors the plugin's shape; a decode test guards the +// tags. +type pruneSaving struct { + BytesRemoved int `json:"bytesRemoved"` + BodyBytesAfter int `json:"bodyBytesAfter"` + RateInput float64 `json:"rateInput"` + RateCacheWrite float64 `json:"rateCacheWrite"` + RateCacheRead float64 `json:"rateCacheRead"` + RateSource string `json:"rateSource"` +} + +// decodePruneSaving pulls the tool-prune event off a request event, if present. +func decodePruneSaving(e *pipeline.SessionEvent) (pruneSaving, bool) { + if e == nil || len(e.Plugins) == 0 { + return pruneSaving{}, false + } + raw, ok := e.Plugins["tool-prune"] + if !ok { + return pruneSaving{}, false + } + var ps pruneSaving + if err := json.Unmarshal(raw, &ps); err != nil || ps.BytesRemoved <= 0 || ps.BodyBytesAfter <= 0 { + return pruneSaving{}, false + } + return ps, true +} + +// savedTokensAndCost converts a request's byte saving into tokens and dollars, +// using the response's own usage. +// +// The two halves live on different events by necessity: the byte saving is known +// when the request is rewritten, and the tier it came out of — and the ratio to +// convert bytes to tokens — only from the response. So this is the last step of +// an arithmetic the plugin starts. +// +// Tier matters more than it looks: providers charge ~1.25x the input rate for a +// cache write and ~0.1x for a cache read, so the same saved bytes are worth over +// 12x more on a cache miss than a hit. Picking the tier the request actually +// used is the difference between a figure and a guess. +func savedTokensAndCost(ps pruneSaving, resp *pipeline.InferenceExtension) (tokens, usd float64, ok bool) { + if resp == nil { + return 0, 0, false + } + prompt := resp.InputTokens + resp.CacheReadTokens + resp.CacheWriteTokens + if prompt == 0 { + prompt = resp.PromptTokens // provider reported only an aggregate + } + if prompt <= 0 { + return 0, 0, false + } + // The plugin calibrates on the request it just sent: prompt tokens over the + // post-prune body size, both measured on the same request so the two sides + // agree. + tokens = float64(ps.BytesRemoved) * float64(prompt) / float64(ps.BodyBytesAfter) + + rate := ps.RateInput + switch { + case resp.CacheWriteTokens > resp.CacheReadTokens && resp.CacheWriteTokens > 0: + rate = ps.RateCacheWrite + case resp.CacheReadTokens > 0: + rate = ps.RateCacheRead + } + return tokens, tokens * rate, true +} + +// formatSavedCell renders the TOKENS cell for a response row: the request's own +// total, then what tool-prune removed from it, then what that was worth. +// +// Shown per request rather than only as a pane aggregate because an average +// hides the thing an operator wants to see — cache-miss turns save an order of +// magnitude more than cache-hit turns, and the aggregate flattens that. +func formatSavedCell(total int, tokens, usd float64, rateSource string) string { + cell := formatCount(total) + if tokens <= 0 { + return cell + } + cell += fmt.Sprintf(" −%s", formatCompact(tokens)) + if usd > 0 && rateSource != "none" { + cell += fmt.Sprintf(" $%s", formatUSD(usd)) + } + return cell +} + +// formatCompact renders a token count tersely enough for a table cell: 10577 +// becomes "10.6k". Exact below 1000, where the extra digits still fit. +func formatCompact(v float64) string { + switch { + case v >= 1_000_000: + return fmt.Sprintf("%.1fM", v/1_000_000) + case v >= 1_000: + return fmt.Sprintf("%.1fk", v/1_000) + default: + return fmt.Sprintf("%.0f", v) + } +} + +// formatUSD keeps small amounts legible: a per-request saving is often fractions +// of a cent, where %.2f would round every row to "0.00". +func formatUSD(v float64) string { + switch { + case v >= 1: + return fmt.Sprintf("%.2f", v) + case v >= 0.01: + return fmt.Sprintf("%.3f", v) + default: + return fmt.Sprintf("%.4f", v) + } +} diff --git a/authbridge/cmd/abctl/tui/prune_saving_test.go b/authbridge/cmd/abctl/tui/prune_saving_test.go new file mode 100644 index 000000000..143b5f12e --- /dev/null +++ b/authbridge/cmd/abctl/tui/prune_saving_test.go @@ -0,0 +1,124 @@ +package tui + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" +) + +// wire is the exact JSON the plugin publishes under "tool-prune". +const wire = `{"toolsRemoved":["NotebookEdit","WebSearch"],"bytesRemoved":28568, + "bodyBytesAfter":90635,"model":"claude-opus-5","rateInput":3.8e-06, + "rateCacheWrite":4.75e-06,"rateCacheRead":3.8e-07,"rateSource":"default"}` + +func reqEvent(t *testing.T, raw string) *pipeline.SessionEvent { + t.Helper() + return &pipeline.SessionEvent{ + Phase: pipeline.SessionRequest, + Plugins: map[string]json.RawMessage{"tool-prune": json.RawMessage(raw)}, + } +} + +// TestDecodePruneSaving guards the tags against drift with the plugin's struct. +// A silent decode failure would show a plain token total and look like the plugin +// having saved nothing. +func TestDecodePruneSaving(t *testing.T) { + ps, ok := decodePruneSaving(reqEvent(t, wire)) + if !ok { + t.Fatal("failed to decode the published event") + } + if ps.BytesRemoved != 28568 || ps.BodyBytesAfter != 90635 { + t.Errorf("byte fields = %d/%d", ps.BytesRemoved, ps.BodyBytesAfter) + } + if ps.RateCacheWrite != 4.75e-06 || ps.RateCacheRead != 3.8e-07 { + t.Errorf("rates did not decode: %+v", ps) + } + if ps.RateSource != "default" { + t.Errorf("RateSource = %q", ps.RateSource) + } + // Absent, malformed, and zero-valued all decline rather than render zeros. + for _, bad := range []*pipeline.SessionEvent{ + nil, + {Phase: pipeline.SessionRequest}, + reqEvent(t, `{"bytesRemoved":0,"bodyBytesAfter":100}`), + reqEvent(t, `{"bytesRemoved":5,"bodyBytesAfter":0}`), + reqEvent(t, `not json`), + } { + if _, ok := decodePruneSaving(bad); ok { + t.Errorf("should not decode: %+v", bad) + } + } +} + +// TestSavedTokensAndCost_TierDecidesTheValue is the point of doing this per +// request. The same saved bytes are worth over 12x more on a cache miss than on +// a cache hit, because providers charge ~1.25x input for a write and ~0.1x for a +// read. An aggregate that averages the two describes neither turn. +func TestSavedTokensAndCost_TierDecidesTheValue(t *testing.T) { + ps, _ := decodePruneSaving(reqEvent(t, wire)) + + miss := &pipeline.InferenceExtension{InputTokens: 8881, CacheWriteTokens: 24701} + hit := &pipeline.InferenceExtension{InputTokens: 26, CacheReadTokens: 24701, CacheWriteTokens: 8907} + + tMiss, usdMiss, ok := savedTokensAndCost(ps, miss) + if !ok || tMiss <= 0 || usdMiss <= 0 { + t.Fatalf("miss: tokens=%v usd=%v ok=%v", tMiss, usdMiss, ok) + } + _, usdHit, ok := savedTokensAndCost(ps, hit) + if !ok || usdHit <= 0 { + t.Fatalf("hit: usd=%v ok=%v", usdHit, ok) + } + if r := usdMiss / usdHit; r < 11 || r > 14 { + t.Errorf("miss/hit cost ratio = %.2f, want ~12.5 — the tier must pick the rate", r) + } + + // A provider that reports only an aggregate still works. + agg := &pipeline.InferenceExtension{PromptTokens: 33582} + if _, _, ok := savedTokensAndCost(ps, agg); !ok { + t.Error("should fall back to PromptTokens when the split is absent") + } + // No usage at all declines rather than dividing by zero. + if _, _, ok := savedTokensAndCost(ps, &pipeline.InferenceExtension{}); ok { + t.Error("no usage should not produce a figure") + } + if _, _, ok := savedTokensAndCost(ps, nil); ok { + t.Error("nil usage should not produce a figure") + } +} + +func TestFormatSavedCell(t *testing.T) { + got := formatSavedCell(33650, 10577.5, 0.05024, "default") + for _, want := range []string{"33,650", "−10.6k", "$0.050"} { + if !strings.Contains(got, want) { + t.Errorf("cell %q missing %q", got, want) + } + } + // No saving: the plain total, so unrelated rows are untouched. + if got := formatSavedCell(33650, 0, 0, "default"); got != "33,650" { + t.Errorf("no-saving cell = %q, want the bare total", got) + } + // Unpriced model: tokens shown, no dollar figure invented. + got = formatSavedCell(33650, 10577.5, 0, "none") + if strings.Contains(got, "$") { + t.Errorf("cell %q shows a price for an unpriced model", got) + } + if !strings.Contains(got, "−10.6k") { + t.Errorf("cell %q should still show the token saving", got) + } +} + +func TestFormatCompactAndUSD(t *testing.T) { + for in, want := range map[float64]string{0: "0", 950: "950", 10577.5: "10.6k", 2_400_000: "2.4M"} { + if got := formatCompact(in); got != want { + t.Errorf("formatCompact(%v) = %q, want %q", in, got, want) + } + } + // Sub-cent savings must not all round to 0.00. + for in, want := range map[float64]string{1.5: "1.50", 0.05: "0.050", 0.0004: "0.0004"} { + if got := formatUSD(in); got != want { + t.Errorf("formatUSD(%v) = %q, want %q", in, got, want) + } + } +} diff --git a/authbridge/docs/tool-prune-plugin.md b/authbridge/docs/tool-prune-plugin.md index 174daf80f..db1c7649b 100644 --- a/authbridge/docs/tool-prune-plugin.md +++ b/authbridge/docs/tool-prune-plugin.md @@ -91,6 +91,29 @@ Metrics: In observe mode `requests projected` replaces `requests pruned`, so a projection is never mistaken for a realised saving. +### Per request, in the events timeline + +The events pane's `TOKENS / SAVED` column shows each request's own saving beside +its token total: + +``` +# TIME PHASE PLUGIN TOKENS / SAVED CODE +7 16:42:53.56 req tool-prune +7 16:42:57.27 resp inference-parser 33,604 −24.7k $0.117 200 +8 16:43:11.02 resp inference-parser 33,656 −24.7k $0.009 200 +``` + +Those two rows removed the same bytes and differ ~12x in value: the first was a +cache miss (the manifest was written to cache, ~1.25x the input rate), the second +a hit (read from it, ~0.1x). An aggregate averages the two into a number that +describes neither turn, which is why the saving is shown per row. + +The plugin publishes the byte saving and the applicable rates on the request +event; the tier and the bytes-to-tokens ratio come from the paired response, so +`abctl` finishes the arithmetic. Pairing is exact, on the proxy-stamped request +id. A model with no rate shows the token saving without a dollar figure rather +than one computed at another model's rate. + ### Why tokens are reported per tier and never summed Byte counts are exact. Tokens are an estimate, and — more importantly — they From 1fc8cf6ee163208a18a9a38e1f325aacc63092d6 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 18:19:13 -0400 Subject: [PATCH 22/28] fix(abctl): Attribute the prune saving to the request, not the response MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The saving was rendered on the response row, beside the billed token total. That reads as though the response had been reduced. It had not — tool-prune rewrites the outbound request, and the plugin's own `modify` invocation is already on the request row. Split across the rows the two figures belong to: the request row shows what was removed and what it was worth, the response row shows the token count the provider billed. They share a # so they are still read together. The response is what makes the request-side figure computable — it supplies the prompt token total behind the bytes-to-tokens ratio and the tier that picks the rate — so a request row looks forward to its paired response. That is a rendering detail, not a reason to attribute the saving there. A request whose response has not arrived yet renders an empty cell rather than a partial figure: without the tier there is no rate, and without the ratio no token count. Signed-off-by: Hai Huang --- authbridge/cmd/abctl/tui/events_pane.go | 49 +++++++++++-------- authbridge/cmd/abctl/tui/prune_saving.go | 16 +++--- authbridge/cmd/abctl/tui/prune_saving_test.go | 21 +++++--- authbridge/docs/tool-prune-plugin.md | 33 +++++++------ 4 files changed, 68 insertions(+), 51 deletions(-) diff --git a/authbridge/cmd/abctl/tui/events_pane.go b/authbridge/cmd/abctl/tui/events_pane.go index 95b4560dd..cd4de5b57 100644 --- a/authbridge/cmd/abctl/tui/events_pane.go +++ b/authbridge/cmd/abctl/tui/events_pane.go @@ -747,34 +747,43 @@ func truncateScopes(scopes []string, n int) string { return strings.Join(scopes[:n], ", ") + fmt.Sprintf(" +%d more", len(scopes)-n) } -// tokensCellWithSaving renders the TOKENS cell. For a response row that pairs -// with a request tool-prune rewrote, it appends the tokens removed and their -// cost — the saving belongs on the row it happened to, not only in an aggregate -// pane where cache-miss and cache-hit turns average into a number that describes -// neither. +// tokensCellWithSaving renders the TOKENS / SAVED cell, splitting the two halves +// across the rows they actually belong to: // -// Returns the plain total when the row is not such a response, so every other -// event type renders exactly as before. +// - a REQUEST row that tool-prune rewrote shows what was removed from it, +// which is where the plugin's own `modify` invocation already sits; +// - a RESPONSE row shows the token total the provider billed. +// +// The saving deliberately does NOT go on the response row. Nothing about the +// response was reduced, and showing it there reads as though it were — the +// pruning happened on the way out. The two rows share a # so they are read +// together anyway. +// +// The response is still what makes the request-side figure computable: it +// supplies the prompt token total behind the bytes-to-tokens ratio and the tier +// that sets the rate. So a request row looks forward to its paired response. func (m *model) tokensCellWithSaving(rows []eventRow, partner map[int]int, i int, ev *pipeline.SessionEvent) string { - base := tokensCell(*ev) - if base == "" { + if ev.Phase == pipeline.SessionResponse { + return tokensCell(*ev) + } + if ev.Phase != pipeline.SessionRequest { + return "" + } + ps, ok := decodePruneSaving(ev) + if !ok { return "" } j, ok := partner[i] if !ok || j < 0 || j >= len(rows) { - return base - } - req := rows[j].event - if req == nil || req.Phase != pipeline.SessionRequest { - return base + return "" // no response yet: the ratio and tier are not known } - ps, ok := decodePruneSaving(req) - if !ok { - return base + resp := rows[j].event + if resp == nil || resp.Phase != pipeline.SessionResponse { + return "" } - tokens, usd, ok := savedTokensAndCost(ps, ev.Inference) + tokens, usd, ok := savedTokensAndCost(ps, resp.Inference) if !ok { - return base + return "" } - return formatSavedCell(ev.Inference.TotalTokens, tokens, usd, ps.RateSource) + return formatSavedOnly(tokens, usd, ps.RateSource) } diff --git a/authbridge/cmd/abctl/tui/prune_saving.go b/authbridge/cmd/abctl/tui/prune_saving.go index df7388989..17bec22be 100644 --- a/authbridge/cmd/abctl/tui/prune_saving.go +++ b/authbridge/cmd/abctl/tui/prune_saving.go @@ -73,18 +73,14 @@ func savedTokensAndCost(ps pruneSaving, resp *pipeline.InferenceExtension) (toke return tokens, tokens * rate, true } -// formatSavedCell renders the TOKENS cell for a response row: the request's own -// total, then what tool-prune removed from it, then what that was worth. -// -// Shown per request rather than only as a pane aggregate because an average -// hides the thing an operator wants to see — cache-miss turns save an order of -// magnitude more than cache-hit turns, and the aggregate flattens that. -func formatSavedCell(total int, tokens, usd float64, rateSource string) string { - cell := formatCount(total) +// formatSavedOnly renders a request row's saving: what was removed and what it +// was worth. No total, because a request has no billed token count — that +// belongs to the response, on its own row. +func formatSavedOnly(tokens, usd float64, rateSource string) string { if tokens <= 0 { - return cell + return "" } - cell += fmt.Sprintf(" −%s", formatCompact(tokens)) + cell := "−" + formatCompact(tokens) if usd > 0 && rateSource != "none" { cell += fmt.Sprintf(" $%s", formatUSD(usd)) } diff --git a/authbridge/cmd/abctl/tui/prune_saving_test.go b/authbridge/cmd/abctl/tui/prune_saving_test.go index 143b5f12e..ae7e2aa18 100644 --- a/authbridge/cmd/abctl/tui/prune_saving_test.go +++ b/authbridge/cmd/abctl/tui/prune_saving_test.go @@ -88,19 +88,26 @@ func TestSavedTokensAndCost_TierDecidesTheValue(t *testing.T) { } } -func TestFormatSavedCell(t *testing.T) { - got := formatSavedCell(33650, 10577.5, 0.05024, "default") - for _, want := range []string{"33,650", "−10.6k", "$0.050"} { +// TestFormatSavedOnly: a request row carries the saving, not a total — the +// billed token count belongs to the response, on its own row. Showing a saving +// beside a response total read as though the response had shrunk, which it had +// not. +func TestFormatSavedOnly(t *testing.T) { + got := formatSavedOnly(10577.5, 0.05024, "default") + for _, want := range []string{"−10.6k", "$0.050"} { if !strings.Contains(got, want) { t.Errorf("cell %q missing %q", got, want) } } - // No saving: the plain total, so unrelated rows are untouched. - if got := formatSavedCell(33650, 0, 0, "default"); got != "33,650" { - t.Errorf("no-saving cell = %q, want the bare total", got) + if strings.Contains(got, ",") { + t.Errorf("cell %q should carry no billed total", got) + } + // Nothing saved: an empty cell, so unrelated request rows stay blank. + if got := formatSavedOnly(0, 0, "default"); got != "" { + t.Errorf("no-saving cell = %q, want empty", got) } // Unpriced model: tokens shown, no dollar figure invented. - got = formatSavedCell(33650, 10577.5, 0, "none") + got = formatSavedOnly(10577.5, 0, "none") if strings.Contains(got, "$") { t.Errorf("cell %q shows a price for an unpriced model", got) } diff --git a/authbridge/docs/tool-prune-plugin.md b/authbridge/docs/tool-prune-plugin.md index db1c7649b..f6abbdbcf 100644 --- a/authbridge/docs/tool-prune-plugin.md +++ b/authbridge/docs/tool-prune-plugin.md @@ -93,26 +93,31 @@ projection is never mistaken for a realised saving. ### Per request, in the events timeline -The events pane's `TOKENS / SAVED` column shows each request's own saving beside -its token total: +The events pane's `TOKENS / SAVED` column splits the two halves across the rows +they belong to — the saving on the request that was rewritten, the billed total on +the response: ``` -# TIME PHASE PLUGIN TOKENS / SAVED CODE -7 16:42:53.56 req tool-prune -7 16:42:57.27 resp inference-parser 33,604 −24.7k $0.117 200 -8 16:43:11.02 resp inference-parser 33,656 −24.7k $0.009 200 +# PHASE ACTION PLUGIN TOKENS / SAVED CODE +12 req modify tool-prune −24.7k $0.117 +12 resp observe inference-parser 34,702 200 ``` -Those two rows removed the same bytes and differ ~12x in value: the first was a -cache miss (the manifest was written to cache, ~1.25x the input rate), the second -a hit (read from it, ~0.1x). An aggregate averages the two into a number that -describes neither turn, which is why the saving is shown per row. +The saving is not shown on the response row: nothing about the response was +reduced, and putting it there reads as though it had been. The two rows share a +`#` so they are read together anyway. + +Two turns that removed the same bytes can still differ ~12x in value — a cache +miss writes the manifest to cache (~1.25x the input rate), a hit reads it (~0.1x). +An aggregate averages the two into a number that describes neither, which is why +this is per row. The plugin publishes the byte saving and the applicable rates on the request -event; the tier and the bytes-to-tokens ratio come from the paired response, so -`abctl` finishes the arithmetic. Pairing is exact, on the proxy-stamped request -id. A model with no rate shows the token saving without a dollar figure rather -than one computed at another model's rate. +event; the paired response supplies the prompt token total behind the +bytes-to-tokens ratio and the tier that picks the rate, so `abctl` finishes the +arithmetic. Pairing is exact, on the proxy-stamped request id. A model with no +rate shows the token saving with no dollar figure rather than one priced at +another model's rate. ### Why tokens are reported per tier and never summed From d85b4e335e02d89498863a00772d720d84dfa487 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 18:49:21 -0400 Subject: [PATCH 23/28] fix: Address code review on the tool-prune series MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Correctness and privacy: - snapshot.go: gjson's String() on an object or array returns that node's RAW JSON, so a structured error.type put response body content — anything the provider quoted from the request — into the unauthenticated session store, defeating the reason error.message is excluded. Now accepts only a JSON string or number. A test plants a credential in a structured value. - toolprune: modelRates.rateFor reports whether a usable rate exists. set() ORs three fields, so a model configured with only a cache-read rate resolved as priced and then charged a cache-write request zero — vanishing from the total with no `requests unpriced` row. - toolprune: the built-in per-model table now precedes the flat fallback. The flat fields are documented as covering models "absent from pricing", and a model in the table is not absent; one flat rate shadowing every per-model default reintroduced flat-rate mispricing, silently, and claimed to be operator-configured. - toolprune: forcedToolChoice replaces forcedToolName. An object tool_choice naming nothing recognisable (Bedrock Converse nests it as tool.name) now declines to prune rather than reading it as "nothing forced" and risking removal of the one required tool. - abctl: a response carrying a RequestID that fails to pair exactly — a retry, or a stream recorded twice — no longer falls through to the adjacency heuristic, where it could claim an unrelated earlier request. The same guard gates pricing, since a mismatched response supplies the wrong cache tier and the tiers are 12.5x apart. - toolscan: PatchConfig writes via temp file + Sync + rename. os.WriteFile truncates in place, so a crash left a truncated config with no recovery copy, and the proxy's fsnotify reloader could observe the partial file. - demo.go: writeDemoConfig keeps an existing demo.yaml. It runs before any port binds, so an unconditional write meant a --demo start that then failed on a port clash destroyed the operator's edits — including a prune list written by `abctl tools scan --write`, which the config's own comment recommends. sparc declared WritesRequestBody but calls pctx.SetBody nowhere; the flag was stale from the undirected capability and occupied the single request-mutator slot, so [sparc, tool-prune] could not build. Dropped — which is the payoff this series was arguing for, now pinned by a test. Tests: real byte-exactness for the prune (reconstructing expected output from the original bytes, covering first/middle/last element and validating with encoding/json — the old test asserted only fragments and a shorter length, and never removed a first or last element); tool_choice string forms; OpenAI-dialect all-removed; and a reflection-driven clone check that fails if a future slice/map capability is aliased. Also: bytes.Contains on the scan hot path; names_unresolved distinguished from no_configured_tool_present; an in-flight guard so the 2s refresh tick cannot stack fetches against a 10s timeout; the dead crypto/rand fallback removed (cannot fail as of Go 1.24); and docs corrected — WritesResponseBody added to the capability snippet, the duplicated rate-derivation section removed, the "ships with on_error: observe" claim replaced with the empty remove list that is the actual guard, counters noted as resetting on hot-reload too, the BodyAccess changelog line marked as since-removed, and the README's 20-25% figure attributed to the traffic it was measured on. Signed-off-by: Hai Huang --- README.md | 5 +- .../authlib/pipeline/bodydirection_test.go | 23 +++ authbridge/authlib/pipeline/errorkind_test.go | 38 ++++- authbridge/authlib/pipeline/requestid.go | 13 +- authbridge/authlib/pipeline/snapshot.go | 22 ++- .../plugins/registry_capsclone_test.go | 37 ++++- authbridge/authlib/plugins/sparc/plugin.go | 12 +- .../authlib/plugins/sparc/plugin_test.go | 16 ++- .../authlib/plugins/toolprune/metrics.go | 2 +- .../authlib/plugins/toolprune/plugin.go | 133 +++++++++++++----- .../authlib/plugins/toolprune/plugin_test.go | 128 +++++++++++++++++ authbridge/cmd/abctl/toolscan/patch.go | 45 +++++- authbridge/cmd/abctl/toolscan/scan.go | 6 +- authbridge/cmd/abctl/tui/app.go | 18 ++- authbridge/cmd/abctl/tui/events_pane.go | 14 ++ authbridge/cmd/abctl/tui/prune_saving_test.go | 53 +++++++ authbridge/cmd/authbridge-proxy/demo.go | 24 +++- authbridge/cmd/authbridge-proxy/demo_test.go | 33 +++++ authbridge/docs/framework-architecture.md | 2 +- authbridge/docs/plugin-reference.md | 4 +- authbridge/docs/tool-prune-plugin.md | 12 +- authbridge/install-demo.sh | 5 +- docs/proposals/tool-prune.md | 14 +- 23 files changed, 568 insertions(+), 91 deletions(-) diff --git a/README.md b/README.md index e1f1e3b67..f3eb76d17 100644 --- a/README.md +++ b/README.md @@ -40,8 +40,9 @@ Watch an AI agent's traffic — its model, tool, and agent-to-agent calls — de ## Cut Claude Code token cost on your laptop Already using Claude Code? Cortex can strip the tool definitions your agent never -calls out of every request — typically 20–25% of the prompt you pay for on each -turn. Four steps, about two minutes: +calls out of every request. On the traffic this was measured against that is +20–25% of the prompt billed per turn; your share depends on how many of the +tools you actually use. Four steps, about two minutes: **[Cut Claude Code token cost](./authbridge/docs/laptop-token-savings.md)**. ## Running on Kubernetes diff --git a/authbridge/authlib/pipeline/bodydirection_test.go b/authbridge/authlib/pipeline/bodydirection_test.go index 28bb8f016..8b4874508 100644 --- a/authbridge/authlib/pipeline/bodydirection_test.go +++ b/authbridge/authlib/pipeline/bodydirection_test.go @@ -163,3 +163,26 @@ func TestValidateCapabilities_Directional(t *testing.T) { }) } } + +// TestValidateCapabilities_ResponseAndRequestMutatorsCoexist is the payoff the +// directional split was arguing for. Before it, SPARC's undirected flag occupied +// the only mutator slot, so a request-only mutator could not share a chain with +// it even though the two write different bodies. Now the real in-tree shape — +// parser, response mutator, request mutator — builds. +func TestValidateCapabilities_ResponseAndRequestMutatorsCoexist(t *testing.T) { + err := validateCapabilities([]Plugin{ + &stubPlugin{name: "inference-parser", caps: PluginCapabilities{ReadsBody: true}}, + &stubPlugin{name: "sparc", caps: PluginCapabilities{WritesResponseBody: true}}, + &stubPlugin{name: "tool-prune", caps: PluginCapabilities{WritesRequestBody: true}}, + }) + if err != nil { + t.Errorf("[parser, sparc, tool-prune] should build: %v", err) + } + // Two mutators on the SAME side are still rejected. + if err := validateCapabilities([]Plugin{ + &stubPlugin{name: "sparc", caps: PluginCapabilities{WritesResponseBody: true}}, + &stubPlugin{name: "cpex", caps: PluginCapabilities{WritesResponseBody: true}}, + }); err == nil { + t.Error("two response mutators must still be rejected") + } +} diff --git a/authbridge/authlib/pipeline/errorkind_test.go b/authbridge/authlib/pipeline/errorkind_test.go index 8e95994a9..f7031c607 100644 --- a/authbridge/authlib/pipeline/errorkind_test.go +++ b/authbridge/authlib/pipeline/errorkind_test.go @@ -1,6 +1,9 @@ package pipeline -import "testing" +import ( + "strings" + "testing" +) // TestUpstreamErrorKind: a bare "backend_error / 400" gives an operator nothing // to act on. The provider's own classification does — and it must be the @@ -75,3 +78,36 @@ func TestDeriveError_PopulatesKindFrom4xxBody(t *testing.T) { t.Errorf("bare 5xx = %+v, want backend_error/503 with empty message", bare) } } + +// TestUpstreamErrorKind_RefusesStructuredValues is the privacy regression for a +// leak the earlier test could not see: gjson's String() on an object or array +// returns that node's RAW JSON. A provider (or a proxy in between) returning a +// structured error.type therefore put response body content — including anything +// quoted from the request — straight into the unauthenticated session store, +// defeating the whole reason error.message is excluded. +func TestUpstreamErrorKind_RefusesStructuredValues(t *testing.T) { + secret := "sk-live-DEADBEEF" + for _, body := range []string{ + `{"error":{"type":{"secret":"` + secret + `","nested":true}}}`, + `{"error":{"type":["` + secret + `"]}}`, + `{"error":{"code":{"inner":"` + secret + `"}}}`, + `{"error":{"type":true}}`, + `{"error":{"type":null}}`, + } { + got := upstreamErrorKind([]byte(body)) + if got != "" { + t.Errorf("structured value leaked %q from %s", got, body) + } + if strings.Contains(got, secret) { + t.Fatalf("CREDENTIAL LEAK: %q", got) + } + } + // A numeric code carries no payload and stays useful. + if got := upstreamErrorKind([]byte(`{"error":{"code":429}}`)); got != "429" { + t.Errorf("numeric code = %q, want 429", got) + } + // The normal string path is unaffected. + if got := upstreamErrorKind([]byte(`{"error":{"type":"rate_limit_error"}}`)); got != "rate_limit_error" { + t.Errorf("string type = %q", got) + } +} diff --git a/authbridge/authlib/pipeline/requestid.go b/authbridge/authlib/pipeline/requestid.go index 395dbd097..01f6fadf6 100644 --- a/authbridge/authlib/pipeline/requestid.go +++ b/authbridge/authlib/pipeline/requestid.go @@ -3,15 +3,8 @@ package pipeline import ( "crypto/rand" "encoding/hex" - "strconv" - "sync/atomic" ) -// requestIDCounter is the fallback when crypto/rand is unavailable, so an id is -// always produced rather than an empty string that would silently disable -// pairing. -var requestIDCounter atomic.Uint64 - // newRequestID returns a short, unique-per-process request identifier. // // Not a UUID on purpose: it exists to pair a request event with its response @@ -20,9 +13,9 @@ var requestIDCounter atomic.Uint64 // cryptographically meaningful. func newRequestID() string { var b [6]byte - if _, err := rand.Read(b[:]); err != nil { - return "r" + strconv.FormatUint(requestIDCounter.Add(1), 36) - } + // crypto/rand.Read never returns an error as of Go 1.24 — it panics on an + // unusable system source instead — so there is no failure branch to write. + _, _ = rand.Read(b[:]) return hex.EncodeToString(b[:]) } diff --git a/authbridge/authlib/pipeline/snapshot.go b/authbridge/authlib/pipeline/snapshot.go index 693421422..d144b183f 100644 --- a/authbridge/authlib/pipeline/snapshot.go +++ b/authbridge/authlib/pipeline/snapshot.go @@ -165,9 +165,15 @@ func upstreamErrorKind(body []byte) string { if !gjson.ValidBytes(body) { return "" } - t := gjson.GetBytes(body, "error.type").String() + // Only accept a JSON string. gjson's String() on an object or array returns + // that node's RAW JSON, so {"error":{"type":{...}}} would put response body + // content — quoted request data, credentials — straight into the + // unauthenticated session store, defeating the whole point of excluding + // error.message. A numeric code is accepted because a number carries no + // payload; anything structured is refused. + t := stringOrNumber(gjson.GetBytes(body, "error.type")) if t == "" { - t = gjson.GetBytes(body, "error.code").String() + t = stringOrNumber(gjson.GetBytes(body, "error.code")) } if t == "" { return "" @@ -177,3 +183,15 @@ func upstreamErrorKind(body []byte) string { } return t } + +// stringOrNumber returns the value only when the node is a JSON string or +// number. Every other type — object, array, true/false, absent — yields "", +// because String() on a container returns its raw JSON and that is body content. +func stringOrNumber(r gjson.Result) string { + switch r.Type { + case gjson.String, gjson.Number: + return r.String() + default: + return "" + } +} diff --git a/authbridge/authlib/plugins/registry_capsclone_test.go b/authbridge/authlib/plugins/registry_capsclone_test.go index e0011577e..05ba61a72 100644 --- a/authbridge/authlib/plugins/registry_capsclone_test.go +++ b/authbridge/authlib/plugins/registry_capsclone_test.go @@ -60,8 +60,41 @@ func TestCloneCatalog_PreservesEveryCapabilityField(t *testing.T) { } } -// TestCloneCatalog_DeepCopiesSlices: the clone must not alias the caller's -// slices, or a mutation through /v1/plugins would reach into the registry. +// TestCloneCatalog_DeepCopiesEveryReferenceField walks PluginCapabilities by +// reflection and asserts that no field of a reference kind is aliased. The +// struct copy in cloneCatalog is correct for today's two slices, but a future +// map or slice capability would be silently shared with the registry — the same +// class of bug the field-by-field copy had, which is why this is driven by the +// struct rather than by a hand-written list. +func TestCloneCatalog_DeepCopiesEveryReferenceField(t *testing.T) { + caps := nonZeroCaps(t) + in := []CatalogEntry{{Name: "probe", Capabilities: caps}} + out := cloneCatalog(in) + + src := reflect.ValueOf(&in[0].Capabilities).Elem() + dst := reflect.ValueOf(&out[0].Capabilities).Elem() + for i := 0; i < src.NumField(); i++ { + name := src.Type().Field(i).Name + switch src.Field(i).Kind() { + case reflect.Slice: + if src.Field(i).Len() == 0 { + t.Fatalf("%s: nonZeroCaps left it empty, so aliasing cannot be detected", name) + } + if src.Field(i).UnsafePointer() == dst.Field(i).UnsafePointer() { + t.Errorf("%s aliases the registry's slice", name) + } + case reflect.Map, reflect.Pointer: + if src.Field(i).UnsafePointer() == dst.Field(i).UnsafePointer() { + t.Errorf("%s is a %s shared with the registry — cloneCatalog needs to copy it", + name, src.Field(i).Kind()) + } + } + } +} + +// TestCloneCatalog_DeepCopiesSlices keeps the concrete mutation check: the clone +// must not alias the caller's slices, or a mutation through /v1/plugins would +// reach into the registry. func TestCloneCatalog_DeepCopiesSlices(t *testing.T) { in := []CatalogEntry{{ Name: "probe", diff --git a/authbridge/authlib/plugins/sparc/plugin.go b/authbridge/authlib/plugins/sparc/plugin.go index 7d0d72618..f99ab0852 100644 --- a/authbridge/authlib/plugins/sparc/plugin.go +++ b/authbridge/authlib/plugins/sparc/plugin.go @@ -209,10 +209,14 @@ func (p *SPARC) Capabilities() pipeline.PluginCapabilities { // conversation + tool specs (both modes); mcp-parser provides the tool // call (mcp mode). RequiresAny is a static "at least one" check; the // per-mode runtime requirements are validated/handled below. - RequiresAny: []string{"inference-parser", "mcp-parser"}, - ReadsBody: true, - WritesRequestBody: true, // MCP result (mcp mode) / completion rewrite (inference mode) - WritesResponseBody: true, // respond.go rewrites the upstream response + RequiresAny: []string{"inference-parser", "mcp-parser"}, + ReadsBody: true, + // Response-only: SPARC rewrites the upstream response (respond.go), and + // calls pctx.SetBody nowhere. Declaring WritesRequestBody was carried over + // from the undirected flag and cost it the single request-mutator slot for + // nothing, so a chain like [sparc, tool-prune] could not build even though + // the two write different bodies. + WritesResponseBody: true, Description: "SPARC pre-tool reflection: blocks ungrounded/hallucinated tool calls.", } } diff --git a/authbridge/authlib/plugins/sparc/plugin_test.go b/authbridge/authlib/plugins/sparc/plugin_test.go index 5457e580e..aa55db535 100644 --- a/authbridge/authlib/plugins/sparc/plugin_test.go +++ b/authbridge/authlib/plugins/sparc/plugin_test.go @@ -335,10 +335,22 @@ func TestInference_MCPModeOnResponseIsNoop(t *testing.T) { } } +// TestCapabilities pins SPARC as a RESPONSE-side mutator. It rewrites the +// upstream response (respond.go) and calls pctx.SetBody nowhere, so declaring +// WritesRequestBody was carried over from the undirected flag and cost it the +// single request-mutator slot for nothing — a chain like [sparc, tool-prune] +// could not build even though the two write different bodies. func TestCapabilities(t *testing.T) { caps := NewSPARC().Capabilities() - if !caps.WritesRequestBody || !caps.ReadsBody { - t.Error("expected ReadsBody+WritesRequestBody") + if !caps.WritesResponseBody { + t.Error("expected WritesResponseBody — SPARC rewrites the response") + } + if caps.WritesRequestBody { + t.Error("must not declare WritesRequestBody: SPARC never calls pctx.SetBody, " + + "and the claim blocks any real request mutator from sharing the chain") + } + if !caps.Normalize().ReadsBody { + t.Error("a write flag must promote ReadsBody") } if len(caps.RequiresAny) == 0 { t.Error("expected RequiresAny parsers") diff --git a/authbridge/authlib/plugins/toolprune/metrics.go b/authbridge/authlib/plugins/toolprune/metrics.go index 67e0c3b55..750d389db 100644 --- a/authbridge/authlib/plugins/toolprune/metrics.go +++ b/authbridge/authlib/plugins/toolprune/metrics.go @@ -119,7 +119,7 @@ func (m *metrics) observeSaving(tokens float64, t tier, usd float64, src rateSou // snapshot renders the counters as operator-facing metrics. Every derived row // carries the sample it was computed from, so a figure can never be read as // more certain than it is. -func (m *metrics) snapshot(cfg *config) []pipeline.Metric { +func (m *metrics) snapshot() []pipeline.Metric { m.mu.Lock() defer m.mu.Unlock() diff --git a/authbridge/authlib/plugins/toolprune/plugin.go b/authbridge/authlib/plugins/toolprune/plugin.go index 270840050..bfe165461 100644 --- a/authbridge/authlib/plugins/toolprune/plugin.go +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -85,18 +85,25 @@ type modelRates struct { CacheReadCostPerToken float64 `json:"cache_read_cost_per_token" description:"USD per cache-read token; defaults to input_cost_per_token."` } -func (r modelRates) rateFor(t tier) float64 { +// rateFor returns the rate for a tier and whether one is actually available. +// +// The bool matters: set() is an OR across three fields, so a model configured +// with only cache_read_cost_per_token used to resolve as "priced" and then +// return 0 for a cache-write request — pricing it at zero while still counting +// toward the priced denominator, so the saving silently vanished with no +// `requests unpriced` row to show it had. +func (r modelRates) rateFor(t tier) (float64, bool) { switch t { case tierCacheWrite: if r.CacheWriteCostPerToken > 0 { - return r.CacheWriteCostPerToken + return r.CacheWriteCostPerToken, true } case tierCacheRead: if r.CacheReadCostPerToken > 0 { - return r.CacheReadCostPerToken + return r.CacheReadCostPerToken, true } } - return r.InputCostPerToken + return r.InputCostPerToken, r.InputCostPerToken > 0 } func (r modelRates) set() bool { @@ -122,6 +129,15 @@ func (c *config) ratesFor(model string) (modelRates, rateSource) { if r, ok := c.pricing[key]; ok && r.set() { return r, rateConfigured } + // The built-in table comes BEFORE the flat fallback. The flat fields are + // documented as covering "models absent from pricing", and a model in the + // built-in table is not absent — letting one flat input rate shadow every + // per-model default would reintroduce exactly the flat-rate mispricing the + // per-model table exists to avoid, and silently, since the figure would then + // claim to be configured. + if r, ok := defaultPricing[key]; ok { + return r, rateDefault + } fallback := modelRates{ InputCostPerToken: c.InputCostPerToken, CacheWriteCostPerToken: c.CacheWriteCostPerToken, @@ -130,26 +146,9 @@ func (c *config) ratesFor(model string) (modelRates, rateSource) { if fallback.set() { return fallback, rateConfigured } - if r, ok := defaultPricing[key]; ok { - return r, rateDefault - } return modelRates{}, rateNone } -// configuredPricing reports whether the operator supplied any rates of their -// own, as opposed to relying on the built-in defaults. -func (c *config) configuredPricing() bool { - if c.InputCostPerToken > 0 || c.CacheWriteCostPerToken > 0 || c.CacheReadCostPerToken > 0 { - return true - } - for _, r := range c.Pricing { - if r.set() { - return true - } - } - return false -} - func (c *config) applyDefaults() { if len(c.Paths) == 0 { c.Paths = append([]string(nil), defaultPaths...) @@ -264,22 +263,40 @@ func toolNameAt(body []byte, i int) string { return gjson.GetBytes(body, fmt.Sprintf("tools.%d.function.name", i)).String() } -// forcedToolName returns the tool a forced tool_choice names, or "" when the -// request does not force one. Anthropic spells it tool_choice.name, OpenAI -// tool_choice.function.name; "auto" / "none" / "any" carry no name. +// forcedToolChoice reports the tool a forced tool_choice names, and whether the +// tool_choice could be interpreted at all. +// +// resolvable is false only when tool_choice is an object from which no name can +// be read. That is the dangerous case: the request forces *some* tool the plugin +// cannot identify, so pruning risks removing it and producing an invalid request. +// Dialects nest this differently — Anthropic tool_choice.name, OpenAI +// tool_choice.function.name, Bedrock Converse tool_choice.tool.name — and an +// unknown shape must not be read as "nothing is forced". // -// This tool can never be removed: a tool_choice naming a tool absent from the -// manifest is an invalid request, so pruning it would turn a cost optimisation -// into a 400. -func forcedToolName(body []byte) string { +// A string form ("auto", "none", "any", "required") forces no *specific* tool, so +// it is resolvable with an empty name: pruning is safe. +func forcedToolChoice(body []byte) (name string, resolvable bool) { tc := gjson.GetBytes(body, "tool_choice") + if !tc.Exists() { + return "", true + } if !tc.IsObject() { - return "" // "auto" / "none" / absent + return "", true // "auto" / "none" / "any" / "required" } - if n := tc.Get("name"); n.Exists() { - return n.String() + for _, path := range []string{"name", "function.name", "tool.name"} { + if n := tc.Get(path); n.Type == gjson.String && n.String() != "" { + return n.String(), true + } + } + // An object naming nothing we recognise. It may still be a plain + // {"type":"auto"}, which is safe — accept only that narrow shape. + if t := tc.Get("type"); t.Type == gjson.String { + switch t.String() { + case "auto", "none", "any", "required": + return "", true + } } - return tc.Get("function.name").String() + return "", false } // OnRequest prunes the manifest. Every failure path returns Continue with the @@ -346,14 +363,29 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action // Resolve indices from the raw bytes rather than from the parsed manifest: // inference-parser drops unnamed tools, so manifest position does not // reliably map back to array position. - forced := forcedToolName(body) + forced, resolvable := forcedToolChoice(body) + if !resolvable { + // tool_choice is an object but names no tool we recognise — e.g. a + // dialect that nests it differently (Bedrock Converse's + // {"tool":{"name":X}}). Treating that as "nothing is forced" risks + // pruning the one tool the request requires, so decline instead. A + // missed saving is the cheap direction of failure. + pctx.Record(pipeline.Invocation{ + Action: pipeline.ActionSkip, + Reason: "tool_choice_unresolved", + Path: pctx.Path, + }) + return action + } var victims []int + var anyNameResolved bool names := make([]string, 0, len(raw)) for i := range raw { name := toolNameAt(body, i) if name == "" { continue } + anyNameResolved = true if name == forced { // Removing the tool tool_choice forces would make the request // invalid. Keep it and prune the rest. @@ -366,7 +398,15 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action } } if len(victims) == 0 { - pctx.Record(pipeline.Invocation{Action: pipeline.ActionSkip, Reason: "no_configured_tool_present"}) + // Distinguish "the manifest had none of the configured tools" from "no + // tool name could be read at all" — the latter means an unrecognised + // dialect (Gemini, Bedrock toolSpec nesting), where the plugin is inert + // for a reason an operator would want to know about. + reason := "no_configured_tool_present" + if len(names) == 0 && !anyNameResolved { + reason = "names_unresolved" + } + pctx.Record(pipeline.Invocation{Action: pipeline.ActionSkip, Reason: reason, Path: pctx.Path}) return action } @@ -425,14 +465,23 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action // (matching on RequestID) to get the prompt token total and which tier the // saving came out of, and finishes the arithmetic. rates, src := p.cfg.ratesFor(inferenceModel(pctx)) + rateInput, _ := rates.rateFor(tierInput) + rateWrite, okW := rates.rateFor(tierCacheWrite) + rateRead, okR := rates.rateFor(tierCacheRead) + if !okW { + rateWrite = 0 + } + if !okR { + rateRead = 0 + } p.publish(pctx, pruneEvent{ ToolsRemoved: names, BytesRemoved: removedBytes, BodyBytesAfter: len(out), Model: inferenceModel(pctx), - RateInput: rates.InputCostPerToken, - RateCacheWrite: rates.rateFor(tierCacheWrite), - RateCacheRead: rates.rateFor(tierCacheRead), + RateInput: rateInput, + RateCacheWrite: rateWrite, + RateCacheRead: rateRead, RateSource: src.String(), }) // Carry the saving to OnFinish, where the response reveals which token tier @@ -498,7 +547,13 @@ func (p *ToolPrune) OnFinish(_ context.Context, pctx *pipeline.Context) { } t := tierOf(inf) rates, src := p.cfg.ratesFor(inf.Model) - p.m.observeSaving(tokens, t, tokens*rates.rateFor(t), src, inf.Model) + rate, ok := rates.rateFor(t) + if !ok { + // No usable rate for the tier this request actually used. Count it + // unpriced rather than charging zero into the total. + src = rateNone + } + p.m.observeSaving(tokens, t, tokens*rate, src, inf.Model) } // tier names which prompt token tier a request's saving came out of. @@ -551,4 +606,4 @@ func (p *ToolPrune) noteDrift(observed []pipeline.InferenceTool) { } // Metrics implements pipeline.MetricsProvider. -func (p *ToolPrune) Metrics() []pipeline.Metric { return p.m.snapshot(&p.cfg) } +func (p *ToolPrune) Metrics() []pipeline.Metric { return p.m.snapshot() } diff --git a/authbridge/authlib/plugins/toolprune/plugin_test.go b/authbridge/authlib/plugins/toolprune/plugin_test.go index 238893944..f3ea93192 100644 --- a/authbridge/authlib/plugins/toolprune/plugin_test.go +++ b/authbridge/authlib/plugins/toolprune/plugin_test.go @@ -3,6 +3,7 @@ package toolprune import ( "context" "encoding/json" + "fmt" "strings" "sync" "testing" @@ -772,3 +773,130 @@ func TestPricing_ModelMatchIsCaseInsensitive(t *testing.T) { t.Error("model lookup should be case-insensitive") } } + +// TestPrune_ByteExactAgainstJSONReconstruction is the real byte-exactness check. +// The earlier test asserted only that some fragments survived and the body got +// shorter, which passes even if the rewrite reflows the whole document — and it +// removed only a middle element, so the two comma cases that actually differ +// (first and last) were never exercised. +// +// Here the expected output is built by deleting the same elements from the +// ORIGINAL bytes by hand, so any reformatting, key reordering or whitespace +// change fails. Also validates the result with encoding/json, which nothing did. +func TestPrune_ByteExactAgainstJSONReconstruction(t *testing.T) { + const orig = `{"model":"m","tools":[{"name":"A","x":1},{"name":"B","x":2},{"name":"C","x":3}],"max_tokens":8}` + cases := []struct { + remove []string + want string + }{ + {[]string{"A"}, `{"model":"m","tools":[{"name":"B","x":2},{"name":"C","x":3}],"max_tokens":8}`}, + {[]string{"C"}, `{"model":"m","tools":[{"name":"A","x":1},{"name":"B","x":2}],"max_tokens":8}`}, + {[]string{"B"}, `{"model":"m","tools":[{"name":"A","x":1},{"name":"C","x":3}],"max_tokens":8}`}, + {[]string{"A", "C"}, `{"model":"m","tools":[{"name":"B","x":2}],"max_tokens":8}`}, + } + for _, tc := range cases { + t.Run(strings.Join(tc.remove, "+"), func(t *testing.T) { + p := configured(t, tc.remove...) + pctx := inferenceCtx("/v1/messages", orig, "A", "B", "C") + run(t, p, pctx) + if got := string(pctx.Body); got != tc.want { + t.Errorf("byte mismatch\n got: %s\nwant: %s", got, tc.want) + } + var any map[string]any + if err := json.Unmarshal(pctx.Body, &any); err != nil { + t.Errorf("result is not valid JSON: %v", err) + } + }) + } +} + +// TestPrune_ToolChoiceStringForms: "required" / "any" force no specific tool, so +// they must not suppress pruning; an object naming nothing recognisable must. +func TestPrune_ToolChoiceStringForms(t *testing.T) { + body := `{"tools":[{"name":"Read"},{"name":"NotebookEdit"}],"tool_choice":%s}` + for _, tc := range []struct { + choice string + wantPruned bool + }{ + {`"required"`, true}, + {`"any"`, true}, + {`{"type":"required"}`, true}, + {`{"tool":{"name":"NotebookEdit"}}`, false}, // Bedrock-style forced tool: kept + {`{"unknown_shape":true}`, false}, // cannot interpret: decline + } { + t.Run(tc.choice, func(t *testing.T) { + p := configured(t, "NotebookEdit") + pctx := inferenceCtx("/v1/messages", fmt.Sprintf(body, tc.choice), "Read", "NotebookEdit") + run(t, p, pctx) + pruned := !strings.Contains(string(pctx.Body), "NotebookEdit") + if pruned != tc.wantPruned { + t.Errorf("tool_choice %s: pruned=%v want %v — body: %s", tc.choice, pruned, tc.wantPruned, pctx.Body) + } + }) + } +} + +// TestPrune_OpenAIDialectAllRemoved: the all-removed path drops tools and +// tool_choice, and must do so for the OpenAI shape too. +func TestPrune_OpenAIDialectAllRemoved(t *testing.T) { + body := `{"model":"m","tools":[{"type":"function","function":{"name":"A"}},` + + `{"type":"function","function":{"name":"B"}}],"tool_choice":"auto"}` + p := configured(t, "A", "B") + pctx := inferenceCtx("/v1/chat/completions", body, "A", "B") + run(t, p, pctx) + got := string(pctx.Body) + if strings.Contains(got, "tools") || strings.Contains(got, "tool_choice") { + t.Errorf("both keys should be dropped: %s", got) + } + var any map[string]any + if err := json.Unmarshal(pctx.Body, &any); err != nil { + t.Errorf("result is not valid JSON: %v", err) + } +} + +// TestPricing_PartialModelConfigIsUnpriced: set() ORs the three rate fields, so a +// model configured with only a cache-read rate used to resolve as "priced" and +// then return 0 for a cache-write request — charging zero into the total while +// counting toward the priced denominator, so the saving vanished with no +// `requests unpriced` row to show it had. +func TestPricing_PartialModelConfigIsUnpriced(t *testing.T) { + p := configuredJSON(t, `{"remove":["NotebookEdit"], + "pricing":{"some-model":{"cache_read_cost_per_token":1e-06}}}`) + pruneWithModel(t, p, "some-model") // pruneWithModel finishes as a cache WRITE + + gap := findMetric(t, p.Metrics(), "requests unpriced") + if gap.Value != 1 { + t.Errorf("requests unpriced = %v, want 1 — no cache-write rate is configured", gap.Value) + } + for _, m := range p.Metrics() { + if m.Name == "$ saved" { + t.Errorf("$ saved = %v, want no row rather than a zero charged into the total", m.Value) + } + } +} + +// TestPricing_BuiltInTableBeatsFlatFallback: the flat fields are documented as +// covering "models absent from pricing", and a model in the built-in table is not +// absent. Letting one flat input rate shadow every per-model default would +// reintroduce the flat-rate mispricing the table exists to avoid — and silently, +// since the figure would then claim to be operator-configured. +func TestPricing_BuiltInTableBeatsFlatFallback(t *testing.T) { + p := configuredJSON(t, `{"remove":["NotebookEdit"],"input_cost_per_token":9e-05}`) + rates, src := p.cfg.ratesFor("claude-opus-5") + if src != rateDefault { + t.Errorf("source = %v, want rateDefault for a model in the built-in table", src) + } + if rates.InputCostPerToken == 9e-05 { + t.Error("flat fallback shadowed the built-in per-model rate") + } + // A model in neither table still uses the flat fallback. + _, src2 := p.cfg.ratesFor("no-such-model") + if src2 != rateConfigured { + t.Errorf("source = %v, want rateConfigured via the flat fallback", src2) + } + // And the caveat is still attached, because defaults were used. + pruneWithModel(t, p, "claude-opus-5") + if m := findMetric(t, p.Metrics(), "$ saved"); !strings.Contains(m.Note, "default rates") { + t.Errorf("note = %q, want the default-rates caveat", m.Note) + } +} diff --git a/authbridge/cmd/abctl/toolscan/patch.go b/authbridge/cmd/abctl/toolscan/patch.go index 67f9ea380..400482ca1 100644 --- a/authbridge/cmd/abctl/toolscan/patch.go +++ b/authbridge/cmd/abctl/toolscan/patch.go @@ -80,8 +80,7 @@ func PatchConfig(path string, candidates []string) (changed bool, err error) { return false, nil // already current — idempotent } lines[i] = replacement - out := strings.Join(lines, "\n") - if err := os.WriteFile(path, []byte(out), 0o600); err != nil { + if err := writeFileAtomic(path, []byte(strings.Join(lines, "\n"))); err != nil { return false, err } return true, nil @@ -97,3 +96,45 @@ func leadingSpaces(s string) int { } return len(s) } + +// writeFileAtomic replaces path's contents via a temp file and a rename. +// +// os.WriteFile truncates in place, which has two failure modes on a live config: +// a crash mid-write leaves a truncated file with no copy to recover from, and +// even on the success path the proxy's fsnotify reloader can wake on the +// truncated intermediate state and reject its own config. A rename is atomic, so +// a reader sees either the old file or the new one. +// +// The temp file is created in the same directory so the rename stays within one +// filesystem, and the destination's existing mode is preserved — the file already +// exists (PatchConfig read it), so its permissions are the operator's to keep. +func writeFileAtomic(path string, data []byte) error { + dir := filepath.Dir(path) + mode := os.FileMode(0o600) + if fi, err := os.Stat(path); err == nil { + mode = fi.Mode().Perm() + } + tmp, err := os.CreateTemp(dir, "."+filepath.Base(path)+".tmp") + if err != nil { + return err + } + tmpName := tmp.Name() + defer os.Remove(tmpName) // no-op once the rename succeeds + if _, err := tmp.Write(data); err != nil { + tmp.Close() + return err + } + // Sync before rename: without it a crash after the rename can leave the + // new name pointing at unflushed (zero-length) content. + if err := tmp.Sync(); err != nil { + tmp.Close() + return err + } + if err := tmp.Close(); err != nil { + return err + } + if err := os.Chmod(tmpName, mode); err != nil { + return err + } + return os.Rename(tmpName, path) +} diff --git a/authbridge/cmd/abctl/toolscan/scan.go b/authbridge/cmd/abctl/toolscan/scan.go index be8f78c9f..e99605d15 100644 --- a/authbridge/cmd/abctl/toolscan/scan.go +++ b/authbridge/cmd/abctl/toolscan/scan.go @@ -2,6 +2,7 @@ package toolscan import ( "bufio" + "bytes" "encoding/json" "fmt" "io/fs" @@ -119,7 +120,10 @@ func scanFile(path string, since time.Time, seenIDs map[string]struct{}, res *Re line := sc.Bytes() // Hot path: the overwhelming majority of lines carry no tool call. // A literal substring check is far cheaper than parsing them. - if !strings.Contains(string(line), `"tool_use"`) { + // bytes.Contains, not strings.Contains(string(line), …): converting would + // copy every candidate line, and this is the hot path the prefilter exists + // to keep cheap. + if !bytes.Contains(line, []byte(`"tool_use"`)) { continue } res.Lines++ diff --git a/authbridge/cmd/abctl/tui/app.go b/authbridge/cmd/abctl/tui/app.go index 1aa91abed..d133031cc 100644 --- a/authbridge/cmd/abctl/tui/app.go +++ b/authbridge/cmd/abctl/tui/app.go @@ -249,6 +249,10 @@ type model struct { // until then. pipeline *apiclient.PipelineView + // pipelineFetching is set while a /v1/pipeline request is outstanding, so + // the 2s refresh tick cannot stack fetches against a slow endpoint. + pipelineFetching bool + // helpVisible toggles the [?] key-help overlay. Deliberately a flag // rather than a paneID: the overlay must be openable over ANY pane // (picker included) without disturbing m.pane / m.previousPane, which @@ -449,7 +453,9 @@ func (m *model) loadPipelineCmd() tea.Cmd { return func() tea.Msg { pv, err := m.client.GetPipeline(m.ctx) if err != nil { - return errMsg{where: "get pipeline", err: err} + // Report as a load with no view so the in-flight flag clears; a + // failure that left it set would wedge refresh for the session. + return pipelineLoadedMsg(nil) } return pipelineLoadedMsg(pv) } @@ -591,12 +597,20 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { // they were when the session was first opened. Skipped elsewhere: // the composition itself does not change, so polling it while nobody // is looking at metrics would be pure overhead. - if m.pane == panePluginDetail || m.pane == panePipeline { + // Guard against stacking fetches: the tick is 2s and the HTTP timeout is + // 10s, so a stalled endpoint would otherwise accumulate ~5 concurrent + // requests and keep adding one every tick. + if (m.pane == panePluginDetail || m.pane == panePipeline) && !m.pipelineFetching { + m.pipelineFetching = true return m, tea.Batch(m.loadSessionsCmd(), m.loadPipelineCmd(), refreshTickCmd()) } return m, tea.Batch(m.loadSessionsCmd(), refreshTickCmd()) case pipelineLoadedMsg: + m.pipelineFetching = false + if msg == nil { + return m, nil // fetch failed; keep the view we have + } m.pipeline = (*apiclient.PipelineView)(msg) m.rebuildPipelineTable() // Re-render an open plugin detail pane against the new view. Without diff --git a/authbridge/cmd/abctl/tui/events_pane.go b/authbridge/cmd/abctl/tui/events_pane.go index cd4de5b57..4a0ee14a0 100644 --- a/authbridge/cmd/abctl/tui/events_pane.go +++ b/authbridge/cmd/abctl/tui/events_pane.go @@ -532,6 +532,14 @@ func computeEventPairs(rows []eventRow) (map[*pipeline.SessionEvent]int, map[int if _, done := partner[j]; done { continue // already paired exactly by RequestID } + if rj.RequestID != "" { + // It carried an id and still did not pair — a second response for + // the same request (a retry, or a streamed reply recorded twice). + // Letting it fall through would have the heuristic walk back and + // claim an unrelated earlier request, which is exactly the + // mis-attribution the id was added to end. Leave it unpaired. + continue + } for i := j - 1; i >= 0; i-- { if _, taken := partner[i]; taken { continue @@ -781,6 +789,12 @@ func (m *model) tokensCellWithSaving(rows []eventRow, partner map[int]int, i int if resp == nil || resp.Phase != pipeline.SessionResponse { return "" } + // Only price against a response that pairs by id. A heuristically-matched + // response may belong to a different request, and its cache tier would + // then pick the wrong rate — a 12.5x error presented as a measurement. + if ev.RequestID == "" || resp.RequestID != ev.RequestID { + return "" + } tokens, usd, ok := savedTokensAndCost(ps, resp.Inference) if !ok { return "" diff --git a/authbridge/cmd/abctl/tui/prune_saving_test.go b/authbridge/cmd/abctl/tui/prune_saving_test.go index ae7e2aa18..57314241d 100644 --- a/authbridge/cmd/abctl/tui/prune_saving_test.go +++ b/authbridge/cmd/abctl/tui/prune_saving_test.go @@ -129,3 +129,56 @@ func TestFormatCompactAndUSD(t *testing.T) { } } } + +// TestComputeEventPairs_DuplicateResponseStaysUnpaired: a second response sharing +// a RequestID — a retry, or a streamed reply recorded twice — finds the request +// already paired. Falling through to the adjacency heuristic would have it walk +// back and claim an unrelated earlier request, reintroducing exactly the +// mis-attribution the id was added to end. +func TestComputeEventPairs_DuplicateResponseStaysUnpaired(t *testing.T) { + ev := func(phase pipeline.SessionPhase, id string, code int) *pipeline.SessionEvent { + return &pipeline.SessionEvent{ + Direction: pipeline.Outbound, Phase: phase, + Host: "h", RequestID: id, StatusCode: code, + } + } + rows := []eventRow{ + {event: ev(pipeline.SessionRequest, "aaa", 0)}, // 0 + {event: ev(pipeline.SessionRequest, "bbb", 0)}, // 1 + {event: ev(pipeline.SessionResponse, "bbb", 200)}, // 2 pairs with 1 + {event: ev(pipeline.SessionResponse, "bbb", 500)}, // 3 duplicate for bbb + } + _, partner := computeEventPairs(rows) + + if partner[1] != 2 { + t.Errorf("bbb should pair 1↔2, got %v", partner) + } + if j, ok := partner[3]; ok { + t.Errorf("duplicate response paired with row %d; it must stay unpaired", j) + } + if j, ok := partner[0]; ok { + t.Errorf("request aaa was claimed by the duplicate (row %d) — the bug this guards", j) + } +} + +// TestTokensCellWithSaving_RequiresAnIDMatch: pricing against a heuristically +// matched response could take its cache tier from a different request, and the +// tiers are ~12.5x apart — a wrong figure presented as a measurement. +func TestTokensCellWithSaving_RequiresAnIDMatch(t *testing.T) { + req := reqEvent(t, wire) + req.RequestID = "aaa" + resp := &pipeline.SessionEvent{ + Phase: pipeline.SessionResponse, RequestID: "zzz", // different exchange + Inference: &pipeline.InferenceExtension{CacheWriteTokens: 24701, TotalTokens: 33582}, + } + rows := []eventRow{{event: req}, {event: resp}} + m := &model{} + if got := m.tokensCellWithSaving(rows, map[int]int{0: 1, 1: 0}, 0, req); got != "" { + t.Errorf("priced against a mismatched response: %q", got) + } + // Matching ids do price. + resp.RequestID = "aaa" + if got := m.tokensCellWithSaving(rows, map[int]int{0: 1, 1: 0}, 0, req); got == "" { + t.Error("an id-matched pair should price") + } +} diff --git a/authbridge/cmd/authbridge-proxy/demo.go b/authbridge/cmd/authbridge-proxy/demo.go index 1b43fbaa9..3a89d1afd 100644 --- a/authbridge/cmd/authbridge-proxy/demo.go +++ b/authbridge/cmd/authbridge-proxy/demo.go @@ -1,6 +1,8 @@ package main import ( + "errors" + "log/slog" "os" "path/filepath" ) @@ -70,16 +72,28 @@ pipeline: ` } -// writeDemoConfig writes the built-in --demo config next to the CA (in caDir) -// and returns its path, so --demo reuses the normal file-based load + -// hot-reload path — edits to the file are picked up live. caDir is -// caller-resolved (cwd-relative by default, or --ca-dir); no absolute path is -// baked into the binary. Overwrites any prior copy so the preset is canonical. +// writeDemoConfig ensures the built-in --demo config exists next to the CA (in +// caDir) and returns its path, so --demo reuses the normal file-based load + +// hot-reload path. caDir is caller-resolved (cwd-relative by default, or +// --ca-dir); no absolute path is baked into the binary. +// +// An existing file is KEPT, not overwritten. The config's own header invites +// editing it, and `abctl tools scan --write` writes a prune list into it — and +// this function runs before any port is bound, so an unconditional write meant +// that even a --demo start which then failed on a port clash silently destroyed +// those edits. Delete the file to regenerate the preset. func writeDemoConfig(caDir string) (string, error) { if err := os.MkdirAll(caDir, 0o755); err != nil { return "", err } path := filepath.Join(caDir, "demo.yaml") + if _, err := os.Stat(path); err == nil { + slog.Info("demo mode — keeping the existing config (edits and any prune list are preserved)", + "path", path, "hint", "delete it to regenerate the built-in preset") + return path, nil + } else if !errors.Is(err, os.ErrNotExist) { + return "", err + } if err := os.WriteFile(path, []byte(demoConfigYAML(caDir)), 0o644); err != nil { return "", err } diff --git a/authbridge/cmd/authbridge-proxy/demo_test.go b/authbridge/cmd/authbridge-proxy/demo_test.go index 296540bc2..31eae9df0 100644 --- a/authbridge/cmd/authbridge-proxy/demo_test.go +++ b/authbridge/cmd/authbridge-proxy/demo_test.go @@ -1,6 +1,7 @@ package main import ( + "os" "path/filepath" "slices" "strings" @@ -100,3 +101,35 @@ func TestDemoConfig_WriteLoadsAndValidates(t *testing.T) { t.Errorf("tool-prune must ship with an empty remove list, got %s", tp.Config) } } + +// TestWriteDemoConfig_PreservesAnExistingFile: the config's own header invites +// editing it, and `abctl tools scan --write` writes a prune list into it. This +// function also runs before any port is bound, so an unconditional overwrite +// meant a --demo start that then failed on a port clash silently destroyed those +// edits — which is exactly how a populated remove list was lost in practice. +func TestWriteDemoConfig_PreservesAnExistingFile(t *testing.T) { + caDir := t.TempDir() + p, err := writeDemoConfig(caDir) + if err != nil { + t.Fatal(err) + } + edited := "# operator edit\nmode: proxy-sidecar\n" + if err := os.WriteFile(p, []byte(edited), 0o644); err != nil { + t.Fatal(err) + } + // A second call — a restart — must not clobber it. + p2, err := writeDemoConfig(caDir) + if err != nil { + t.Fatal(err) + } + if p2 != p { + t.Errorf("path changed: %q vs %q", p2, p) + } + got, err := os.ReadFile(p) + if err != nil { + t.Fatal(err) + } + if string(got) != edited { + t.Errorf("edits were overwritten:\n%s", got) + } +} diff --git a/authbridge/docs/framework-architecture.md b/authbridge/docs/framework-architecture.md index 2f0889ad5..6265d3773 100644 --- a/authbridge/docs/framework-architecture.md +++ b/authbridge/docs/framework-architecture.md @@ -813,7 +813,7 @@ The plugin interface is **not** semver-stable yet (AuthBridge is pre-1.0). Chang - **`pctx.Record` helpers**: `Allow` / `Skip` / `Observe` / `Modify` / `Record` / `DenyAndRecord` on `Context`. Framework-managed attribution (`currentPlugin`, `currentPhase`, `Path`) fills Invocation fields automatically. - **Open plugin registry**: plugins self-register from `init()` via `plugins.RegisterPlugin`. Third-party plugins in external modules drop in via a side-effect import. Closed `registry` map literal removed. - **Config hot-reload**: new `pipeline.Holder` (atomic wrapper) + `authlib/reloader` package (fsnotify-driven). Listeners receive `*Holder` instead of `*Pipeline`; the reloader atomically swaps the holder's contents when the config file changes. `mode` and `listener.*` edits are refused (pod restart required); any other change is picked up within the kubelet sync window (~60s). See §9. -- **Body mutation**: `PluginCapabilities.BodyAccess` split into `ReadsBody` / `WritesRequestBody`. New `pctx.SetBody` / `pctx.SetResponseBody` helpers flip a mutation flag; all three listeners (extproc / forwardproxy / reverseproxy) propagate the rewrite to the upstream with correct `Content-Length` and cleared `Content-Encoding`. `BodyAccess` kept as deprecated alias. See §6, "Body mutation." +- **Body mutation**: `PluginCapabilities.BodyAccess` split into `ReadsBody` / `WritesRequestBody`. New `pctx.SetBody` / `pctx.SetResponseBody` helpers flip a mutation flag; all three listeners (extproc / forwardproxy / reverseproxy) propagate the rewrite to the upstream with correct `Content-Length` and cleared `Content-Encoding`. `BodyAccess` was kept as a deprecated alias at the time; it has since been removed. See §6, "Body mutation." - **Detyped framework**: `pipeline/` no longer imports plugin-specific packages. **Breaking**: `Context.Claims *validation.Claims` → `Context.Identity Identity` (interface with `Subject()`/`ClientID()`/`Scopes()`); plugins publish adapters. `Context.Route` removed (was dead code). `Invocation`'s nine jwt-validation + token-exchange specific fields (`ExpectedIssuer`, `TokenSubject`, `RouteHost`, `CacheHit`, etc.) collapsed into `Details map[string]string`; built-in plugins migrated to `Details["expected_issuer"]` etc. `SessionEvent.TargetAudience` removed (was only populated from dead `pctx.Route`). Third-party plugins get a clean diagnostic slot they can populate without framework edits. - **Single-owner packages relocated**: `authlib/validation` → `authlib/plugins/jwtvalidation/validation`. `authlib/exchange` / `authlib/cache` / `authlib/spiffe` → `authlib/plugins/tokenexchange/{exchange,cache,spiffe}`. Each plugin now lives in its own directory (`plugins/jwtvalidation/plugin.go`, `plugins/tokenexchange/plugin.go`) and self-registers via its own init(). `authlib/bypass`, `authlib/routing`, `authlib/auth` stay shared. - **Plugin relationship declarations**: `PluginCapabilities` extended with four chain-scoped fields — `Requires` (all-must-be-earlier), `RequiresAny` (at-least-one-earlier), `After` (soft ordering), `Claims` (mutex on a semantic resource). Validated at `plugins.Build` time (startup + hot-reload); all errors per chain are collected into one report. `authlib/contracts/claims.go` ships `ClaimAuthorizationHeader` as the initial canonical claim constant. `token-exchange` and `token-broker` migrated to declare it, so configuring both on the same outbound chain now fails startup instead of silently clobbering each other's Authorization header. See [`plugin-reference.md` "Declaring plugin relationships"](./plugin-reference.md#declaring-plugin-relationships). diff --git a/authbridge/docs/plugin-reference.md b/authbridge/docs/plugin-reference.md index f35266fe9..33d485044 100644 --- a/authbridge/docs/plugin-reference.md +++ b/authbridge/docs/plugin-reference.md @@ -218,6 +218,7 @@ fail loud before serving traffic. type PluginCapabilities struct { ReadsBody bool WritesRequestBody bool + WritesResponseBody bool Requires []string // ALL must be present + earlier (hard) RequiresAny []string // AT LEAST ONE must be present + run after it (hard) @@ -763,7 +764,8 @@ nothing about how the response may be relayed. |---|---|---| | request-only mutator (`tool-prune`, `context-guru`) | `WritesRequestBody` | yes | | response mutator | `WritesResponseBody` | no — buffered | -| both (`sparc`, `cpex`) | both | no — buffered | +| response mutator (`sparc`) | `WritesResponseBody` | no — buffered | +| both (`cpex`) | both | no — buffered | | pure reader (parsers) | `ReadsBody` | yes | ### Build-time validation (enforced by `pipeline.New`) diff --git a/authbridge/docs/tool-prune-plugin.md b/authbridge/docs/tool-prune-plugin.md index f6abbdbcf..64fd46835 100644 --- a/authbridge/docs/tool-prune-plugin.md +++ b/authbridge/docs/tool-prune-plugin.md @@ -20,7 +20,8 @@ pipeline: - inference-parser - mcp-parser - name: tool-prune - on_error: observe # measure only; switch to enforce when trusted + # on_error defaults to enforce; the empty remove list is what gates the + # plugin. Set observe when you want a projection instead — see below. config: remove: [NotebookEdit, ScheduleWakeup, TaskOutput] ``` @@ -179,12 +180,6 @@ model outright: input_cost_per_token: 0.0000038 ``` -**Deriving your own rates.** If your gateway reports cost on non-streaming -responses (LiteLLM's `x-litellm-response-cost`), send two non-streaming requests -of different prompt length and difference them: `rate = Δcost / Δinput_tokens`. -Repeat with a `cache_control` block sent twice for the write and read rates. This -is how the table above was obtained, and it is exact for your deployment. - Model keys match what the parser records (`Extensions.Inference.Model`) and are matched case-insensitively, since gateways vary in how they echo the name and a case mismatch would silently unprice the traffic. @@ -218,7 +213,8 @@ falls back to configured rates for streams. A saving is also a counterfactual — the cost of a request that was never sent — so even with a cost header it could only ever be priced from rates, not measured. -Counters are in-memory and per-process. That is the right trade for the +Counters are in-memory and per-process, and reset on a config hot-reload as well +as a restart — a reload rebuilds the plugin. That is the right trade for the single-laptop case this targets and what keeps the plugin free of a storage dependency; fleet aggregation belongs on the stats server later and would not change the plugin. diff --git a/authbridge/install-demo.sh b/authbridge/install-demo.sh index 7ca0871a9..3c87862b5 100755 --- a/authbridge/install-demo.sh +++ b/authbridge/install-demo.sh @@ -197,8 +197,9 @@ else fi info "" -# tool-prune ships inert: present in the pipeline with an empty remove list and -# on_error: observe. Offer the scan that fills the list in. Only patch a config +# tool-prune ships inert: the remove list is empty, so it does nothing until a +# name is added. That empty list is the guard — on_error is enforce, because the +# list is what gates the plugin. Offer the scan that fills it in. Only patch a config # that already exists, so a first run never rewrites a file it just created # behind the user's back -- print the command instead and let them look first. demo_cfg="${ca_dir}/demo.yaml" diff --git a/docs/proposals/tool-prune.md b/docs/proposals/tool-prune.md index 58ae52c45..019e84728 100644 --- a/docs/proposals/tool-prune.md +++ b/docs/proposals/tool-prune.md @@ -149,7 +149,7 @@ Every in-tree plugin that declares the capability today, and what it actually do | Plugin | Rewrites request | Rewrites response | Evidence | After the change | |---|---|---|---|---| | `context-guru` | yes | **no** | `contextguru/plugin.go:160`; no `SetResponseBody` call anywhere | `WritesRequestBody` — **gains** response streaming | -| `sparc` | yes | yes | `sparc/plugin.go:214`; `sparc/respond.go:111,122` | both flags — unchanged | +| `sparc` | **no** | yes | `sparc/respond.go:111,122`; calls `SetBody` nowhere | `WritesResponseBody` only — the request flag was stale, and dropping it frees the request-mutator slot so `[sparc, tool-prune]` builds | | `cpex` | yes | yes | `cpex/plugin.go:122`; `cmf_body.go:609`, `cmf_a2a.go:216`, `cmf_inference.go:218` | both flags — unchanged | | `tool-prune` | yes | no | new | `WritesRequestBody` — streams | @@ -253,7 +253,7 @@ pipeline: - mcp-parser - a2a-parser - name: tool-prune - on_error: observe # measure only; switch to enforce when trusted + # on_error defaults to enforce; the empty remove list is the gate. config: remove: [NotebookEdit, ScheduleWakeup, TaskOutput] ``` @@ -273,9 +273,11 @@ This is why one registration suffices: the same plugin code serves measure and enforce, selected by one word of configuration. `context.go` states the intent directly — "Plugin code therefore looks identical under enforce and observe." -Off-by-default is satisfied structurally: the plugin is absent from the shipped -pipeline list until a user adds it, and the documented first step adds it with -`on_error: observe`. +Off-by-default is satisfied structurally, and by the remove list rather than the +policy: an empty `remove` is a no-op whatever `on_error` says, so filling the list +is the single deliberate act that enables the plugin. `on_error: observe` remains +available as a projection mode, but is not the shipped default — two guards where +one suffices only added a step operators skipped. ### Where the list comes from: `abctl tools scan` @@ -546,7 +548,7 @@ For part 3: | Risk | Mitigation | |---|---| -| Removing a tool the agent needs | Unknown names always kept; `--keep` override; ship with `on_error: observe`; fail open on any error | +| Removing a tool the agent needs | Unknown names always kept; a tool forced by `tool_choice` never pruned; `--keep` override; empty `remove` ships as the off switch; fail open on any error | | Stale bundled tool set as Claude Code evolves | Drift reduces savings only; plugin warns on configured names never observed in `ext.Tools` | | One-off prompt-cache invalidation when the list changes | Inherent and bounded: static list means it happens once, then the prefix is stable | | Commit 1 conflicts with in-flight branches declaring `WritesBody` | One-line fix per branch; the compile error makes it self-evident | From 1a71a91311a092c54d75eeabb2b001564b3d598a Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 19:21:48 -0400 Subject: [PATCH 24/28] =?UTF-8?q?fix:=20Address=20review=20pass=202=20?= =?UTF-8?q?=E2=80=94=20six=20blockers=20on=20tool-prune?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All six verified before fixing; two did not reproduce and are recorded as such rather than claimed. Blockers: 1. Tools cited by conversation history are no longer pruned. A provider may reject a request whose tool_use / tool_result blocks reference a tool the manifest no longer defines, and enabling the plugin mid-conversation is exactly when that arises: the config hot-reloads, and the scan's rolling window can propose a tool used earlier in the same session. NOT reproducible against the gateway available here (a pruned manifest with citing history returned 200), so this is a guard against a plausible provider difference, not a demonstrated 400 — but a few unpruned definitions against a failed request is not a trade worth making. 2. A prompt-cache breakpoint carried by a pruned element is moved to the last surviving tool. Claude Code marks the last tool with cache_control; deleting that element deleted the breakpoint, turning every later turn into a full cache write — which costs far more than the definitions saved, so the plugin could have made spend worse. Also guarded against duplicating a marker when a survivor already has one, which would exceed the provider's limit. The available gateway reports no cache tokens either way, so this rests on the structural argument, not a measurement. 3. toolscan now checks sc.Err(). A scanner error silently stopped iteration, under-reporting which tools were CALLED and so proposing MORE for removal — failing toward removing a tool the agent needs, the one direction this must not fail in. 4. PatchConfig refuses a block-style remove: list instead of corrupting it. Replacing only the `remove:` line left its `- item` children dangling under an inline value: invalid YAML the proxy rejects on reload, leaving the operator with a file this tool broke. 5. laptop-token-savings.md had `AUTHBRIDGE_INSTALL_ONLY=1 curl … | sh`, which sets the variable for curl, not for the shell that runs the script. Verified: the piped sh sees it unset, so step 1 would have started the demo, bound 47600-47602, written the cortex-ca/demo.yaml the doc warns about, and made step 3 fail on the port clash. Moved onto sh, with the reason inline. I had claimed to run every step verbatim; I ran steps 2-4 and skipped 1 because the released binaries lack the plugin, so the claim was wrong. 6. Bare claude-sonnet-5 and claude-haiku-4-5 added to the rate table — without them the "no extra configuration" promise failed for the doc's own direct-to-vendor path. The caveat now names the direction of the error: "built-in rates (discounted gateway; understates list pricing)". Anyone paying vendor list is under-credited several-fold, and the laptop doc says to read the figure as a floor. Also from the review: DeriveError documents that it needs a body-reading plugin and silently no-ops on auth-only and lite chains; MetricsProvider carries a producer contract that Name/Note never hold request content, with a length bound in the session API — redact.JSON was not the answer, since it filters by key and the exposure is a value; and the deferred reverse-order reader gap now logs a startup warning naming the reader, so its only record is no longer a code comment an operator would never read. Signed-off-by: Hai Huang --- authbridge/authlib/pipeline/metrics.go | 8 ++ authbridge/authlib/pipeline/pipeline.go | 32 ++++++++ authbridge/authlib/pipeline/snapshot.go | 8 ++ authbridge/authlib/plugins/toolprune/event.go | 33 +++++++- .../authlib/plugins/toolprune/metrics.go | 5 +- .../authlib/plugins/toolprune/plugin.go | 48 +++++++++++ .../authlib/plugins/toolprune/plugin_test.go | 82 +++++++++++++++++-- .../authlib/plugins/toolprune/pricing.go | 18 +++- authbridge/authlib/sessionapi/server.go | 35 +++++++- authbridge/cmd/abctl/toolscan/patch.go | 25 ++++++ authbridge/cmd/abctl/toolscan/patch_test.go | 52 ++++++++++++ authbridge/cmd/abctl/toolscan/scan.go | 8 ++ authbridge/docs/laptop-token-savings.md | 26 ++++-- 13 files changed, 360 insertions(+), 20 deletions(-) diff --git a/authbridge/authlib/pipeline/metrics.go b/authbridge/authlib/pipeline/metrics.go index 1fd5322d5..1cab473e6 100644 --- a/authbridge/authlib/pipeline/metrics.go +++ b/authbridge/authlib/pipeline/metrics.go @@ -22,6 +22,14 @@ type Metric struct { // must not block — take a mutex, copy, release. Returning nil is fine and // renders as "(none)". // +// CONTRACT ON Name AND Note: these are short operator-facing labels, surfaced on +// an endpoint with no authentication. They must never carry request or response +// content — no prompts, no completions, no header or credential values, nothing +// derived from a body. A caveat naming a sample size or a configuration key is +// fine; a caveat quoting the traffic is not. The framework caps their length but +// cannot inspect their meaning, so this is a producer obligation, the same one +// body-mutation events carry when they publish only lengths and hashes. +// // This is deliberately separate from plugins.StatsSource / auth.Stats, which // are auth-shaped: they carry typed approval and denial enums and a custom // MarshalJSON, so routing "bytes removed" through them would distort their diff --git a/authbridge/authlib/pipeline/pipeline.go b/authbridge/authlib/pipeline/pipeline.go index e77542fc6..ced57ec15 100644 --- a/authbridge/authlib/pipeline/pipeline.go +++ b/authbridge/authlib/pipeline/pipeline.go @@ -621,8 +621,40 @@ func validateCapabilities(plugins []Plugin) error { readerAfterMutator = plugin.Name() } } + warnResponseReaderOrdering(plugins) if readerAfterMutator != "" { return fmt.Errorf("pipeline: plugin %q reads body after mutator %q — body readers must precede the mutator so they see the original bytes", readerAfterMutator, firstMutator) } return nil } + +// warnResponseReaderOrdering logs the chain shape that the documented +// reverse-order gap makes unsafe: a non-streaming body reader placed BEFORE a +// response mutator. RunResponse iterates in reverse, so the mutator runs first +// and the reader sees rewritten response bytes — for a policy plugin that means +// authorizing against content it did not receive. +// +// A warning rather than a rejection: enforcing it would fail chains that +// validate today (see the gap comment in validateCapabilities), and this change +// promised no working configuration starts failing. But the deferral should not +// be invisible — until now its only record was a code comment, which an operator +// running the shape would never read. +func warnResponseReaderOrdering(plugins []Plugin) { + var respMutator string + for _, p := range plugins { + caps := p.Capabilities().Normalize() + if caps.WritesResponseBody { + respMutator = p.Name() + continue + } + if respMutator != "" || !caps.ReadsBody { + continue + } + if _, streaming := p.(StreamingResponder); streaming { + continue // RunResponse skips these entirely + } + slog.Warn("pipeline: body reader precedes a response mutator — on the response pass the mutator runs first, so this reader sees rewritten bytes", + "reader", p.Name(), + "hint", "place the reader after the response mutator, or confirm it does not read pctx.ResponseBody") + } +} diff --git a/authbridge/authlib/pipeline/snapshot.go b/authbridge/authlib/pipeline/snapshot.go index d144b183f..e3d05bd81 100644 --- a/authbridge/authlib/pipeline/snapshot.go +++ b/authbridge/authlib/pipeline/snapshot.go @@ -142,6 +142,14 @@ func DeriveError(pctx *Context) *EventError { // upstreamErrorKind extracts the provider's machine-readable error type from an // error response body, or "" when there isn't one. // +// REQUIRES A BUFFERED BODY. pctx.ResponseBody is only populated when some +// plugin in the chain declares ReadsBody, so on an auth-only chain — and in +// the authbridge-lite build, where the parsers are compiled out — this yields +// "" and the event stays the bare backend_error/ it was before. That is +// precisely where an operator has the fewest other diagnostics; closing it +// would mean buffering error responses on chains that otherwise never read a +// body, which is a listener-level decision, not one to make here. +// // A bare `backend_error / 400` tells an operator nothing about why, which turns // every upstream rejection into a guessing exercise. The provider already // classifies its own failures, and the classification is what an operator acts diff --git a/authbridge/authlib/plugins/toolprune/event.go b/authbridge/authlib/plugins/toolprune/event.go index 19310efb7..4e47075a7 100644 --- a/authbridge/authlib/plugins/toolprune/event.go +++ b/authbridge/authlib/plugins/toolprune/event.go @@ -1,6 +1,10 @@ package toolprune -import "github.com/rossoctl/cortex/authbridge/authlib/pipeline" +import ( + "github.com/tidwall/gjson" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" +) // pruneEvent is the per-request record published under "tool-prune/event", so a // consumer can show what this one request saved instead of only an aggregate. @@ -43,3 +47,30 @@ func inferenceModel(pctx *pipeline.Context) string { } return pctx.Extensions.Inference.Model } + +// toolsCitedByHistory returns the tool names the conversation already used, from +// tool_use blocks in assistant messages. +// +// Those tools must survive pruning: a provider may reject a request whose history +// references a tool the manifest no longer defines. Enabling the plugin +// mid-conversation is exactly when this arises, because the config hot-reloads and +// the scan's window can propose a tool that was used earlier in the same session. +// +// Scanned with gjson paths rather than a full unmarshal — a Claude Code body runs +// to hundreds of KB and this is the request hot path. +func toolsCitedByHistory(body []byte) map[string]struct{} { + out := map[string]struct{}{} + gjson.GetBytes(body, "messages").ForEach(func(_, msg gjson.Result) bool { + msg.Get("content").ForEach(func(_, block gjson.Result) bool { + if block.Get("type").String() != "tool_use" { + return true + } + if n := block.Get("name"); n.Type == gjson.String && n.String() != "" { + out[n.String()] = struct{}{} + } + return true + }) + return true + }) + return out +} diff --git a/authbridge/authlib/plugins/toolprune/metrics.go b/authbridge/authlib/plugins/toolprune/metrics.go index 750d389db..7b65d62ec 100644 --- a/authbridge/authlib/plugins/toolprune/metrics.go +++ b/authbridge/authlib/plugins/toolprune/metrics.go @@ -191,7 +191,10 @@ func (m *metrics) snapshot() []pipeline.Metric { // Provenance travels with the number. Built-in rates are // gateway-specific and not refreshed, so a figure derived from them // must not read as one measured on this account. - costNote = "default rates — set pricing. to use yours" + // Name the provenance, not just the fact. "default rates" alone reads + // as a rounding caveat; these were measured on a discounted gateway, + // so for anyone paying vendor list the figure is several times low. + costNote = "built-in rates (discounted gateway; understates list pricing) — set pricing." if note != "" { costNote = note + "; " + costNote } diff --git a/authbridge/authlib/plugins/toolprune/plugin.go b/authbridge/authlib/plugins/toolprune/plugin.go index bfe165461..89f6df504 100644 --- a/authbridge/authlib/plugins/toolprune/plugin.go +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -377,6 +377,18 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action }) return action } + // Tools the conversation already used must stay in the manifest. A provider + // may reject a tool_use / tool_result block that references a tool the + // request no longer defines, and enabling the plugin mid-conversation (the + // config hot-reloads) is exactly when history can cite a tool the scan + // proposed — the scan only looks at a rolling window, so a tool used earlier + // in this very session can be on the remove list. + // + // Not reproducible against every provider (one gateway accepts it), but the + // cost of the guard is a few unpruned definitions and the cost of being wrong + // is a failed request, so it is not a trade worth making. + used := toolsCitedByHistory(body) + var victims []int var anyNameResolved bool names := make([]string, 0, len(raw)) @@ -392,6 +404,10 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action slog.Debug("tool-prune: keeping tool forced by tool_choice", "tool", name) continue } + if _, cited := used[name]; cited { + slog.Debug("tool-prune: keeping tool cited by conversation history", "tool", name) + continue + } if _, ok := p.remove[name]; ok { victims = append(victims, i) names = append(names, name) @@ -426,6 +442,38 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action } } } else { + // A prompt-cache breakpoint rides on one element (Claude Code marks the + // last tool). Deleting that element deletes the breakpoint, and losing + // it turns every subsequent turn into a full cache write — which costs + // far more than the definitions saved. Carry the marker to the last + // surviving tool instead. + victimSet := make(map[int]bool, len(victims)) + for _, v := range victims { + victimSet[v] = true + } + var orphanedCacheControl gjson.Result + for _, v := range victims { + if cc := gjson.GetBytes(body, fmt.Sprintf("tools.%d.cache_control", v)); cc.Exists() { + orphanedCacheControl = cc + } + } + lastSurvivor := -1 + for i := len(raw) - 1; i >= 0; i-- { + if !victimSet[i] { + lastSurvivor = i + break + } + } + if orphanedCacheControl.Exists() && lastSurvivor >= 0 && + !gjson.GetBytes(body, fmt.Sprintf("tools.%d.cache_control", lastSurvivor)).Exists() { + if out, err = sjson.SetRawBytes(out, + fmt.Sprintf("tools.%d.cache_control", lastSurvivor), + []byte(orphanedCacheControl.Raw)); err != nil { + slog.Warn("tool-prune: could not preserve cache_control, forwarding original", "err", err) + return action + } + slog.Debug("tool-prune: moved cache_control to the last surviving tool", "index", lastSurvivor) + } // Descending, so an earlier deletion never shifts a later index. for i := len(victims) - 1; i >= 0; i-- { if out, err = sjson.DeleteBytes(out, fmt.Sprintf("tools.%d", victims[i])); err != nil { diff --git a/authbridge/authlib/plugins/toolprune/plugin_test.go b/authbridge/authlib/plugins/toolprune/plugin_test.go index f3ea93192..6d4995676 100644 --- a/authbridge/authlib/plugins/toolprune/plugin_test.go +++ b/authbridge/authlib/plugins/toolprune/plugin_test.go @@ -380,11 +380,14 @@ func TestPricing_DefaultsPriceKnownModelsWithoutConfig(t *testing.T) { } // Provenance must travel with the number: built-in rates are // gateway-specific and never refreshed, so this must not read as measured. - if !strings.Contains(m.Note, "default rates") { - t.Errorf("note = %q, want it to disclose that default rates were used", m.Note) - } - if !strings.Contains(m.Note, "pricing.") { - t.Errorf("note = %q, want it to name how to override", m.Note) + // The note must disclose three things: that the rates are built in, the + // DIRECTION of the error (they came from a discounted gateway, so anyone on + // vendor list is under-credited), and how to override. "default rates" alone + // reads as a rounding caveat rather than a several-fold one. + for _, want := range []string{"built-in rates", "understates", "pricing."} { + if !strings.Contains(m.Note, want) { + t.Errorf("note = %q, missing %q", m.Note, want) + } } } @@ -896,7 +899,72 @@ func TestPricing_BuiltInTableBeatsFlatFallback(t *testing.T) { } // And the caveat is still attached, because defaults were used. pruneWithModel(t, p, "claude-opus-5") - if m := findMetric(t, p.Metrics(), "$ saved"); !strings.Contains(m.Note, "default rates") { - t.Errorf("note = %q, want the default-rates caveat", m.Note) + if m := findMetric(t, p.Metrics(), "$ saved"); !strings.Contains(m.Note, "built-in rates") { + t.Errorf("note = %q, want the built-in-rates caveat", m.Note) + } +} + +// TestPrune_KeepsToolsCitedByHistory: a provider may reject a request whose +// history references a tool the manifest no longer defines. This arises exactly +// when the plugin is enabled mid-conversation — the config hot-reloads, and the +// scan's rolling window can propose a tool used earlier in the same session. +func TestPrune_KeepsToolsCitedByHistory(t *testing.T) { + body := `{"tools":[{"name":"Read"},{"name":"WebSearch"},{"name":"NotebookEdit"}], + "messages":[ + {"role":"user","content":"go"}, + {"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"WebSearch","input":{}}]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"t1","content":"ok"}]}]}` + p := configured(t, "WebSearch", "NotebookEdit") + pctx := inferenceCtx("/v1/messages", body, "Read", "WebSearch", "NotebookEdit") + run(t, p, pctx) + + got := string(pctx.Body) + if !strings.Contains(got, "WebSearch") { + t.Errorf("WebSearch is cited by history and must survive:\n%s", got) + } + if strings.Contains(got, "NotebookEdit") { + t.Errorf("NotebookEdit is uncited and should still be pruned:\n%s", got) + } +} + +// TestPrune_PreservesCacheControlBreakpoint: a prompt-cache breakpoint rides on +// one element — Claude Code marks the last tool. Deleting that element deletes +// the breakpoint, and losing it turns every later turn into a full cache write, +// which costs far more than the definitions saved. The marker must move to the +// last surviving tool. +func TestPrune_PreservesCacheControlBreakpoint(t *testing.T) { + body := `{"tools":[{"name":"Read"},{"name":"Bash"},` + + `{"name":"NotebookEdit","cache_control":{"type":"ephemeral"}}]}` + p := configured(t, "NotebookEdit") + pctx := inferenceCtx("/v1/messages", body, "Read", "Bash", "NotebookEdit") + run(t, p, pctx) + + got := string(pctx.Body) + if strings.Contains(got, "NotebookEdit") { + t.Fatalf("the tool should be pruned: %s", got) + } + if !strings.Contains(got, "cache_control") { + t.Errorf("cache breakpoint was destroyed — every later turn becomes a full cache write:\n%s", got) + } + // It must land on the LAST surviving tool, where the prefix ends. + if !strings.Contains(got, `{"name":"Bash","cache_control":{"type":"ephemeral"}}`) { + t.Errorf("marker not on the last survivor:\n%s", got) + } + var any map[string]any + if err := json.Unmarshal(pctx.Body, &any); err != nil { + t.Errorf("invalid JSON after the move: %v", err) + } +} + +// TestPrune_DoesNotDuplicateCacheControl: when a surviving tool already carries a +// breakpoint, adding another would exceed the provider's cache_control limit. +func TestPrune_DoesNotDuplicateCacheControl(t *testing.T) { + body := `{"tools":[{"name":"Read","cache_control":{"type":"ephemeral"}},` + + `{"name":"NotebookEdit","cache_control":{"type":"ephemeral"}}]}` + p := configured(t, "NotebookEdit") + pctx := inferenceCtx("/v1/messages", body, "Read", "NotebookEdit") + run(t, p, pctx) + if n := strings.Count(string(pctx.Body), "cache_control"); n != 1 { + t.Errorf("cache_control appears %d times, want 1: %s", n, pctx.Body) } } diff --git a/authbridge/authlib/plugins/toolprune/pricing.go b/authbridge/authlib/plugins/toolprune/pricing.go index 4b8c2ae8c..65f9da757 100644 --- a/authbridge/authlib/plugins/toolprune/pricing.go +++ b/authbridge/authlib/plugins/toolprune/pricing.go @@ -10,9 +10,11 @@ package toolprune // between a number an operator sees and one they never get around to enabling. // They are a starting point, not a fact about your account: // -// - Rates are gateway-specific. This gateway bills well below Anthropic list; -// a deployment talking straight to the vendor pays more, so these would -// understate its saving. +// - Rates are gateway-specific. This gateway bills well below vendor list, so +// a deployment talking straight to Anthropic pays more and these +// UNDERSTATE its saving — by roughly 4x on the input tier at the time of +// measurement. That is the common case for a laptop install, so the +// caveat travels with every figure rather than living only here. // - Rates change. Nothing here refreshes them. // // A figure derived from these is therefore labelled as coming from default @@ -33,6 +35,16 @@ var defaultPricing = map[string]modelRates{ CacheWriteCostPerToken: 0.00000475, CacheReadCostPerToken: 0.00000038, }, + "claude-sonnet-5": { + InputCostPerToken: 0.00000152, + CacheWriteCostPerToken: 0.0000019, + CacheReadCostPerToken: 0.000000152, + }, + "claude-haiku-4-5": { + InputCostPerToken: 0.00000076, + CacheWriteCostPerToken: 0.00000095, + CacheReadCostPerToken: 0.000000076, + }, "aws/claude-sonnet-5": { InputCostPerToken: 0.00000152, CacheWriteCostPerToken: 0.0000019, diff --git a/authbridge/authlib/sessionapi/server.go b/authbridge/authlib/sessionapi/server.go index ccba81017..90a4c71b7 100644 --- a/authbridge/authlib/sessionapi/server.go +++ b/authbridge/authlib/sessionapi/server.go @@ -239,7 +239,11 @@ func describePipeline(h *pipeline.Holder, direction string) []pipelinePluginView view.Config = redact.JSON(rc.RawConfig()) } if mp, ok := pl.(pipeline.MetricsProvider); ok { - view.Metrics = mp.Metrics() + // Bounded, not redacted: Metric.Name and Metric.Note are free-text + // and plugin-controlled, and this endpoint has no authentication. A + // key-based redactor cannot help with a value, so the framework caps + // the length and MetricsProvider carries the contract. + view.Metrics = boundMetrics(mp.Metrics()) } out[i] = view } @@ -342,3 +346,32 @@ func (s *Server) handleStream(w http.ResponseWriter, r *http.Request) { } } } + +// boundMetrics caps each metric's free-text fields. +// +// redact.JSON is not the right tool here: it filters by KEY name (api_key, +// token, …), and the exposure on this channel is a VALUE — a plugin putting +// request-derived text into Metric.Name or Metric.Note. Running metrics through +// a key-based filter would be a no-op that looked like a control. +// +// What the framework can enforce is a bound, so a plugin cannot stream content +// through a field meant for short labels. The rest is a producer contract, stated +// on pipeline.MetricsProvider: these fields carry labels and caveats, never +// request or response content. The session API has no authentication. +func boundMetrics(in []pipeline.Metric) []pipeline.Metric { + const maxLabel = 120 + if len(in) == 0 { + return nil + } + out := make([]pipeline.Metric, len(in)) + for i, m := range in { + if len(m.Name) > maxLabel { + m.Name = m.Name[:maxLabel] + } + if len(m.Note) > maxLabel { + m.Note = m.Note[:maxLabel] + } + out[i] = m + } + return out +} diff --git a/authbridge/cmd/abctl/toolscan/patch.go b/authbridge/cmd/abctl/toolscan/patch.go index 400482ca1..7aca6e019 100644 --- a/authbridge/cmd/abctl/toolscan/patch.go +++ b/authbridge/cmd/abctl/toolscan/patch.go @@ -75,6 +75,15 @@ func PatchConfig(path string, candidates []string) (changed bool, err error) { if m == nil { continue } + // Refuse the block-list form. Replacing just the `remove:` line would + // leave its `- item` children dangling under a now-inline value, which + // is invalid YAML — the proxy would reject the config on reload and the + // operator would be left with a file this tool broke. + if isBlockList(lines, i, end) { + return false, fmt.Errorf("%s: the tool-prune `remove:` list is in block form (one `- item` per line);\n"+ + " this tool only rewrites the inline form. Replace those lines with `remove: []` and re-run,\n"+ + " or paste the block this command prints without --write", path) + } replacement := m[1] + want if lines[i] == replacement { return false, nil // already current — idempotent @@ -138,3 +147,19 @@ func writeFileAtomic(path string, data []byte) error { } return os.Rename(tmpName, path) } + +// isBlockList reports whether the `remove:` at lines[i] is followed by YAML +// block-sequence items rather than carrying an inline value. +func isBlockList(lines []string, i, end int) bool { + if strings.TrimSpace(strings.SplitN(lines[i], ":", 2)[1]) != "" { + return false // has an inline value on the same line + } + for j := i + 1; j < end && j < len(lines); j++ { + t := strings.TrimSpace(lines[j]) + if t == "" || strings.HasPrefix(t, "#") { + continue + } + return strings.HasPrefix(t, "- ") + } + return false +} diff --git a/authbridge/cmd/abctl/toolscan/patch_test.go b/authbridge/cmd/abctl/toolscan/patch_test.go index fa2f5ce3c..7389ea5a0 100644 --- a/authbridge/cmd/abctl/toolscan/patch_test.go +++ b/authbridge/cmd/abctl/toolscan/patch_test.go @@ -193,3 +193,55 @@ func TestPatchConfig_MissingFileExplainsWhere(t *testing.T) { t.Errorf("should replace the bare os error, not wrap it:\n%s", msg) } } + +// TestPatchConfig_RefusesBlockStyleList: replacing only the `remove:` line would +// leave its `- item` children dangling under an inline value — invalid YAML the +// proxy rejects on reload, leaving the operator with a file this tool broke. +func TestPatchConfig_RefusesBlockStyleList(t *testing.T) { + cfg := `pipeline: + outbound: + plugins: + - name: tool-prune + config: + remove: + - NotebookEdit + - WebSearch + - name: token-exchange +` + p := writeConfig(t, cfg) + _, err := PatchConfig(p, []string{"LSP"}) + if err == nil { + t.Fatal("expected a refusal for the block-list form") + } + for _, want := range []string{"block form", "remove: []"} { + if !strings.Contains(err.Error(), want) { + t.Errorf("error should mention %q: %v", want, err) + } + } + // And it must not have touched the file. + got, _ := os.ReadFile(p) + if string(got) != cfg { + t.Errorf("file was modified despite the refusal:\n%s", got) + } +} + +// TestIsBlockList distinguishes the two spellings. +func TestIsBlockList(t *testing.T) { + inline := []string{" remove: [A, B]"} + if isBlockList(inline, 0, 1) { + t.Error("inline form misdetected as a block list") + } + empty := []string{" remove: []"} + if isBlockList(empty, 0, 1) { + t.Error("empty inline form misdetected") + } + block := []string{" remove:", " - A", " - B"} + if !isBlockList(block, 0, 3) { + t.Error("block form not detected") + } + // A bare `remove:` with a following key (not a list) is not a block list. + bare := []string{" remove:", " other: 1"} + if isBlockList(bare, 0, 2) { + t.Error("bare key followed by another key misdetected as a block list") + } +} diff --git a/authbridge/cmd/abctl/toolscan/scan.go b/authbridge/cmd/abctl/toolscan/scan.go index e99605d15..488194779 100644 --- a/authbridge/cmd/abctl/toolscan/scan.go +++ b/authbridge/cmd/abctl/toolscan/scan.go @@ -65,6 +65,7 @@ func Scan(dir string, days int, keep []string) (*Result, error) { return nil } res.Files++ + // Propagate: a partial scan silently proposes more tools. return scanFile(path, since, seenIDs, res) }) if err != nil { @@ -148,6 +149,13 @@ func scanFile(path string, since time.Time, seenIDs map[string]struct{}, res *Re res.CallCounts[c.Name]++ } } + // A scanner error (a line past the 16MB cap, a read fault) silently stops + // iteration. Swallowing it under-reports which tools were CALLED, which + // makes the scan propose MORE for removal — failing toward removing a tool + // the agent needs, the one direction this must not fail in. + if err := sc.Err(); err != nil { + return fmt.Errorf("%s: %w (the tool list would be incomplete, so refusing to guess)", path, err) + } return nil } diff --git a/authbridge/docs/laptop-token-savings.md b/authbridge/docs/laptop-token-savings.md index f68129977..429c89e71 100644 --- a/authbridge/docs/laptop-token-savings.md +++ b/authbridge/docs/laptop-token-savings.md @@ -11,12 +11,17 @@ Four steps, about two minutes. ## 1. Install the binaries ```sh -AUTHBRIDGE_INSTALL_ONLY=1 \ - curl -fsSL https://raw.githubusercontent.com/rossoctl/cortex/main/authbridge/install-demo.sh | sh +curl -fsSL https://raw.githubusercontent.com/rossoctl/cortex/main/authbridge/install-demo.sh \ + | AUTHBRIDGE_INSTALL_ONLY=1 sh ``` -Puts `authbridge-proxy` and `abctl` in `~/.local/bin`. `INSTALL_ONLY` skips the -demo — you want a config that persists, which the next step writes. +Puts `authbridge-proxy` and `abctl` in `~/.local/bin`. `INSTALL_ONLY` skips +starting the demo — you want a config that persists, which the next step writes. + +The variable goes on `sh`, not on `curl`. Written the other way round +(`VAR=1 curl … | sh`) it only reaches `curl`, the script runs without it and +starts the demo — binding 47600-47602 and writing a `cortex-ca/demo.yaml`, so +step 3 then fails on the port clash. ## 2. Write a config @@ -87,11 +92,18 @@ Stop it with `pkill -f 'authbridge-proxy --config'`. ## Seeing the saving in money -`$ saved` and `$ saved / request` appear with no extra configuration — the plugin -ships rates for the Claude models on the rossoctl gateway. The figure is labelled -`default rates` to be clear it comes from a built-in table rather than your own +`$ saved` and `$ saved / request` appear with no extra configuration, labelled +`default rates` to be clear they come from a built-in table rather than your own account. +**Read that figure as a floor, not a measurement.** The built-in rates were +measured on a shared gateway that bills below vendor list. If your Claude Code +talks straight to Anthropic — which it does unless you have set +`ANTHROPIC_BASE_URL` — you are paying list, so the real saving is several times +what the column shows. Set your own rates to make it accurate; the number is +useful as-is only for confirming the plugin is working and comparing turns +against each other. + Token savings are reported per prompt-cache tier, never as one blended number: providers charge ~1.25x the input rate for a cache write and ~0.1x for a cache read, so identical saved bytes differ by more than 12x depending on cache state. From a4128d29707493d34f7b1a69ef008e4749bb912f Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 2 Sep 2026 19:52:39 -0400 Subject: [PATCH 25/28] =?UTF-8?q?fix:=20Address=20review=20pass=203=20?= =?UTF-8?q?=E2=80=94=20accounting,=20predicate=20split,=20drift=20guard?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pass-3 table lists blockers 1-6 as unchanged; they were fixed in 1a71a913 and the cited line numbers are pre-fix positions. The new findings were all unaddressed and are fixed here. NeedsBody was the one predicate left undirected, and it gated both request buffering and response buffering — so a response-only mutator still forced the request body to be buffered and a request-only mutator still forced non-SSE response buffering, the mirror image of the waste this series set out to remove. Split into NeedsRequestBody / NeedsResponseBody, with the listeners using the matching one. They read RAW capabilities, not Normalize(): the ReadsBody promotion means "you may read the body you write", which is directional, so going through the undirected field would let WritesResponseBody imply a need for the request body and undo the split. An explicitly declared ReadsBody still counts for both, because that field genuinely does not say which body — closing that needs direction-specific read capabilities, the same prerequisite as the reverse-order reader gap. noteDrift's sync.Once was consumed by an empty first manifest, because Once.Do marks itself done however the closure returns. The precondition now runs before the guard. An empty first manifest is the norm on the dialects whose tool names the plugin already documents it cannot read (Gemini functionDeclarations, Bedrock toolSpec nesting), so a stale remove list stayed silent in exactly the deployments most likely to have one. The SSE buffered-path WARN fired per response when a WritesResponseBody plugin was present. cpex and sparc are supported configurations, not misconfigurations, so that was log spam at request rate; now one Info per process. The comment beside it claimed "WritesRequestBody is already false in this branch" — reaching that else only rules out WritesResponseBody and HasStreamingResponders, so a request-only mutator can be there. Corrected, and the check now asks about the response side specifically rather than asserting what NeedsBody implies. Cost accounting: `$ saved` is gross and now says so. Changing the remove list changes the cached prefix, so the next request re-writes it at ~1.25x input while the recurring saving accrues at ~0.1x on a small delta — tens of requests to break even per change. Worse, applying a change hot-reloads the config, which rebuilds the plugin and resets the counters, so the re-warm is invisible exactly when it is paid. Documented with the break-even reasoning. The "20-25% of your prompt" figure was wrong. Measured over the 99 requests of one real session: 4-20%, median 6%. It decays with conversation length because the removed bytes are a fixed size against a growing prompt (13% early, 4% by the end), and full-manifest requests save 15-20% where reduced-manifest ones save 4-6%. I had derived 24% from a single early turn and shipped it unqualified in the README and the quickstart. Signed-off-by: Hai Huang --- README.md | 7 ++-- .../authlib/listener/forwardproxy/server.go | 23 ++++++++--- .../authlib/pipeline/bodydirection_test.go | 35 ++++++++++++++++ authbridge/authlib/pipeline/holder.go | 6 +++ authbridge/authlib/pipeline/pipeline.go | 39 +++++++++++++++++- .../authlib/plugins/toolprune/metrics.go | 16 +++++++- .../authlib/plugins/toolprune/plugin.go | 17 ++++++-- .../authlib/plugins/toolprune/plugin_test.go | 40 +++++++++++++++++++ authbridge/docs/laptop-token-savings.md | 16 +++++++- authbridge/docs/tool-prune-plugin.md | 22 ++++++++++ 10 files changed, 203 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index f3eb76d17..3c2093f66 100644 --- a/README.md +++ b/README.md @@ -40,9 +40,10 @@ Watch an AI agent's traffic — its model, tool, and agent-to-agent calls — de ## Cut Claude Code token cost on your laptop Already using Claude Code? Cortex can strip the tool definitions your agent never -calls out of every request. On the traffic this was measured against that is -20–25% of the prompt billed per turn; your share depends on how many of the -tools you actually use. Four steps, about two minutes: +calls out of every request. Measured over 99 requests in one session: **4–20% of +the prompt billed per turn, median 6%**. The share is highest early — the removed +bytes are a fixed size, so as the conversation grows they shrink as a fraction of +it — and depends on how many of the tools you actually use. Four steps, about two minutes: **[Cut Claude Code token cost](./authbridge/docs/laptop-token-savings.md)**. ## Running on Kubernetes diff --git a/authbridge/authlib/listener/forwardproxy/server.go b/authbridge/authlib/listener/forwardproxy/server.go index 40e1dde6f..99130dfcf 100644 --- a/authbridge/authlib/listener/forwardproxy/server.go +++ b/authbridge/authlib/listener/forwardproxy/server.go @@ -70,6 +70,10 @@ type Server struct { TLSBridge *tlsbridge.Engine // nil = disabled; set by caller after NewServer + // bufferedFallbackOnce keeps the SSE-buffered-path notice to one line per + // process; the condition is a supported chain shape, not an error. + bufferedFallbackOnce sync.Once + // Bridge-health counters. When the TLS bridge is enabled but the client // does not trust its CA, every HTTPS request opens a CONNECT tunnel and // nothing is ever decrypted: the pipeline sees opaque tunnels, every @@ -267,7 +271,7 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, isBridge }() } - if !skipped && s.OutboundPipeline.NeedsBody() && r.Body != nil { + if !skipped && s.OutboundPipeline.NeedsRequestBody() && r.Body != nil { r.Body = http.MaxBytesReader(w, r.Body, maxBodySize) body, err := io.ReadAll(r.Body) if err != nil { @@ -425,7 +429,12 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, isBridge // it can't stream — fall back to the buffered path with a warning. // A request-only mutator does NOT land here: it never touches // these bytes, so the relay stays incremental. - slog.Warn("forward-proxy: text/event-stream response with WritesResponseBody plugin — falling back to buffered path", "host", r.Host) + // Once, not per response: a response mutator on an SSE chain is a + // supported configuration (cpex, sparc), not a misconfiguration, so + // warning every request is log spam at request rate. + s.bufferedFallbackOnce.Do(func() { + slog.Info("forward-proxy: text/event-stream responses will use the buffered path — a WritesResponseBody plugin is in the chain", "host", r.Host) + }) } else if s.OutboundPipeline.HasStreamingResponders() { // Streaming-aware plugins (inference-parser, a2a-parser) parse // each SSE frame; handleStreamingResponse re-frames via sseframe. @@ -446,9 +455,11 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, isBridge // streamed body should implement StreamingResponder. Warn // (mirroring the WritesRequestBody fallback above) so the // misconfiguration surfaces instead of the plugin silently seeing - // no body. WritesRequestBody is already false in this branch, so - // NeedsBody() here implies ReadsBody. - if s.OutboundPipeline.NeedsBody() { + // no body. Reaching this branch only rules out WritesResponseBody and + // HasStreamingResponders — a request-only mutator can still be here — + // so ask about the response side specifically rather than asserting + // what NeedsBody implies. + if s.OutboundPipeline.NeedsResponseBody() { slog.Warn("forward-proxy: text/event-stream response with a ReadsBody plugin that is not a StreamingResponder — streaming byte-for-byte; its OnResponse will see an empty body (implement StreamingResponder to inspect a streamed body)", "host", r.Host) } s.streamPassthrough(w, r, resp, pctx) @@ -456,7 +467,7 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, isBridge } } - if s.OutboundPipeline.NeedsBody() && resp.Body != nil { + if s.OutboundPipeline.NeedsResponseBody() && resp.Body != nil { respBody, err := io.ReadAll(io.LimitReader(resp.Body, maxBodySize+1)) if err != nil { slog.Warn("forward-proxy: response body read error", "host", r.Host, "error", err) diff --git a/authbridge/authlib/pipeline/bodydirection_test.go b/authbridge/authlib/pipeline/bodydirection_test.go index 8b4874508..3ede431c0 100644 --- a/authbridge/authlib/pipeline/bodydirection_test.go +++ b/authbridge/authlib/pipeline/bodydirection_test.go @@ -186,3 +186,38 @@ func TestValidateCapabilities_ResponseAndRequestMutatorsCoexist(t *testing.T) { t.Error("two response mutators must still be rejected") } } + +// TestNeedsBody_DirectionalDoesNotCrossContaminate: the undirected NeedsBody made +// each write flag force the other direction's buffering — a response-only mutator +// had the request body buffered for nothing, and a request-only mutator had +// non-SSE responses buffered for nothing. That is the mirror image of the waste +// the directional capabilities exist to remove. +func TestNeedsBody_DirectionalDoesNotCrossContaminate(t *testing.T) { + tests := []struct { + name string + caps PluginCapabilities + wantReq, wantRsp bool + }{ + {"request-only mutator", PluginCapabilities{WritesRequestBody: true}, true, false}, + {"response-only mutator", PluginCapabilities{WritesResponseBody: true}, false, true}, + {"both", PluginCapabilities{WritesRequestBody: true, WritesResponseBody: true}, true, true}, + // ReadsBody is itself undirected, so it must still count for both. + {"pure reader", PluginCapabilities{ReadsBody: true}, true, true}, + {"neither", PluginCapabilities{}, false, false}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + p := mustBuild(t, &stubPlugin{name: "p", caps: tc.caps}) + if got := p.NeedsRequestBody(); got != tc.wantReq { + t.Errorf("NeedsRequestBody() = %v, want %v", got, tc.wantReq) + } + if got := p.NeedsResponseBody(); got != tc.wantRsp { + t.Errorf("NeedsResponseBody() = %v, want %v", got, tc.wantRsp) + } + // The aggregate stays the OR, for callers that need either. + if got := p.NeedsBody(); got != (tc.wantReq || tc.wantRsp) { + t.Errorf("NeedsBody() = %v, want %v", got, tc.wantReq || tc.wantRsp) + } + }) + } +} diff --git a/authbridge/authlib/pipeline/holder.go b/authbridge/authlib/pipeline/holder.go index 0201edec9..b27bab81e 100644 --- a/authbridge/authlib/pipeline/holder.go +++ b/authbridge/authlib/pipeline/holder.go @@ -81,6 +81,12 @@ func (h *Holder) RunFinish(ctx context.Context, pctx *Context, outcome Outcome) // that decide whether to buffer the request/response body. func (h *Holder) NeedsBody() bool { return h.p.Load().NeedsBody() } +// NeedsRequestBody is equivalent to h.Load().NeedsRequestBody(). +func (h *Holder) NeedsRequestBody() bool { return h.p.Load().NeedsRequestBody() } + +// NeedsResponseBody is equivalent to h.Load().NeedsResponseBody(). +func (h *Holder) NeedsResponseBody() bool { return h.p.Load().NeedsResponseBody() } + // WritesRequestBody is equivalent to h.Load().WritesRequestBody(). // Listeners read this when deciding whether to propagate a rewritten // request body to the wire. diff --git a/authbridge/authlib/pipeline/pipeline.go b/authbridge/authlib/pipeline/pipeline.go index ced57ec15..3f246e599 100644 --- a/authbridge/authlib/pipeline/pipeline.go +++ b/authbridge/authlib/pipeline/pipeline.go @@ -366,9 +366,44 @@ func (p *Pipeline) NotReadyPlugin() string { // NeedsBody returns true if any plugin in the pipeline needs the body // buffered — either to read it (ReadsBody) or to mutate it (WritesRequestBody). func (p *Pipeline) NeedsBody() bool { + return p.NeedsRequestBody() || p.NeedsResponseBody() +} + +// NeedsRequestBody reports whether the request body must be buffered. +// +// Split from NeedsBody because the undirected version made each write flag force +// the other direction's buffering: a response-only mutator had the request body +// buffered for nothing, and a request-only mutator had non-SSE responses +// buffered for nothing — the mirror image of the waste the directional +// capabilities exist to remove. +// +// ReadsBody still counts toward both, and deliberately: it is itself undirected +// ("reads pctx.Body and/or pctx.ResponseBody"), so a plugin that only reads +// responses cannot be distinguished from one that only reads requests. Closing +// that needs direction-specific READ capabilities — the same prerequisite as the +// reverse-order reader gap noted in validateCapabilities. +func (p *Pipeline) NeedsRequestBody() bool { for _, plugin := range p.plugins { - caps := plugin.Capabilities().Normalize() - if caps.ReadsBody || caps.WritesRequestBody || caps.WritesResponseBody { + // RAW capabilities, not Normalize(). The ReadsBody promotion means "you + // may read the body you write", which is inherently directional — so + // reading it back through the undirected ReadsBody field would let + // WritesResponseBody imply a need for the REQUEST body and undo the + // split. An explicitly declared ReadsBody still counts for both, because + // that field genuinely does not say which body. + caps := plugin.Capabilities() + if caps.ReadsBody || caps.WritesRequestBody { + return true + } + } + return false +} + +// NeedsResponseBody reports whether the response body must be buffered. See +// NeedsRequestBody for why ReadsBody counts toward both. +func (p *Pipeline) NeedsResponseBody() bool { + for _, plugin := range p.plugins { + caps := plugin.Capabilities() // raw — see NeedsRequestBody + if caps.ReadsBody || caps.WritesResponseBody { return true } } diff --git a/authbridge/authlib/plugins/toolprune/metrics.go b/authbridge/authlib/plugins/toolprune/metrics.go index 7b65d62ec..a24ba73ef 100644 --- a/authbridge/authlib/plugins/toolprune/metrics.go +++ b/authbridge/authlib/plugins/toolprune/metrics.go @@ -199,13 +199,25 @@ func (m *metrics) snapshot() []pipeline.Metric { costNote = note + "; " + costNote } } - out = append(out, pipeline.Metric{Name: "$ saved", Value: m.usdSaved, Unit: "usd", Note: costNote}) + // GROSS, not net. Changing the remove list changes the cached prefix, so + // the next request re-writes the whole prefix at the cache-write rate + // (~1.25x input) while the recurring saving is at the cache-read rate + // (~0.1x) on a small delta — tens of requests to break even after each + // change. Counters also reset on the reload that applies the change, so + // the re-warm is invisible exactly when it is paid. Say so on the row + // rather than presenting a gross figure as a net one. + grossNote := costNote + if grossNote != "" { + grossNote += "; " + } + grossNote += "gross — excludes cache re-warm after a remove-list change" + out = append(out, pipeline.Metric{Name: "$ saved", Value: m.usdSaved, Unit: "usd", Note: grossNote}) if priced := m.requestsCosted - m.unpriced; priced > 0 { out = append(out, pipeline.Metric{ Name: "$ saved / request", Value: m.usdSaved / float64(priced), Unit: "usd", - Note: costNote, + Note: grossNote, }) } } diff --git a/authbridge/authlib/plugins/toolprune/plugin.go b/authbridge/authlib/plugins/toolprune/plugin.go index 89f6df504..ff1a37d0f 100644 --- a/authbridge/authlib/plugins/toolprune/plugin.go +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -171,6 +171,9 @@ type ToolPrune struct { m metrics driftOnce sync.Once + // driftChecked records that the stale-list check actually ran, so a test can + // tell "guard consumed" from "guard consumed without checking anything". + driftChecked bool } func New() *ToolPrune { return &ToolPrune{} } @@ -631,10 +634,18 @@ func tierOf(inf *pipeline.InferenceExtension) tier { // plugin actually sees. A stale list costs savings rather than correctness, so // it surfaces as a warning instead of a failure. func (p *ToolPrune) noteDrift(observed []pipeline.InferenceTool) { + // Check the precondition BEFORE consuming the Once. sync.Once marks itself + // done however the closure returns, so an early return on an empty manifest + // used to disable this warning permanently — and an empty first manifest is + // the norm on the dialects the plugin already knows it cannot read names + // from (Gemini functionDeclarations, Bedrock toolSpec nesting), which is a + // live path here. The result was that a stale remove list stayed silent in + // exactly the deployments most likely to have one. + if len(observed) == 0 { + return + } p.driftOnce.Do(func() { - if len(observed) == 0 { - return - } + p.driftChecked = true present := make(map[string]struct{}, len(observed)) for _, t := range observed { present[t.Name] = struct{}{} diff --git a/authbridge/authlib/plugins/toolprune/plugin_test.go b/authbridge/authlib/plugins/toolprune/plugin_test.go index 6d4995676..786cdcf13 100644 --- a/authbridge/authlib/plugins/toolprune/plugin_test.go +++ b/authbridge/authlib/plugins/toolprune/plugin_test.go @@ -968,3 +968,43 @@ func TestPrune_DoesNotDuplicateCacheControl(t *testing.T) { t.Errorf("cache_control appears %d times, want 1: %s", n, pctx.Body) } } + +// TestNoteDrift_EmptyFirstManifestDoesNotSpendTheOnce: sync.Once marks itself +// done however the closure returns, so an early return on an empty manifest used +// to disable the stale-list warning permanently — and an empty first manifest is +// the norm on dialects whose tool names the plugin cannot read (Gemini +// functionDeclarations, Bedrock toolSpec nesting). A stale remove list then stayed +// silent in exactly the deployments most likely to have one. +func TestNoteDrift_EmptyFirstManifestDoesNotSpendTheOnce(t *testing.T) { + p := configured(t, "NeverOffered") + + // Nothing observed: must not consume the guard. + p.noteDrift(nil) + p.noteDrift([]pipeline.InferenceTool{}) + if p.driftChecked { + t.Fatal("the check claims to have run on an empty manifest") + } + + // A real manifest afterwards must still reach the check. + p.noteDrift([]pipeline.InferenceTool{{Name: "Read"}}) + if !p.driftChecked { + t.Error("the check never ran on the first non-empty manifest — an empty one had spent the Once") + } +} + +// TestMetrics_DollarRowsDiscloseTheyAreGross: changing the remove list changes the +// cached prefix, so the next request re-writes the whole prefix at ~1.25x input +// while the recurring saving is ~0.1x on a small delta — tens of requests to break +// even after each change. Counters also reset on the reload that applies the +// change, so the re-warm is invisible exactly when it is paid. A figure that does +// not say it is gross reads as net. +func TestMetrics_DollarRowsDiscloseTheyAreGross(t *testing.T) { + p := configured(t, "NotebookEdit") + pruneWithModel(t, p, "claude-opus-5") + for _, name := range []string{"$ saved", "$ saved / request"} { + m := findMetric(t, p.Metrics(), name) + if !strings.Contains(m.Note, "gross") || !strings.Contains(m.Note, "re-warm") { + t.Errorf("%s note = %q, want it to disclose the figure is gross of cache re-warm", name, m.Note) + } + } +} diff --git a/authbridge/docs/laptop-token-savings.md b/authbridge/docs/laptop-token-savings.md index 429c89e71..37b4d5867 100644 --- a/authbridge/docs/laptop-token-savings.md +++ b/authbridge/docs/laptop-token-savings.md @@ -85,8 +85,20 @@ Then watch what it saved: abctl --endpoint http://localhost:47601 ``` -Plugin pane → `tool-prune` → `Metrics`. Expect `bytes removed / request` around -25–30 KB and a token saving of roughly 20–25% of your prompt. +Plugin pane → `tool-prune` → `Metrics`. + +What to expect, measured over 99 requests of one real session: **4–20% of the +prompt per turn, median 6%**. Two things move it, and neither is a defect: + +- **How much of the manifest is yours to prune.** Requests carrying the full tool + set saved 15–20%; most requests in that session offered a reduced set and saved + 4–6%. +- **How far into the conversation you are.** The removed bytes are a fixed size, + so their share of a growing prompt falls — 13% early in that session, 4% by the + end. + +A single early turn can read ~24%, which is why a figure quoted from one request +is not the number to plan with. Stop it with `pkill -f 'authbridge-proxy --config'`. diff --git a/authbridge/docs/tool-prune-plugin.md b/authbridge/docs/tool-prune-plugin.md index 64fd46835..299963b3f 100644 --- a/authbridge/docs/tool-prune-plugin.md +++ b/authbridge/docs/tool-prune-plugin.md @@ -140,6 +140,28 @@ The bytes-to-tokens ratio is calibrated on your own traffic — prompt tokens ov request bytes for the same request, both post-pruning so the two sides agree — rather than bundling a tokenizer or assuming a constant. +### The figure is gross, not net + +Changing the `remove` list changes the cached prompt prefix, so the first request +after a change re-writes the whole prefix at the cache-**write** rate (~1.25x +input) while the recurring saving accrues at the cache-**read** rate (~0.1x) on a +small delta. Order of magnitude: re-warming a ~30k-token prefix costs on the order +of tens of thousands of input-equivalents against a few hundred saved per +subsequent cache-read request — **tens of requests to break even after each list +change.** + +Two consequences worth being blunt about: + +- **`$ saved` is gross.** It counts what the removed definitions would have cost + and subtracts nothing for the re-warm. The row says so. +- **The re-warm is invisible exactly when it is paid.** Applying a list change + hot-reloads the config, which rebuilds the plugin and resets its counters — so + the run that incurs the cost starts from zero. + +Practical reading: change the list rarely, and treat a figure gathered over a few +requests immediately after a change as optimistic. Over a long steady session the +gross figure converges on the net one, because the re-warm is paid once. + ### Costing it **Dollars work out of the box.** The plugin ships a rate table measured from the From c4fd112d61cd699fb4393e90c24dd5eef35cd01b Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Thu, 3 Sep 2026 09:57:32 -0400 Subject: [PATCH 26/28] feat: Key tool-prune pricing by model family, not version MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The built-in rate table held seven exact model keys, so every provider version bump — opus 4.6, 4.7, 4.8, 5 — silently dropped that traffic to "unpriced" until someone edited Go and rebuilt the image. That is not a maintenance burden an operator can be asked to carry, and the failure is quiet: the readout just stops showing dollars. Pricing keys are now globs, and the built-ins are keyed by family: "*claude-opus-*" / "*claude-sonnet-*" / "*claude-haiku-*". Three patterns replace the seven exact entries and cover every model on the gateway's allowlist, including provider prefixes (aws/claude-opus-5), dated suffixes (claude-haiku-4-5-20251001), and versions not yet released. Config keys may be globs too. Resolution is ordered so the more specific statement wins: exact config key, then longest matching config glob, then built-in pattern, then the flat fallback, then unpriced. Exact-before-glob is what keeps the escape hatch open — if a future version ever bills differently from its family, pin that one model and the family default keeps serving the rest. Two details worth naming: - Globs compile with no separator, unlike the host globs elsewhere in authlib. Model names are delimited by "-" and "/", so "*" has to span both; hostglob's "."-delimited semantics would not match aws/ prefixes. - Glob candidates are sorted longest-pattern-first, so overlapping patterns resolve deterministically rather than by map iteration order. An invalid pattern now fails Configure with the offending key named. A typo'd glob and a genuinely unknown model would otherwise both surface as the same silent "unpriced" row. Tests pin the real model list off the gateway allowlist, the precedence ladder, case-insensitivity on both sides, and the bad-glob rejection. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- .../authlib/plugins/toolprune/plugin.go | 38 ++++-- .../authlib/plugins/toolprune/plugin_test.go | 107 ++++++++++++++++ .../authlib/plugins/toolprune/pricing.go | 119 ++++++++++++------ authbridge/docs/plugin-catalog.md | 2 +- authbridge/docs/tool-prune-plugin.md | 37 +++++- 5 files changed, 250 insertions(+), 53 deletions(-) diff --git a/authbridge/authlib/plugins/toolprune/plugin.go b/authbridge/authlib/plugins/toolprune/plugin.go index ff1a37d0f..099e3a2e9 100644 --- a/authbridge/authlib/plugins/toolprune/plugin.go +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -64,6 +64,11 @@ type config struct { // pricing is Pricing with keys lower-cased; built by applyDefaults. pricing map[string]modelRates `json:"-"` + // pricingGlobs are the Pricing keys containing glob metacharacters, + // compiled in match order — the mechanism that makes a version bump need + // no edit anywhere. + pricingGlobs []patternRates `json:"-"` + pricingGlobErr error `json:"-"` // The flat fields are the fallback for models absent from Pricing. Names and // semantics match litellm-budget-track. All optional; with nothing set no @@ -126,16 +131,22 @@ const ( // per model without deleting anything. func (c *config) ratesFor(model string) (modelRates, rateSource) { key := strings.ToLower(model) + // Exact config key first, so an operator can pin one version even when a + // broader pattern would also match it. if r, ok := c.pricing[key]; ok && r.set() { return r, rateConfigured } - // The built-in table comes BEFORE the flat fallback. The flat fields are - // documented as covering "models absent from pricing", and a model in the - // built-in table is not absent — letting one flat input rate shadow every - // per-model default would reintroduce exactly the flat-rate mispricing the - // per-model table exists to avoid, and silently, since the figure would then - // claim to be configured. - if r, ok := defaultPricing[key]; ok { + // Then a config glob. This is what lets a model version bump need no edit at + // all: one "*claude-opus-*" entry covers every opus release. + if r, ok := lookupPattern(c.pricingGlobs, key); ok { + return r, rateConfigured + } + // Then the built-in family patterns — before the flat fallback, because the + // flat fields are documented as covering models "absent from pricing", and a + // model the built-in table knows is not absent. Letting one flat rate shadow + // every family default would reintroduce flat-rate mispricing, silently, and + // the figure would claim to be operator-configured. + if r, ok := lookupPattern(defaultPatterns, key); ok { return r, rateDefault } fallback := modelRates{ @@ -157,9 +168,17 @@ func (c *config) applyDefaults() { // without allocating per request. Gateways vary in how they echo model // names, and a case mismatch would silently unprice the traffic. c.pricing = make(map[string]modelRates, len(c.Pricing)) + globs := map[string]modelRates{} for k, v := range c.Pricing { - c.pricing[strings.ToLower(k)] = v + lk := strings.ToLower(k) + if strings.ContainsAny(lk, "*?[") { + globs[lk] = v + continue + } + c.pricing[lk] = v } + // A bad glob is surfaced by validate(), not silently dropped. + c.pricingGlobs, c.pricingGlobErr = compilePatterns(globs) } // ToolPrune is the plugin. Counters live in metrics, guarded by its own mutex; @@ -213,6 +232,9 @@ func (p *ToolPrune) Configure(raw json.RawMessage) error { } } c.applyDefaults() + if c.pricingGlobErr != nil { + return fmt.Errorf("tool-prune config: invalid pricing pattern: %w", c.pricingGlobErr) + } p.cfg = c p.raw = raw diff --git a/authbridge/authlib/plugins/toolprune/plugin_test.go b/authbridge/authlib/plugins/toolprune/plugin_test.go index 786cdcf13..e1aba7dfa 100644 --- a/authbridge/authlib/plugins/toolprune/plugin_test.go +++ b/authbridge/authlib/plugins/toolprune/plugin_test.go @@ -1008,3 +1008,110 @@ func TestMetrics_DollarRowsDiscloseTheyAreGross(t *testing.T) { } } } + +// TestPricingPatternsCoverRealModels pins the pattern keys against the actual +// model list the rossoctl LiteLLM gateway serves, including provider prefixes +// and dated suffixes. The point of this test is the regression it prevents: a +// provider version bump must not silently drop a family to unpriced. +func TestPricingPatternsCoverRealModels(t *testing.T) { + cases := []struct { + model string + want float64 // input rate + }{ + // opus family, across versions and prefixes + {"claude-opus-5", 0.0000038}, + {"claude-opus-4-8", 0.0000038}, + {"claude-opus-4-7", 0.0000038}, + {"claude-opus-4-6", 0.0000038}, + {"aws/claude-opus-5", 0.0000038}, + {"aws/claude-opus-4-7", 0.0000038}, + // a version that does not exist yet must still price + {"claude-opus-9", 0.0000038}, + {"claude-opus-5-20260901", 0.0000038}, + // sonnet + {"claude-sonnet-5", 0.00000152}, + {"claude-sonnet-4-6", 0.00000152}, + {"aws/claude-sonnet-4-5", 0.00000152}, + {"claude-sonnet-4-5-20250929", 0.00000152}, + // haiku + {"claude-haiku-4-5", 0.00000076}, + {"aws/claude-haiku-4-5", 0.00000076}, + {"claude-haiku-4-5-20251001", 0.00000076}, + } + c := &config{} + c.applyDefaults() + for _, tc := range cases { + rates, src := c.ratesFor(tc.model) + if src != rateDefault { + t.Errorf("%s: source = %v, want rateDefault", tc.model, src) + continue + } + if rates.InputCostPerToken != tc.want { + t.Errorf("%s: input rate = %g, want %g", tc.model, rates.InputCostPerToken, tc.want) + } + } + // A non-Claude model has no built-in rate and must report so rather than + // borrowing a Claude family's numbers. + if _, src := c.ratesFor("gpt-4o"); src != rateNone { + t.Errorf("gpt-4o: source = %v, want rateNone", src) + } +} + +// TestPricingPatternPrecedence covers the resolution order that lets an operator +// pin one version without giving up family coverage. +func TestPricingPatternPrecedence(t *testing.T) { + c := &config{Pricing: map[string]modelRates{ + // exact key: must win over both globs below + "claude-opus-5": {InputCostPerToken: 1}, + // broad glob + "*claude-opus-*": {InputCostPerToken: 2}, + // narrower glob: longer pattern wins among globs + "*claude-opus-4-8*": {InputCostPerToken: 3}, + }} + c.applyDefaults() + if c.pricingGlobErr != nil { + t.Fatalf("compile: %v", c.pricingGlobErr) + } + for _, tc := range []struct { + model string + want float64 + src rateSource + }{ + {"claude-opus-5", 1, rateConfigured}, // exact beats glob + {"claude-opus-4-8", 3, rateConfigured}, // longest glob wins + {"claude-opus-4-6", 2, rateConfigured}, // broad glob + // built-in pattern still covers a family the operator said nothing about + {"claude-haiku-4-5", 0.00000076, rateDefault}, + } { + rates, src := c.ratesFor(tc.model) + if src != tc.src || rates.InputCostPerToken != tc.want { + t.Errorf("%s: got (%g, %v), want (%g, %v)", + tc.model, rates.InputCostPerToken, src, tc.want, tc.src) + } + } +} + +// TestPricingPatternMatchesCase guards the lowercasing on both sides: config +// keys and the model name off the wire. +func TestPricingPatternMatchesCase(t *testing.T) { + c := &config{Pricing: map[string]modelRates{ + "*CLAUDE-OPUS-*": {InputCostPerToken: 7}, + }} + c.applyDefaults() + rates, src := c.ratesFor("AWS/Claude-Opus-5") + if src != rateConfigured || rates.InputCostPerToken != 7 { + t.Errorf("got (%g, %v), want (7, configured)", rates.InputCostPerToken, src) + } +} + +// TestPricingBadPatternRejected: a malformed glob must fail Configure loudly, +// not degrade to unpriced with no explanation. +func TestPricingBadPatternRejected(t *testing.T) { + c := &config{Pricing: map[string]modelRates{ + "*claude-[opus": {InputCostPerToken: 1}, + }} + c.applyDefaults() + if c.pricingGlobErr == nil { + t.Fatal("want a compile error for an unterminated character class") + } +} diff --git a/authbridge/authlib/plugins/toolprune/pricing.go b/authbridge/authlib/plugins/toolprune/pricing.go index 65f9da757..bd906fd31 100644 --- a/authbridge/authlib/plugins/toolprune/pricing.go +++ b/authbridge/authlib/plugins/toolprune/pricing.go @@ -1,65 +1,104 @@ package toolprune -// defaultPricing holds per-token rates for the models seen on the rossoctl -// LiteLLM gateway, measured from its own x-litellm-response-cost headers: -// send two non-streaming requests of differing prompt length and difference -// them, rate = Δcost / Δinput_tokens; the cache tiers were obtained the same -// way with a cache_control block sent twice. +import ( + "sort" + + "github.com/gobwas/glob" +) + +// patternRates is one pricing entry whose key may be a glob. +type patternRates struct { + pattern string + glob glob.Glob + rates modelRates +} + +// defaultPatterns holds per-token rates for the Claude families seen on the +// rossoctl LiteLLM gateway, measured from its own x-litellm-response-cost +// headers: send two non-streaming requests of differing prompt length and +// difference them, rate = Δcost / Δinput_tokens; the cache tiers were obtained +// the same way with a cache_control block sent twice. // -// These exist so `$ saved` works with no configuration, which is the difference -// between a number an operator sees and one they never get around to enabling. -// They are a starting point, not a fact about your account: +// Keyed by FAMILY, not by version. Model names churn — opus 4.6, 4.7, 4.8, 5 — +// and a table of exact versions means a code change and a rebuild every time a +// provider ships one, which is not a thing an operator can be asked to do. One +// pattern per family survives version bumps, and matches provider prefixes +// (aws/, azure/) and dated suffixes (-20251001) alike. +// +// The tradeoff is stated plainly: this assumes a family bills at one rate. That +// has held across the Claude versions measured, but if a version ever differs, +// pin it with an exact `pricing:` key in config — exact always beats a pattern. +// +// These exist so `$ saved` works with no configuration. They are a starting +// point, not a fact about your account: // // - Rates are gateway-specific. This gateway bills well below vendor list, so // a deployment talking straight to Anthropic pays more and these // UNDERSTATE its saving — by roughly 4x on the input tier at the time of -// measurement. That is the common case for a laptop install, so the -// caveat travels with every figure rather than living only here. +// measurement. That is the common case for a laptop install, so the caveat +// travels with every figure rather than living only here. // - Rates change. Nothing here refreshes them. // -// A figure derived from these is therefore labelled as coming from default -// rates wherever it is reported, so it cannot be mistaken for one measured on -// the operator's own account. Any `pricing` entry in config overrides the -// matching model outright. -// -// Keys are lower-case; lookup folds the observed model name the same way. -var defaultPricing = map[string]modelRates{ +// Any `pricing` entry in config overrides the matching model. +var defaultPatterns = mustCompilePatterns(map[string]modelRates{ // input 1.00x / cache write 1.25x / cache read 0.10x - "claude-opus-5": { - InputCostPerToken: 0.0000038, - CacheWriteCostPerToken: 0.00000475, - CacheReadCostPerToken: 0.00000038, - }, - "aws/claude-opus-5": { + "*claude-opus-*": { InputCostPerToken: 0.0000038, CacheWriteCostPerToken: 0.00000475, CacheReadCostPerToken: 0.00000038, }, - "claude-sonnet-5": { - InputCostPerToken: 0.00000152, - CacheWriteCostPerToken: 0.0000019, - CacheReadCostPerToken: 0.000000152, - }, - "claude-haiku-4-5": { - InputCostPerToken: 0.00000076, - CacheWriteCostPerToken: 0.00000095, - CacheReadCostPerToken: 0.000000076, - }, - "aws/claude-sonnet-5": { + "*claude-sonnet-*": { InputCostPerToken: 0.00000152, CacheWriteCostPerToken: 0.0000019, CacheReadCostPerToken: 0.000000152, }, - "aws/claude-haiku-4-5": { - InputCostPerToken: 0.00000076, - CacheWriteCostPerToken: 0.00000095, - CacheReadCostPerToken: 0.000000076, - }, - "claude-haiku-4-5-20251001": { + "*claude-haiku-*": { InputCostPerToken: 0.00000076, CacheWriteCostPerToken: 0.00000095, CacheReadCostPerToken: 0.000000076, }, +}) + +// compilePatterns compiles glob keys into match order. No separator is passed to +// glob.Compile: model names are delimited by "-" and "/", so "*" must span both +// — unlike the host globs elsewhere in authlib, which are "."-delimited. +// +// Sorted longest-pattern-first so the most specific glob wins deterministically +// when two match: "*claude-opus-4-8*" beats "*claude-opus-*". +func compilePatterns(in map[string]modelRates) ([]patternRates, error) { + out := make([]patternRates, 0, len(in)) + for pat, r := range in { + g, err := glob.Compile(pat) + if err != nil { + return nil, err + } + out = append(out, patternRates{pattern: pat, glob: g, rates: r}) + } + sort.Slice(out, func(i, j int) bool { + if len(out[i].pattern) != len(out[j].pattern) { + return len(out[i].pattern) > len(out[j].pattern) + } + return out[i].pattern < out[j].pattern + }) + return out, nil +} + +func mustCompilePatterns(in map[string]modelRates) []patternRates { + out, err := compilePatterns(in) + if err != nil { + panic("toolprune: invalid built-in pricing pattern: " + err.Error()) + } + return out +} + +// lookupPattern returns the rates for the first pattern matching model. +func lookupPattern(pats []patternRates, model string) (modelRates, bool) { + for _, p := range pats { + if p.glob.Match(model) && p.rates.set() { + return p.rates, true + } + } + return modelRates{}, false } func (s rateSource) String() string { diff --git a/authbridge/docs/plugin-catalog.md b/authbridge/docs/plugin-catalog.md index 62c5ee8db..534c36c10 100644 --- a/authbridge/docs/plugin-catalog.md +++ b/authbridge/docs/plugin-catalog.md @@ -240,7 +240,7 @@ body-reading plugin (it rewrites the request body). Declares - `remove` (`[]string`) — tool names to delete from the manifest. The complete verdict: no learning, no state, no storage. Names absent from a given request are ignored. - `paths` (`[]string`) — request paths to act on, matched exactly or by suffix. Defaults to `/v1/chat/completions`, `/v1/completions`, `/v1/messages`. -- `pricing` (`map[model]rates`) — per-token rates keyed by model name, each with `input_cost_per_token`, `cache_write_cost_per_token`, `cache_read_cost_per_token`. **Optional**: a built-in table covers the Claude models on the rossoctl gateway, so `$ saved` works unconfigured; any entry here overrides the built-in for that model. Per model because rates differ ~5x across opus/sonnet/haiku. Keys matched case-insensitively. +- `pricing` (`map[model]rates`) — per-token rates keyed by model name **or glob** (`*claude-opus-*`), each with `input_cost_per_token`, `cache_write_cost_per_token`, `cache_read_cost_per_token`. **Optional**: built-in patterns cover the Claude families on the rossoctl gateway, so `$ saved` works unconfigured; any entry here overrides the built-in. Per model because rates differ ~5x across opus/sonnet/haiku. Built-ins are keyed by *family*, not version, so an opus 4.8 → 5 rename needs no code change. Resolution: exact key → longest matching glob → built-in pattern → flat fallback → unpriced; keys matched case-insensitively. An invalid glob fails startup with the key named. - `input_cost_per_token`, `cache_write_cost_per_token`, `cache_read_cost_per_token` (`float`) — optional flat fallback for models absent from `pricing`. A figure from built-in rates is labelled as such; a model in neither the table nor config is counted in a `requests unpriced` row instead of charged at another model's rate. No output rate: pruning only shrinks the prompt. Generate the list from local transcripts with `abctl tools scan`, which diff --git a/authbridge/docs/tool-prune-plugin.md b/authbridge/docs/tool-prune-plugin.md index 299963b3f..0b0c09407 100644 --- a/authbridge/docs/tool-prune-plugin.md +++ b/authbridge/docs/tool-prune-plugin.md @@ -167,11 +167,11 @@ gross figure converges on the net one, because the re-warm is paid once. **Dollars work out of the box.** The plugin ships a rate table measured from the rossoctl LiteLLM gateway, so `$ saved` appears with no configuration: -| model | input | cache write (1.25x) | cache read (0.10x) | +| pattern | input | cache write (1.25x) | cache read (0.10x) | |---|---|---|---| -| `claude-opus-5`, `aws/claude-opus-5` | $3.80/Mtok | $4.75/Mtok | $0.38/Mtok | -| `aws/claude-sonnet-5` | $1.52/Mtok | $1.90/Mtok | $0.152/Mtok | -| `aws/claude-haiku-4-5`, `claude-haiku-4-5-20251001` | $0.76/Mtok | $0.95/Mtok | $0.076/Mtok | +| `*claude-opus-*` | $3.80/Mtok | $4.75/Mtok | $0.38/Mtok | +| `*claude-sonnet-*` | $1.52/Mtok | $1.90/Mtok | $0.152/Mtok | +| `*claude-haiku-*` | $0.76/Mtok | $0.95/Mtok | $0.076/Mtok | Rates are keyed **per model** because they differ far more than the tiers do — 5x across this family — so a single flat rate would misprice the saving by that @@ -179,6 +179,20 @@ factor depending on which model served the request. Each request is priced at it own model's rate and the dollars accumulated, never a blended token total multiplied by one number. +Keys are **globs, and the built-ins are keyed by family rather than by version**, +which is what stops a model rename from becoming a code change. Model names churn +— opus 4.6, 4.7, 4.8, 5 — and a table of exact versions would go stale on every +release and need a rebuild to fix, which is not something an operator can be +asked to do. One pattern per family absorbs the churn and also covers provider +prefixes (`aws/claude-opus-5`) and dated suffixes +(`claude-haiku-4-5-20251001`) without separate entries. + +The tradeoff, stated plainly: this assumes a family bills at one rate. That has +held across the Claude versions measured. If a future version differs, pin it — +an exact key always beats a pattern, so `claude-opus-6:` overrides +`*claude-opus-*` for that one model and leaves the family default doing its job +for the rest. + Any figure derived from these carries `default rates — set pricing. to use yours` in its note, because they are a starting point rather than a fact about your account: they are specific to that gateway (which bills below vendor list), @@ -206,6 +220,21 @@ Model keys match what the parser records (`Extensions.Inference.Model`) and are matched case-insensitively, since gateways vary in how they echo the name and a case mismatch would silently unprice the traffic. +Config keys may be globs too (`*`, `?`, `[...]` — `gobwas/glob` with no separator, +so `*` spans the `-` and `/` in a model name). Resolution is deliberately ordered +so the more specific statement wins: + +1. exact key in your `pricing` +2. glob in your `pricing` — **longest pattern first**, so `*claude-opus-4-8*` + beats `*claude-opus-*` deterministically rather than by map iteration luck +3. built-in family pattern +4. the flat `input_cost_per_token` fallback +5. unpriced + +An invalid pattern fails startup with the offending key named, rather than +silently dropping to unpriced — a typo'd glob and a genuinely unknown model +should not look the same in the readout. + A model with no entry and no fallback is **counted, not guessed**: the readout grows a `requests unpriced` row naming the models, so an incomplete table shows as a visible gap rather than a quietly understated total. Tokens are still From 7d63fbad4e801fd7df9ce00a587f386ba917433a Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Thu, 3 Sep 2026 10:09:48 -0400 Subject: [PATCH 27/28] feat: Accept tool-prune pricing per million tokens MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rates had to be given per token — 0.0000038 — while every provider publishes per million tokens. That put a division in the operator's head for no reason, and six leading zeros is a bad place to ask for accuracy: 0.000038 is a plausible-looking typo that misprices by 10x, in either direction, with nothing in the readout to reveal it. Pricing now accepts input_cost_per_million / cache_write_cost_per_million / cache_read_cost_per_million, on both the per-model entries and the flat fallback. "3.80" copied off a price list is now a valid config value. The per-token names remain accepted, for parity with litellm-budget-track and because LiteLLM's own model_prices_and_context_window.json is per-token — rates get copied out of it verbatim, and breaking that would trade one papercut for another. Different tiers may use different units. Setting both units for the same tier is a startup error rather than a precedence question. They differ by 10^6, so silently honouring one would either overstate a saving a millionfold or bury it under rounding, and the readout gives no clue which happened. The error names the entry and the tier, so an operator with three tiers configured doesn't bisect. The built-in table is now written per-Mtok too, which is the point of the exercise: "3.80 / tokensPerMillion" can be checked against a price list at a glance where 0.0000038 cannot. Those divisions are CONSTANT expressions so the compiler folds them exactly — a runtime division lands a ulp low (3.7999999999999996e-06) and would leave the table disagreeing with the documented $3.80/Mtok in the last digit. A test pins each built-in against the documented figure as a constant expression, and fails if anyone reintroduces a runtime conversion; I mutation-checked that it does. Config-supplied per-million values necessarily divide at runtime, so those tests compare with a tolerance — a 1e-16 relative difference on a dollar figure is not a property worth pinning. Also sorts the pricing keys while normalizing, so when several entries are malformed the reported one is stable across restarts instead of depending on map iteration order. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- .../authlib/plugins/toolprune/plugin.go | 124 +++++++++-- .../authlib/plugins/toolprune/plugin_test.go | 201 +++++++++++++++++- .../authlib/plugins/toolprune/pricing.go | 23 +- authbridge/docs/plugin-catalog.md | 4 +- authbridge/docs/tool-prune-plugin.md | 29 ++- 5 files changed, 339 insertions(+), 42 deletions(-) diff --git a/authbridge/authlib/plugins/toolprune/plugin.go b/authbridge/authlib/plugins/toolprune/plugin.go index 099e3a2e9..1e0f0fd80 100644 --- a/authbridge/authlib/plugins/toolprune/plugin.go +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -30,6 +30,7 @@ import ( "encoding/json" "fmt" "log/slog" + "sort" "strings" "sync" @@ -60,15 +61,21 @@ type config struct { // that factor depending on which model served the request. // Keys match the model name the parser records // (pctx.Extensions.Inference.Model), matched case-insensitively. - Pricing map[string]modelRates `json:"pricing" description:"Per-token rates keyed by model name."` + Pricing map[string]modelRates `json:"pricing" description:"Rates keyed by model name or glob; prefer the per-million fields."` // pricing is Pricing with keys lower-cased; built by applyDefaults. pricing map[string]modelRates `json:"-"` // pricingGlobs are the Pricing keys containing glob metacharacters, // compiled in match order — the mechanism that makes a version bump need // no edit anywhere. - pricingGlobs []patternRates `json:"-"` - pricingGlobErr error `json:"-"` + pricingGlobs []patternRates `json:"-"` + // flat is the normalized flat fallback, so ratesFor doesn't rebuild it per + // request and the unit folding happens exactly once. + flat modelRates `json:"-"` + // pricingErr carries any pricing config fault — bad glob, or both units set + // for one tier — for Configure to reject. Faults are not dropped: an unpriced + // row from a typo and one from a genuinely unknown model must not look alike. + pricingErr error `json:"-"` // The flat fields are the fallback for models absent from Pricing. Names and // semantics match litellm-budget-track. All optional; with nothing set no @@ -76,18 +83,71 @@ type config struct { // // There is deliberately no output rate: pruning only ever shrinks the // prompt, so attributing output cost to it would be false. - InputCostPerToken float64 `json:"input_cost_per_token" description:"Fallback USD per uncached input token, for models absent from pricing."` - CacheWriteCostPerToken float64 `json:"cache_write_cost_per_token" description:"Fallback USD per cache-write token; defaults to input_cost_per_token."` - CacheReadCostPerToken float64 `json:"cache_read_cost_per_token" description:"Fallback USD per cache-read token; defaults to input_cost_per_token."` + InputCostPerMillion float64 `json:"input_cost_per_million" description:"Fallback USD per million uncached input tokens, for models absent from pricing."` + CacheWriteCostPerMillion float64 `json:"cache_write_cost_per_million" description:"Fallback USD per million cache-write tokens; defaults to input_cost_per_million."` + CacheReadCostPerMillion float64 `json:"cache_read_cost_per_million" description:"Fallback USD per million cache-read tokens; defaults to input_cost_per_million."` + + InputCostPerToken float64 `json:"input_cost_per_token" description:"Fallback USD per uncached input token. Alternative to input_cost_per_million; set one, not both."` + CacheWriteCostPerToken float64 `json:"cache_write_cost_per_token" description:"Fallback USD per cache-write token; defaults to input rate."` + CacheReadCostPerToken float64 `json:"cache_read_cost_per_token" description:"Fallback USD per cache-read token; defaults to input rate."` } // modelRates is one model's prompt-tier pricing. Cache rates fall back to the // input rate, matching litellm-budget-track — though on Anthropic-family models // that fallback is poor (a real cache read is 0.1x input), so set them when known. type modelRates struct { - InputCostPerToken float64 `json:"input_cost_per_token" description:"USD per uncached input token."` - CacheWriteCostPerToken float64 `json:"cache_write_cost_per_token" description:"USD per cache-write token; defaults to input_cost_per_token."` - CacheReadCostPerToken float64 `json:"cache_read_cost_per_token" description:"USD per cache-read token; defaults to input_cost_per_token."` + // The per-million fields are the ones to reach for. Every provider publishes + // prices per million tokens ("$3.80 / Mtok"), so this is the unit an operator + // already has in hand — no dividing by a million by hand, and no + // 0.0000038-vs-0.000038 typo that misprices by 10x and looks plausible in + // either direction. + InputCostPerMillion float64 `json:"input_cost_per_million" description:"USD per million uncached input tokens (the unit providers publish)."` + CacheWriteCostPerMillion float64 `json:"cache_write_cost_per_million" description:"USD per million cache-write tokens; defaults to input_cost_per_million."` + CacheReadCostPerMillion float64 `json:"cache_read_cost_per_million" description:"USD per million cache-read tokens; defaults to input_cost_per_million."` + + // The per-token fields remain accepted, for parity with + // litellm-budget-track's config and with LiteLLM's own + // model_prices_and_context_window.json — both are per-token, and rates get + // copied straight out of them. Setting both units for one tier is an error, + // not a precedence question: see normalize. + InputCostPerToken float64 `json:"input_cost_per_token" description:"USD per uncached input token. Alternative to input_cost_per_million; set one, not both."` + CacheWriteCostPerToken float64 `json:"cache_write_cost_per_token" description:"USD per cache-write token; defaults to input rate."` + CacheReadCostPerToken float64 `json:"cache_read_cost_per_token" description:"USD per cache-read token; defaults to input rate."` +} + +// tokensPerMillion converts the published unit to the per-token one all the +// downstream arithmetic uses. +const tokensPerMillion = 1_000_000 + +// normalize folds the per-million fields into the per-token ones, so everything +// after Configure deals in a single unit. +// +// Setting both units for the same tier is rejected rather than resolved by +// precedence. The two differ by 10^6, so picking a winner silently would either +// overstate a saving by a millionfold or bury it below rounding — and the +// readout gives an operator no way to tell which unit was honoured. A startup +// error naming the tier is the only outcome that can't be misread. what +// identifies the offending entry ("pricing[\"claude-opus-5\"]"). +func (r modelRates) normalize(what string) (modelRates, error) { + for _, f := range []struct { + name string + million float64 + token *float64 + }{ + {"input", r.InputCostPerMillion, &r.InputCostPerToken}, + {"cache_write", r.CacheWriteCostPerMillion, &r.CacheWriteCostPerToken}, + {"cache_read", r.CacheReadCostPerMillion, &r.CacheReadCostPerToken}, + } { + if f.million <= 0 { + continue + } + if *f.token > 0 { + return r, fmt.Errorf("%s: %s rate set as both %s_cost_per_million and %s_cost_per_token; set one", + what, f.name, f.name, f.name) + } + *f.token = f.million / tokensPerMillion + } + return r, nil } // rateFor returns the rate for a tier and whether one is actually available. @@ -149,13 +209,8 @@ func (c *config) ratesFor(model string) (modelRates, rateSource) { if r, ok := lookupPattern(defaultPatterns, key); ok { return r, rateDefault } - fallback := modelRates{ - InputCostPerToken: c.InputCostPerToken, - CacheWriteCostPerToken: c.CacheWriteCostPerToken, - CacheReadCostPerToken: c.CacheReadCostPerToken, - } - if fallback.set() { - return fallback, rateConfigured + if c.flat.set() { + return c.flat, rateConfigured } return modelRates{}, rateNone } @@ -169,16 +224,43 @@ func (c *config) applyDefaults() { // names, and a case mismatch would silently unprice the traffic. c.pricing = make(map[string]modelRates, len(c.Pricing)) globs := map[string]modelRates{} - for k, v := range c.Pricing { + // Sorted so that when several entries are malformed the reported one is + // stable across restarts, instead of whichever map iteration reached first. + keys := make([]string, 0, len(c.Pricing)) + for k := range c.Pricing { + keys = append(keys, k) + } + sort.Strings(keys) + for _, k := range keys { lk := strings.ToLower(k) + v, err := c.Pricing[k].normalize(fmt.Sprintf("pricing[%q]", k)) + if err != nil && c.pricingErr == nil { + c.pricingErr = err + } if strings.ContainsAny(lk, "*?[") { globs[lk] = v continue } c.pricing[lk] = v } - // A bad glob is surfaced by validate(), not silently dropped. - c.pricingGlobs, c.pricingGlobErr = compilePatterns(globs) + flat, err := modelRates{ + InputCostPerMillion: c.InputCostPerMillion, + CacheWriteCostPerMillion: c.CacheWriteCostPerMillion, + CacheReadCostPerMillion: c.CacheReadCostPerMillion, + InputCostPerToken: c.InputCostPerToken, + CacheWriteCostPerToken: c.CacheWriteCostPerToken, + CacheReadCostPerToken: c.CacheReadCostPerToken, + }.normalize("config") + if err != nil && c.pricingErr == nil { + c.pricingErr = err + } + c.flat = flat + + globsCompiled, err := compilePatterns(globs) + if err != nil && c.pricingErr == nil { + c.pricingErr = fmt.Errorf("invalid pricing pattern: %w", err) + } + c.pricingGlobs = globsCompiled } // ToolPrune is the plugin. Counters live in metrics, guarded by its own mutex; @@ -232,8 +314,8 @@ func (p *ToolPrune) Configure(raw json.RawMessage) error { } } c.applyDefaults() - if c.pricingGlobErr != nil { - return fmt.Errorf("tool-prune config: invalid pricing pattern: %w", c.pricingGlobErr) + if c.pricingErr != nil { + return fmt.Errorf("tool-prune config: %w", c.pricingErr) } p.cfg = c diff --git a/authbridge/authlib/plugins/toolprune/plugin_test.go b/authbridge/authlib/plugins/toolprune/plugin_test.go index e1aba7dfa..c3de5bbf8 100644 --- a/authbridge/authlib/plugins/toolprune/plugin_test.go +++ b/authbridge/authlib/plugins/toolprune/plugin_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "math" "strings" "sync" "testing" @@ -1069,8 +1070,8 @@ func TestPricingPatternPrecedence(t *testing.T) { "*claude-opus-4-8*": {InputCostPerToken: 3}, }} c.applyDefaults() - if c.pricingGlobErr != nil { - t.Fatalf("compile: %v", c.pricingGlobErr) + if c.pricingErr != nil { + t.Fatalf("compile: %v", c.pricingErr) } for _, tc := range []struct { model string @@ -1111,7 +1112,201 @@ func TestPricingBadPatternRejected(t *testing.T) { "*claude-[opus": {InputCostPerToken: 1}, }} c.applyDefaults() - if c.pricingGlobErr == nil { + if c.pricingErr == nil { t.Fatal("want a compile error for an unterminated character class") } } + +// TestPricingPerMillionUnits is the natural-units path: an operator copies +// "$3.80 / Mtok" off a price list and the plugin prices with it, no hand +// division. Values are compared against the per-token equivalent to prove the +// conversion, not merely that something non-zero landed. +func TestPricingPerMillionUnits(t *testing.T) { + c := &config{Pricing: map[string]modelRates{ + "my-model": { + InputCostPerMillion: 3.80, + CacheWriteCostPerMillion: 4.75, + CacheReadCostPerMillion: 0.38, + }, + }} + c.applyDefaults() + if c.pricingErr != nil { + t.Fatalf("unexpected error: %v", c.pricingErr) + } + rates, src := c.ratesFor("my-model") + if src != rateConfigured { + t.Fatalf("source = %v, want rateConfigured", src) + } + // Compared with a tolerance, not for equality: a config value divides at + // runtime, so 3.80/1e6 lands one ulp below the 3.8e-06 literal. That is a + // 1e-16 relative difference on a dollar figure — not a property worth + // pinning, and pinning it would only invite a fragile test. + for _, tc := range []struct { + tier tier + want float64 + }{ + {tierInput, 0.0000038}, + {tierCacheWrite, 0.00000475}, + {tierCacheRead, 0.00000038}, + } { + got, ok := rates.rateFor(tc.tier) + if !ok || math.Abs(got-tc.want) > tc.want*1e-12 { + t.Errorf("tier %v: got (%g, %v), want (~%g, true)", tc.tier, got, ok, tc.want) + } + } +} + +// TestPricingPerMillionGlobAndFlat covers the two other places a rate can be +// stated, so per-million isn't quietly honoured in only one of the three. +func TestPricingPerMillionGlobAndFlat(t *testing.T) { + c := &config{ + Pricing: map[string]modelRates{ + "*my-family-*": {InputCostPerMillion: 2.0}, + }, + InputCostPerMillion: 9.0, + } + c.applyDefaults() + if c.pricingErr != nil { + t.Fatalf("unexpected error: %v", c.pricingErr) + } + if r, src := c.ratesFor("my-family-7"); src != rateConfigured || r.InputCostPerToken != 2.0/1e6 { + t.Errorf("glob: got (%g, %v), want (%g, configured)", r.InputCostPerToken, src, 2.0/1e6) + } + // A model no pattern claims falls to the flat rate, also stated per-million. + if r, src := c.ratesFor("totally-unknown"); src != rateConfigured || r.InputCostPerToken != 9.0/1e6 { + t.Errorf("flat: got (%g, %v), want (%g, configured)", r.InputCostPerToken, src, 9.0/1e6) + } +} + +// TestPricingUnitConflictRejected is the important one. The two units differ by +// 10^6, so silently preferring either would misprice by a millionfold with +// nothing in the readout to show which was honoured. +func TestPricingUnitConflictRejected(t *testing.T) { + for _, tc := range []struct { + name string + r modelRates + want string + }{ + {"input", modelRates{InputCostPerMillion: 3.8, InputCostPerToken: 0.0000038}, "input"}, + {"cache_write", modelRates{CacheWriteCostPerMillion: 4.75, CacheWriteCostPerToken: 0.00000475}, "cache_write"}, + {"cache_read", modelRates{CacheReadCostPerMillion: 0.38, CacheReadCostPerToken: 0.00000038}, "cache_read"}, + } { + t.Run(tc.name, func(t *testing.T) { + c := &config{Pricing: map[string]modelRates{"m": tc.r}} + c.applyDefaults() + if c.pricingErr == nil { + t.Fatal("want an error when both units are set for one tier") + } + // The message must name the tier, or an operator with three tiers + // configured has to bisect to find the one at fault. + if !strings.Contains(c.pricingErr.Error(), tc.want) { + t.Errorf("error %q does not name tier %q", c.pricingErr, tc.want) + } + }) + } + // Same rule on the flat fallback. + c := &config{InputCostPerMillion: 3.8, InputCostPerToken: 0.0000038} + c.applyDefaults() + if c.pricingErr == nil { + t.Fatal("flat fallback: want an error when both units are set") + } +} + +// TestPricingMixedUnitsAcrossTiersAllowed: stating different tiers in different +// units is odd but unambiguous, so it must not be rejected — the rule is about +// one tier stated twice, not about tidiness. +func TestPricingMixedUnitsAcrossTiersAllowed(t *testing.T) { + c := &config{Pricing: map[string]modelRates{ + "m": {InputCostPerMillion: 3.80, CacheReadCostPerToken: 0.00000038}, + }} + c.applyDefaults() + if c.pricingErr != nil { + t.Fatalf("unexpected error: %v", c.pricingErr) + } + r, _ := c.ratesFor("m") + // per-million tier: runtime division, so tolerance (see TestPricingPerMillionUnits). + if got, _ := r.rateFor(tierInput); math.Abs(got-0.0000038) > 1e-18 { + t.Errorf("input = %g, want ~0.0000038", got) + } + // per-token tier: stored verbatim, so exact. + if got, _ := r.rateFor(tierCacheRead); got != 0.00000038 { + t.Errorf("cache read = %g, want 0.00000038", got) + } +} + +// TestPricingConfigureRejectsUnitConflict proves the fault reaches Configure +// rather than stopping at applyDefaults, since that is what actually fails boot. +func TestPricingConfigureRejectsUnitConflict(t *testing.T) { + p := New() + err := p.Configure([]byte(`{"pricing":{"m":{"input_cost_per_million":3.8,"input_cost_per_token":0.0000038}}}`)) + if err == nil { + t.Fatal("Configure accepted a both-units entry") + } + if !strings.Contains(err.Error(), "input") { + t.Errorf("error %q does not name the tier", err) + } +} + +// TestPricingPerMillionJSONDecodes guards the wire names an operator types. +func TestPricingPerMillionJSONDecodes(t *testing.T) { + p := New() + if err := p.Configure([]byte(`{ + "remove": ["X"], + "pricing": {"*claude-opus-*": { + "input_cost_per_million": 3.80, + "cache_write_cost_per_million": 4.75, + "cache_read_cost_per_million": 0.38 + }}, + "input_cost_per_million": 1.0 + }`)); err != nil { + t.Fatalf("Configure: %v", err) + } + r, src := p.cfg.ratesFor("claude-opus-5") + if src != rateConfigured || math.Abs(r.InputCostPerToken-0.0000038) > 1e-18 { + t.Errorf("got (%g, %v), want (~0.0000038, configured)", r.InputCostPerToken, src) + } +} + +// TestBuiltinRatesMatchDocumentedPerMillion pins the built-in table against the +// exact per-Mtok figures the docs publish. +// +// Each expectation is written as a CONSTANT expression ($3.80 / tokensPerMillion), +// which the compiler folds exactly — the same way pricing.go does. So this fails +// both if a documented figure and the table drift apart, and if anyone converts +// the table with a runtime division instead, which lands a ulp low. +// +// It deliberately does not assert rate*1e6 == 3.80: multiplying back is a second +// rounding that isn't exact for every value (0.076 round-trips to +// 0.07600000000000001), which would make the test fail for a reason that has +// nothing to do with the table being right. +func TestBuiltinRatesMatchDocumentedPerMillion(t *testing.T) { + c := &config{} + c.applyDefaults() + for _, tc := range []struct { + model string + input, cacheWr, cacheRead float64 + }{ + {"claude-opus-5", 3.80 / tokensPerMillion, 4.75 / tokensPerMillion, 0.38 / tokensPerMillion}, + {"claude-sonnet-5", 1.52 / tokensPerMillion, 1.90 / tokensPerMillion, 0.152 / tokensPerMillion}, + {"claude-haiku-4-5", 0.76 / tokensPerMillion, 0.95 / tokensPerMillion, 0.076 / tokensPerMillion}, + } { + r, src := c.ratesFor(tc.model) + if src != rateDefault { + t.Errorf("%s: src = %v, want rateDefault", tc.model, src) + continue + } + for _, f := range []struct { + name string + got, want float64 + }{ + {"input", r.InputCostPerToken, tc.input}, + {"cache write", r.CacheWriteCostPerToken, tc.cacheWr}, + {"cache read", r.CacheReadCostPerToken, tc.cacheRead}, + } { + if f.got != f.want { + t.Errorf("%s %s: %v, want exactly %v ($%v/Mtok)", + tc.model, f.name, f.got, f.want, f.want*tokensPerMillion) + } + } + } +} diff --git a/authbridge/authlib/plugins/toolprune/pricing.go b/authbridge/authlib/plugins/toolprune/pricing.go index bd906fd31..734abcb4a 100644 --- a/authbridge/authlib/plugins/toolprune/pricing.go +++ b/authbridge/authlib/plugins/toolprune/pricing.go @@ -41,21 +41,24 @@ type patternRates struct { // // Any `pricing` entry in config overrides the matching model. var defaultPatterns = mustCompilePatterns(map[string]modelRates{ - // input 1.00x / cache write 1.25x / cache read 0.10x + // Written in the published unit — dollars per million tokens — and divided by + // a CONSTANT, so the compiler folds each one exactly. A runtime division + // lands a ulp low (3.7999999999999996e-06), which would make this table + // disagree with the documented $3.80/Mtok in the last digit for no reason. "*claude-opus-*": { - InputCostPerToken: 0.0000038, - CacheWriteCostPerToken: 0.00000475, - CacheReadCostPerToken: 0.00000038, + InputCostPerToken: 3.80 / tokensPerMillion, + CacheWriteCostPerToken: 4.75 / tokensPerMillion, // 1.25x input + CacheReadCostPerToken: 0.38 / tokensPerMillion, // 0.10x input }, "*claude-sonnet-*": { - InputCostPerToken: 0.00000152, - CacheWriteCostPerToken: 0.0000019, - CacheReadCostPerToken: 0.000000152, + InputCostPerToken: 1.52 / tokensPerMillion, + CacheWriteCostPerToken: 1.90 / tokensPerMillion, + CacheReadCostPerToken: 0.152 / tokensPerMillion, }, "*claude-haiku-*": { - InputCostPerToken: 0.00000076, - CacheWriteCostPerToken: 0.00000095, - CacheReadCostPerToken: 0.000000076, + InputCostPerToken: 0.76 / tokensPerMillion, + CacheWriteCostPerToken: 0.95 / tokensPerMillion, + CacheReadCostPerToken: 0.076 / tokensPerMillion, }, }) diff --git a/authbridge/docs/plugin-catalog.md b/authbridge/docs/plugin-catalog.md index 534c36c10..9fd977b63 100644 --- a/authbridge/docs/plugin-catalog.md +++ b/authbridge/docs/plugin-catalog.md @@ -240,8 +240,8 @@ body-reading plugin (it rewrites the request body). Declares - `remove` (`[]string`) — tool names to delete from the manifest. The complete verdict: no learning, no state, no storage. Names absent from a given request are ignored. - `paths` (`[]string`) — request paths to act on, matched exactly or by suffix. Defaults to `/v1/chat/completions`, `/v1/completions`, `/v1/messages`. -- `pricing` (`map[model]rates`) — per-token rates keyed by model name **or glob** (`*claude-opus-*`), each with `input_cost_per_token`, `cache_write_cost_per_token`, `cache_read_cost_per_token`. **Optional**: built-in patterns cover the Claude families on the rossoctl gateway, so `$ saved` works unconfigured; any entry here overrides the built-in. Per model because rates differ ~5x across opus/sonnet/haiku. Built-ins are keyed by *family*, not version, so an opus 4.8 → 5 rename needs no code change. Resolution: exact key → longest matching glob → built-in pattern → flat fallback → unpriced; keys matched case-insensitively. An invalid glob fails startup with the key named. -- `input_cost_per_token`, `cache_write_cost_per_token`, `cache_read_cost_per_token` (`float`) — optional flat fallback for models absent from `pricing`. A figure from built-in rates is labelled as such; a model in neither the table nor config is counted in a `requests unpriced` row instead of charged at another model's rate. No output rate: pruning only shrinks the prompt. +- `pricing` (`map[model]rates`) — rates keyed by model name **or glob** (`*claude-opus-*`), each with `input_cost_per_million`, `cache_write_cost_per_million`, `cache_read_cost_per_million` (per-million: the unit providers publish, so `3.80` not `0.0000038`). The per-token names are also accepted for `litellm-budget-track` parity; setting both units for one tier fails startup, since they differ by 10^6 and picking a winner silently would misprice by that factor. **Optional**: built-in patterns cover the Claude families on the rossoctl gateway, so `$ saved` works unconfigured; any entry here overrides the built-in. Per model because rates differ ~5x across opus/sonnet/haiku. Built-ins are keyed by *family*, not version, so an opus 4.8 → 5 rename needs no code change. Resolution: exact key → longest matching glob → built-in pattern → flat fallback → unpriced; keys matched case-insensitively. An invalid glob fails startup with the key named. +- `input_cost_per_million`, `cache_write_cost_per_million`, `cache_read_cost_per_million` (`float`) — optional flat fallback for models absent from `pricing` (per-token variants also accepted). A figure from built-in rates is labelled as such; a model in neither the table nor config is counted in a `requests unpriced` row instead of charged at another model's rate. No output rate: pruning only shrinks the prompt. Generate the list from local transcripts with `abctl tools scan`, which proposes only tools it recognises as Claude Code built-ins and never diff --git a/authbridge/docs/tool-prune-plugin.md b/authbridge/docs/tool-prune-plugin.md index 0b0c09407..0cf4e2e13 100644 --- a/authbridge/docs/tool-prune-plugin.md +++ b/authbridge/docs/tool-prune-plugin.md @@ -208,14 +208,31 @@ model outright: config: remove: [CronCreate, NotebookEdit] pricing: - claude-opus-5: - input_cost_per_token: 0.0000038 - cache_write_cost_per_token: 0.00000475 - cache_read_cost_per_token: 0.00000038 + "*claude-opus-*": + input_cost_per_million: 3.80 + cache_write_cost_per_million: 4.75 + cache_read_cost_per_million: 0.38 # optional flat fallback for models absent from the table above - input_cost_per_token: 0.0000038 + input_cost_per_million: 3.80 ``` +**Rates are stated per million tokens**, because that is the unit every provider +publishes and the one you already have in hand — `3.80`, copied straight off a +price list, rather than `0.0000038` arrived at by dividing in your head. That +difference is not just ergonomics: `0.0000038` is six leading zeros, and +`0.000038` is a plausible-looking typo that misprices by 10x with nothing in the +readout to reveal it. + +The per-token field names are still accepted (`input_cost_per_token`, …), for +parity with [`litellm-budget-track`](./litellm-budgettrack-plugin.md) and because +LiteLLM's own `model_prices_and_context_window.json` is per-token, so rates get +copied out of it verbatim. Different tiers may use different units. + +**Setting both units for the same tier is a startup error**, not a precedence +question. The two differ by 106: silently honouring one would either +overstate a saving a millionfold or bury it under rounding, and the readout gives +you no way to tell which happened. The error names the offending entry and tier. + Model keys match what the parser records (`Extensions.Inference.Model`) and are matched case-insensitively, since gateways vary in how they echo the name and a case mismatch would silently unprice the traffic. @@ -228,7 +245,7 @@ so the more specific statement wins: 2. glob in your `pricing` — **longest pattern first**, so `*claude-opus-4-8*` beats `*claude-opus-*` deterministically rather than by map iteration luck 3. built-in family pattern -4. the flat `input_cost_per_token` fallback +4. the flat `input_cost_per_million` fallback 5. unpriced An invalid pattern fails startup with the offending key named, rather than From 0096fa000518af9a0c6a0a9ca35ec974b6321219 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Thu, 3 Sep 2026 12:26:45 -0400 Subject: [PATCH 28/28] =?UTF-8?q?fix:=20Address=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20observe-mode=20accounting,=20id=20collisions,=20bri?= =?UTF-8?q?dge=20warning?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Seven of the thirteen open review findings were real. Taking them in order of what they'd cost someone: Observe mode reported a saving it had not made. The event was published before SetBody, carrying BodyBytesAfter: len(out) — but under on_error: observe SetBody is a no-op and the ORIGINAL body goes upstream. abctl divides the response's prompt tokens by that field to get tokens-per-byte, so it was calibrating against a body that was never sent and inflating the ratio. The event now reports the size actually sent and carries Projected, and abctl renders a projected figure as "~3.4k" instead of "−3.4k" so it can't be added up as money not spent. observe exists to be trusted while it isn't enforcing, which makes a wrong number there worse than no number. Request ids could collide. 48 random bits is birthday-bounded at about 0.2% per million ids, and a collision is not cosmetic: the id is what pairs a response to its request, so a duplicate files a response under the wrong request — precisely the misattribution the field was added to eliminate. A process-local counter carries the uniqueness now, with a short random suffix kept only so a consumer merging two proxies' streams can still tell them apart. The TLS-bridge warning cried wolf and then went quiet. noteTunnel counted every CONNECT, including hosts in TLSBridge.Skip and ones classification deliberately passed through, so five intentional tunnels tripped the "client doesn't trust the CA" warning — and since it's once-only, that false positive then MASKED the real failure later. Counting only bridge-eligible attempts fixes the false positive but would break the true case on its own: a handshake refusal adds the host to Skip, so attempts stop accumulating and the threshold is never crossed. So the refusal itself now warns immediately — it's proof, not evidence, and needs no threshold. The laptop guide overclaimed safety. It said drift "costs savings rather than correctness". That is wrong in the direction that matters: the scan knows what you haven't used, not what you won't need, and a pruned tool the model later needs is simply absent — a functional failure. Now stated plainly, with rescanning and observe mode as the two cheap mitigations. Also: repaired a godoc sentence of mine that trailed off mid-thought; documented that the cache_control victim scan keeps the last marker deliberately; labelled five code fences (MD040), fixed an MD028 blockquote break, and Americanised one spelling. Four findings were rejected, for the record: - "Normalize should imply ReadsBody from either write flag, and framework-architecture.md still references BodyAccess" — Normalize already does this, and the BodyAccess mentions are in the §12 Versioning changelog, one of which literally reads "it has since been removed". Editing history entries to drop a field that historically existed would make the changelog wrong. - "Reject [reader, response-mutator] in pipeline.New" — deliberate and already documented: the check would reject [opa, sparc], which is a valid chain today. A startup WARN names the shape instead; closing it properly needs direction-specific read capabilities. - "DeriveError should capture an error body regardless of plugin body requirements" — a real limitation, already documented at the call site, but it means buffering error responses on chains that opted out of body access. Out of scope for this PR. - "Make the metrics readout say rates are gateway rates" — it already does: "built-in rates (discounted gateway; understates list pricing)". Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- .../forwardproxy/bridgehealth_test.go | 60 ++++++++++++++++--- .../authlib/listener/forwardproxy/server.go | 47 ++++++++++++--- authbridge/authlib/pipeline/requestid.go | 22 ++++++- .../authlib/pipeline/requestid_unique_test.go | 51 ++++++++++++++++ authbridge/authlib/plugins/toolprune/event.go | 17 ++++-- .../authlib/plugins/toolprune/plugin.go | 30 +++++++--- authbridge/cmd/abctl/tui/events_pane.go | 2 +- authbridge/cmd/abctl/tui/prune_saving.go | 13 +++- authbridge/cmd/abctl/tui/prune_saving_test.go | 33 +++++++++- authbridge/docs/laptop-token-savings.md | 20 ++++++- authbridge/docs/plugin-reference.md | 2 +- authbridge/docs/tool-prune-plugin.md | 14 +++-- docs/proposals/tool-prune.md | 4 +- 13 files changed, 269 insertions(+), 46 deletions(-) create mode 100644 authbridge/authlib/pipeline/requestid_unique_test.go diff --git a/authbridge/authlib/listener/forwardproxy/bridgehealth_test.go b/authbridge/authlib/listener/forwardproxy/bridgehealth_test.go index 9c663d50d..de79503c7 100644 --- a/authbridge/authlib/listener/forwardproxy/bridgehealth_test.go +++ b/authbridge/authlib/listener/forwardproxy/bridgehealth_test.go @@ -8,11 +8,14 @@ import ( "github.com/rossoctl/cortex/authbridge/authlib/tlsbridge" ) -// TestNoteTunnel_WarnsOnlyWhenNothingIsEverDecrypted covers the failure that -// looks like a plugin bug: the bridge is on, the client does not trust its CA, -// so every HTTPS request opens an opaque tunnel and every body-reading plugin -// correctly does nothing. Nothing errors — the only symptom is silence. -func TestNoteTunnel_WarnsOnlyWhenNothingIsEverDecrypted(t *testing.T) { +// TestNoteBridgeAttempt_WarnsOnlyWhenNothingIsEverDecrypted covers the failure +// that looks like a plugin bug: the bridge is on, the client does not trust its +// CA, so every HTTPS request opens an opaque tunnel and every body-reading +// plugin correctly does nothing. Nothing errors — the only symptom is silence. +// +// Driven through noteBridgeAttempt, not noteTunnel: the trigger is a CONNECT the +// bridge actually tried to decrypt. See TestNoteTunnel_PassthroughNeverWarns. +func TestNoteBridgeAttempt_WarnsOnlyWhenNothingIsEverDecrypted(t *testing.T) { tests := []struct { name string bridge *tlsbridge.Engine @@ -27,20 +30,20 @@ func TestNoteTunnel_WarnsOnlyWhenNothingIsEverDecrypted(t *testing.T) { wantWarns: 0, }, { - name: "below threshold: a few tunnels are normal (passthrough hosts, startup races)", + name: "below threshold: a few attempts are normal (startup races)", bridge: &tlsbridge.Engine{}, tunnels: tunnelWarnThreshold - 1, wantWarns: 0, }, { - name: "tunnels but something was decrypted: bridge is working", + name: "attempts but something was decrypted: bridge is working", bridge: &tlsbridge.Engine{}, tunnels: 50, bridged: 1, wantWarns: 0, }, { - name: "many tunnels, nothing decrypted: warn", + name: "many attempts, nothing decrypted: warn", bridge: &tlsbridge.Engine{}, tunnels: tunnelWarnThreshold, wantWarns: 1, @@ -61,7 +64,7 @@ func TestNoteTunnel_WarnsOnlyWhenNothingIsEverDecrypted(t *testing.T) { // the guarded block would run by observing the sync.Once directly. for i := 0; i < tc.tunnels; i++ { before := s.warnFired() - s.noteTunnel() + s.noteBridgeAttempt() if !before && s.warnFired() { warns++ } @@ -107,3 +110,42 @@ func TestNoteTunnel_ConcurrentIsRaceFree(t *testing.T) { t.Errorf("tunnelsOpened = %d, want %d", got, 16*64) } } + +// TestNoteTunnel_PassthroughNeverWarns is the regression this split exists for. +// A CONNECT to a host in TLSBridge.Skip, or one classification chose to pass +// through, is intentional — it is not evidence of a broken CA. Counting those +// let a correctly-configured proxy cry wolf, and because the warning is +// once-only, the false positive then MASKED the real failure if it came later. +func TestNoteTunnel_PassthroughNeverWarns(t *testing.T) { + s := &Server{TLSBridge: &tlsbridge.Engine{}} + for i := 0; i < tunnelWarnThreshold*20; i++ { + s.noteTunnel() + } + if s.warnFired() { + t.Error("warned about intentional passthrough tunnels") + } + // The warning must still be available afterwards for a genuine failure — + // i.e. the sync.Once was not burned by the passthrough traffic above. + for i := 0; i < tunnelWarnThreshold; i++ { + s.noteBridgeAttempt() + } + if !s.warnFired() { + t.Error("real bridge failure did not warn after passthrough traffic") + } +} + +// TestNoteBridgeHandshakeFailure_WarnsImmediately: a refused forged certificate +// is proof, so it must not wait for a threshold. It especially must not, because +// the refusal adds the host to Skip — later requests never reach +// noteBridgeAttempt, so the threshold alone would never be crossed. +func TestNoteBridgeHandshakeFailure_WarnsImmediately(t *testing.T) { + s := &Server{TLSBridge: &tlsbridge.Engine{}} + s.noteBridgeAttempt() // one attempt, well below threshold + if s.warnFired() { + t.Fatal("warned on a single attempt") + } + s.noteBridgeHandshakeFailure() + if !s.warnFired() { + t.Error("a rejected bridge certificate did not warn") + } +} diff --git a/authbridge/authlib/listener/forwardproxy/server.go b/authbridge/authlib/listener/forwardproxy/server.go index 99130dfcf..f32eb97b5 100644 --- a/authbridge/authlib/listener/forwardproxy/server.go +++ b/authbridge/authlib/listener/forwardproxy/server.go @@ -81,6 +81,7 @@ type Server struct { // Nothing errors, so the only symptom is silence. These count the two // outcomes so the listener can say so out loud. tunnelsOpened atomic.Uint64 + bridgeAttempts atomic.Uint64 bridgedRequests atomic.Uint64 bridgeWarnOnce sync.Once bridgeWarned atomic.Bool @@ -560,6 +561,11 @@ func (s *Server) bridgeServe(client net.Conn, authority, host string) bool { if err != nil { s.TLSBridge.Skip.Add(host) // pinned client → its retry will passthrough slog.Warn("tls-bridge passthrough", "host", host, "reason", "handshake-fail", "error", err) + // Proof the client doesn't trust the CA. Warn now with the fix, because + // Skip.Add above means this host never reaches noteBridgeAttempt again. + if s.bridgedRequests.Load() == 0 { + s.noteBridgeHandshakeFailure() + } return true // conn is dead post-forge; nothing left to tunnel } @@ -1064,6 +1070,7 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { key := hostOnly(r.Host) if !s.TLSBridge.Skip.Contains(key) { if v, _ := s.TLSBridge.Decision.Classify(key, portOf(r.Host), first); v == tlsbridge.Terminate { + s.noteBridgeAttempt() _ = upstream.Close() // bridgeServe dials its own verified upstream if s.bridgeServe(clientConn, authority, key) { return @@ -1260,24 +1267,48 @@ func portOf(authority string) int { // client is not trusting the CA. const tunnelWarnThreshold = 5 -// noteTunnel counts a CONNECT tunnel and, once, warns if the bridge is enabled -// yet has never decrypted anything. +// noteTunnel counts a CONNECT tunnel. It deliberately does NOT warn: a CONNECT +// says nothing about bridge health yet, because the destination may be in +// TLSBridge.Skip or classified as passthrough on purpose. Warning here counted +// intentional opaque tunnels as evidence of a broken CA — and because the +// warning is once-only, those false positives then masked the real failure when +// it happened later. The warning lives on the bridge-eligible path instead. +func (s *Server) noteTunnel() { s.tunnelsOpened.Add(1) } + +// noteBridgeAttempt counts a CONNECT that classification chose to terminate, and +// warns once if the bridge has been asked to decrypt this many times and never +// managed it. // // This is the failure that looks like a bug in whatever plugin you are testing: // tool-prune, the parsers and every body reader correctly do nothing, because -// there is no plaintext to act on. Naming the trust anchor turns a silent -// dead end into a one-line fix. -func (s *Server) noteTunnel() { - n := s.tunnelsOpened.Add(1) +// there is no plaintext to act on. Naming the trust anchor turns a silent dead +// end into a one-line fix. +func (s *Server) noteBridgeAttempt() { + n := s.bridgeAttempts.Add(1) if s.TLSBridge == nil || n < tunnelWarnThreshold || s.bridgedRequests.Load() > 0 { return } + s.warnBridgeUnused("bridge_attempts", n, "the client does not trust the bridge CA") +} + +// noteBridgeHandshakeFailure reports the unambiguous case: the client refused the +// forged certificate. Unlike the attempt threshold this needs no accumulation — +// one refusal already proves the trust anchor is not installed. It matters that +// this path warns, because a refusal adds the host to Skip, so later requests +// never reach noteBridgeAttempt and the threshold alone would never be crossed. +func (s *Server) noteBridgeHandshakeFailure() { + s.warnBridgeUnused("bridge_attempts", s.bridgeAttempts.Load(), + "the client rejected the bridge certificate, so it does not trust the bridge CA") +} + +func (s *Server) warnBridgeUnused(countKey string, count uint64, cause string) { s.bridgeWarnOnce.Do(func() { s.bridgeWarned.Store(true) slog.Warn("tls-bridge: enabled but nothing has been decrypted — every request is tunnelling through opaquely, so body-reading plugins (parsers, tool-prune) cannot act", - "tunnels_opened", n, + countKey, count, + "tunnels_opened", s.tunnelsOpened.Load(), "bridged_requests", 0, - "likely_cause", "the client does not trust the bridge CA", + "likely_cause", cause, "fix", "point the client at the trust anchor, e.g. NODE_EXTRA_CA_CERTS="+s.caFileHint()) }) } diff --git a/authbridge/authlib/pipeline/requestid.go b/authbridge/authlib/pipeline/requestid.go index 01f6fadf6..e23b82f39 100644 --- a/authbridge/authlib/pipeline/requestid.go +++ b/authbridge/authlib/pipeline/requestid.go @@ -3,20 +3,38 @@ package pipeline import ( "crypto/rand" "encoding/hex" + "strconv" + "sync/atomic" ) +// requestIDSeq makes ids collision-free within a process by construction. +var requestIDSeq atomic.Uint64 + // newRequestID returns a short, unique-per-process request identifier. // // Not a UUID on purpose: it exists to pair a request event with its response // event in a session timeline, so it needs to be unique among in-flight // requests and short enough to read in a terminal — not globally unique or // cryptographically meaningful. +// +// A monotonic counter carries the uniqueness rather than randomness alone. +// Random-only was 48 bits, which sounds ample but is birthday-bounded: about a +// 0.2% chance of at least one collision within a million ids. A collision is not +// cosmetic here — the consumer pairs a response to a request BY this id, so two +// requests sharing one puts a response under the wrong request, which is exactly +// the misattribution this field was added to eliminate. A counter cannot collide +// with itself, so the failure mode is gone rather than made unlikely. +// +// The random suffix stays for cross-process distinction: a consumer can merge +// streams from two authbridge instances (the advanced demo runs an agent-side +// and a tool-side proxy), where both counters start at 1. func newRequestID() string { - var b [6]byte + n := requestIDSeq.Add(1) + var b [3]byte // crypto/rand.Read never returns an error as of Go 1.24 — it panics on an // unusable system source instead — so there is no failure branch to write. _, _ = rand.Read(b[:]) - return hex.EncodeToString(b[:]) + return strconv.FormatUint(n, 36) + "-" + hex.EncodeToString(b[:]) } // RequestID returns a stable identifier for this request, generated on first diff --git a/authbridge/authlib/pipeline/requestid_unique_test.go b/authbridge/authlib/pipeline/requestid_unique_test.go new file mode 100644 index 000000000..27bae35da --- /dev/null +++ b/authbridge/authlib/pipeline/requestid_unique_test.go @@ -0,0 +1,51 @@ +package pipeline + +import ( + "sync" + "testing" +) + +// TestRequestIDNoCollisions: the id is what pairs a response to its request, so +// a duplicate silently files a response under the wrong request. Randomness +// alone made that unlikely; the counter makes it impossible in-process. +func TestRequestIDNoCollisions(t *testing.T) { + const n = 200_000 + seen := make(map[string]struct{}, n) + for i := 0; i < n; i++ { + id := newRequestID() + if _, dup := seen[id]; dup { + t.Fatalf("collision after %d ids: %q", i, id) + } + seen[id] = struct{}{} + } +} + +// TestRequestIDConcurrent: listeners generate ids from many goroutines. +func TestRequestIDConcurrent(t *testing.T) { + const goroutines, each = 32, 2000 + var mu sync.Mutex + seen := make(map[string]struct{}, goroutines*each) + var wg sync.WaitGroup + for g := 0; g < goroutines; g++ { + wg.Add(1) + go func() { + defer wg.Done() + ids := make([]string, 0, each) + for i := 0; i < each; i++ { + ids = append(ids, newRequestID()) + } + mu.Lock() + defer mu.Unlock() + for _, id := range ids { + if _, dup := seen[id]; dup { + t.Errorf("concurrent collision: %q", id) + } + seen[id] = struct{}{} + } + }() + } + wg.Wait() + if len(seen) != goroutines*each { + t.Errorf("got %d unique ids, want %d", len(seen), goroutines*each) + } +} diff --git a/authbridge/authlib/plugins/toolprune/event.go b/authbridge/authlib/plugins/toolprune/event.go index 4e47075a7..b531f919f 100644 --- a/authbridge/authlib/plugins/toolprune/event.go +++ b/authbridge/authlib/plugins/toolprune/event.go @@ -19,10 +19,19 @@ import ( // No body content: the session store is unauthenticated, so this holds counts, // tool names the operator themselves configured, and rates. type pruneEvent struct { - ToolsRemoved []string `json:"toolsRemoved,omitempty"` - BytesRemoved int `json:"bytesRemoved"` - BodyBytesAfter int `json:"bodyBytesAfter"` - Model string `json:"model,omitempty"` + ToolsRemoved []string `json:"toolsRemoved,omitempty"` + BytesRemoved int `json:"bytesRemoved"` + // BodyBytesAfter is the size of the body actually SENT upstream, which is + // not the pruned size under on_error: observe — there SetBody is a no-op and + // the original goes out. A consumer divides the response's prompt-token + // count by this to get tokens-per-byte, so using the pruned size while the + // original was billed would inflate that ratio and overstate the saving. + BodyBytesAfter int `json:"bodyBytesAfter"` + // Projected marks a saving that was measured but NOT applied — observe mode. + // The bytes were not actually removed from the request, so a consumer must + // present this as "would have saved", never as money already not spent. + Projected bool `json:"projected,omitempty"` + Model string `json:"model,omitempty"` // Rates are USD per token for this request's model, already resolved // through config → flat fallback → built-in defaults. diff --git a/authbridge/authlib/plugins/toolprune/plugin.go b/authbridge/authlib/plugins/toolprune/plugin.go index 1e0f0fd80..9ce62d407 100644 --- a/authbridge/authlib/plugins/toolprune/plugin.go +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -126,8 +126,10 @@ const tokensPerMillion = 1_000_000 // precedence. The two differ by 10^6, so picking a winner silently would either // overstate a saving by a millionfold or bury it below rounding — and the // readout gives an operator no way to tell which unit was honoured. A startup -// error naming the tier is the only outcome that can't be misread. what -// identifies the offending entry ("pricing[\"claude-opus-5\"]"). +// error naming the tier is the only outcome that can't be misread. +// +// what names the entry being normalized, so the error can point at it — +// `pricing["claude-opus-5"]` for a map entry, "config" for the flat fallback. func (r modelRates) normalize(what string) (modelRates, error) { for _, f := range []struct { name string @@ -558,6 +560,10 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action for _, v := range victims { victimSet[v] = true } + // Last marker wins: if two pruned tools each carried a breakpoint, only + // one can move to the single last survivor. Claude Code marks exactly one + // tool, so this is not a shape seen in practice — but a future reader + // should know the overwrite is deliberate, not an oversight. var orphanedCacheControl gjson.Result for _, v := range victims { if cc := gjson.GetBytes(body, fmt.Sprintf("tools.%d.cache_control", v)); cc.Exists() { @@ -629,10 +635,22 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action if !okR { rateRead = 0 } + // SetBody BEFORE publishing, so the event can report what was actually sent. + // Under ErrorPolicyObserve it is a no-op on bytes and leaves bodyMutated + // false — this same code path measures without enforcing. + pctx.SetBody(out) + applied := pctx.BodyMutated() + // The body upstream actually sees: the rewrite when it was applied, the + // original when it was only measured. + bodySent := len(out) + if !applied { + bodySent = len(body) + } p.publish(pctx, pruneEvent{ ToolsRemoved: names, BytesRemoved: removedBytes, - BodyBytesAfter: len(out), + BodyBytesAfter: bodySent, + Projected: !applied, Model: inferenceModel(pctx), RateInput: rateInput, RateCacheWrite: rateWrite, @@ -643,11 +661,7 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action // it came out of. SetState keeps it private to this plugin, unlike // Extensions.Custom which is shared. pipeline.SetState(pctx, p.Name(), &requestState{bytesRemoved: removedBytes}) - pctx.SetBody(out) - // Under ErrorPolicyObserve, SetBody is a no-op on bytes and leaves - // bodyMutated false — so this same code path measures without enforcing, - // and the counter it lands in is what distinguishes the two. - if pctx.BodyMutated() { + if applied { p.m.pruned(names, removedBytes) } else { p.m.projected(names, removedBytes) diff --git a/authbridge/cmd/abctl/tui/events_pane.go b/authbridge/cmd/abctl/tui/events_pane.go index 4a0ee14a0..c083f2fce 100644 --- a/authbridge/cmd/abctl/tui/events_pane.go +++ b/authbridge/cmd/abctl/tui/events_pane.go @@ -799,5 +799,5 @@ func (m *model) tokensCellWithSaving(rows []eventRow, partner map[int]int, i int if !ok { return "" } - return formatSavedOnly(tokens, usd, ps.RateSource) + return formatSavedOnly(tokens, usd, ps.RateSource, ps.Projected) } diff --git a/authbridge/cmd/abctl/tui/prune_saving.go b/authbridge/cmd/abctl/tui/prune_saving.go index 17bec22be..490a0b1c4 100644 --- a/authbridge/cmd/abctl/tui/prune_saving.go +++ b/authbridge/cmd/abctl/tui/prune_saving.go @@ -17,6 +17,9 @@ type pruneSaving struct { RateCacheWrite float64 `json:"rateCacheWrite"` RateCacheRead float64 `json:"rateCacheRead"` RateSource string `json:"rateSource"` + // Projected marks observe mode: the saving was measured but the bytes were + // not actually removed, so it must not render as money already not spent. + Projected bool `json:"projected"` } // decodePruneSaving pulls the tool-prune event off a request event, if present. @@ -76,11 +79,19 @@ func savedTokensAndCost(ps pruneSaving, resp *pipeline.InferenceExtension) (toke // formatSavedOnly renders a request row's saving: what was removed and what it // was worth. No total, because a request has no billed token count — that // belongs to the response, on its own row. -func formatSavedOnly(tokens, usd float64, rateSource string) string { +// +// A projected saving (on_error: observe, where the bytes were measured but not +// removed) is prefixed "~" and drops the "−". Rendering it identically to a real +// saving would invite an operator to add up money that was still spent, and +// observe mode exists precisely to be trusted while it is not yet enforcing. +func formatSavedOnly(tokens, usd float64, rateSource string, projected bool) string { if tokens <= 0 { return "" } cell := "−" + formatCompact(tokens) + if projected { + cell = "~" + formatCompact(tokens) + } if usd > 0 && rateSource != "none" { cell += fmt.Sprintf(" $%s", formatUSD(usd)) } diff --git a/authbridge/cmd/abctl/tui/prune_saving_test.go b/authbridge/cmd/abctl/tui/prune_saving_test.go index 57314241d..125c6d377 100644 --- a/authbridge/cmd/abctl/tui/prune_saving_test.go +++ b/authbridge/cmd/abctl/tui/prune_saving_test.go @@ -93,7 +93,7 @@ func TestSavedTokensAndCost_TierDecidesTheValue(t *testing.T) { // beside a response total read as though the response had shrunk, which it had // not. func TestFormatSavedOnly(t *testing.T) { - got := formatSavedOnly(10577.5, 0.05024, "default") + got := formatSavedOnly(10577.5, 0.05024, "default", false) for _, want := range []string{"−10.6k", "$0.050"} { if !strings.Contains(got, want) { t.Errorf("cell %q missing %q", got, want) @@ -103,11 +103,11 @@ func TestFormatSavedOnly(t *testing.T) { t.Errorf("cell %q should carry no billed total", got) } // Nothing saved: an empty cell, so unrelated request rows stay blank. - if got := formatSavedOnly(0, 0, "default"); got != "" { + if got := formatSavedOnly(0, 0, "default", false); got != "" { t.Errorf("no-saving cell = %q, want empty", got) } // Unpriced model: tokens shown, no dollar figure invented. - got = formatSavedOnly(10577.5, 0, "none") + got = formatSavedOnly(10577.5, 0, "none", false) if strings.Contains(got, "$") { t.Errorf("cell %q shows a price for an unpriced model", got) } @@ -182,3 +182,30 @@ func TestTokensCellWithSaving_RequiresAnIDMatch(t *testing.T) { t.Error("an id-matched pair should price") } } + +// TestFormatSavedOnlyProjected: an observe-mode figure must be visually distinct +// from a realized one, or an operator adds up money that was still spent. +func TestFormatSavedOnlyProjected(t *testing.T) { + real := formatSavedOnly(10577.5, 0.05024, "default", false) + proj := formatSavedOnly(10577.5, 0.05024, "default", true) + if real == proj { + t.Fatalf("projected renders identically to realized: %q", real) + } + if !strings.HasPrefix(proj, "~") { + t.Errorf("projected = %q, want a leading ~", proj) + } + if strings.Contains(proj, "−") { + t.Errorf("projected = %q, must not claim bytes were removed", proj) + } +} + +// TestPruneSavingProjectedDecodes guards the wire tag. +func TestPruneSavingProjectedDecodes(t *testing.T) { + ps, ok := decodePruneSaving(reqEvent(t, `{"bytesRemoved":100,"bodyBytesAfter":1000,"projected":true}`)) + if !ok { + t.Fatal("decode failed") + } + if !ps.Projected { + t.Error("projected did not decode") + } +} diff --git a/authbridge/docs/laptop-token-savings.md b/authbridge/docs/laptop-token-savings.md index 37b4d5867..b88f59a49 100644 --- a/authbridge/docs/laptop-token-savings.md +++ b/authbridge/docs/laptop-token-savings.md @@ -66,9 +66,23 @@ abctl tools scan --write ~/.cortex/config.yaml `tools scan` reads your own `~/.claude/projects/*.jsonl` transcripts and proposes the built-in tools you have not called in 30 days. It only ever proposes tools it -recognises, and never one it has seen you call — removing a tool the model needs -is the harmful direction of failure, so drift costs savings rather than -correctness. The config is hot-reloaded; no restart. +recognises, and never one it has seen you call. The config is hot-reloaded; no +restart. + +**What the scan cannot know is the future.** It reports what you have not used, +not what you will not need. If you start a kind of work that needs a pruned tool, +its definition is gone from the request and the model cannot call it — that is a +functional failure, not merely a smaller saving. Two things keep it cheap: + +- **Rescan occasionally** (say monthly, or after your work changes shape) so the + list tracks what you actually use. The list is only ever as current as the last + scan. +- **Reach for `on_error: observe` when in doubt.** It measures the saving and + changes nothing, so you can see what a list would be worth before trusting it. + abctl marks those figures with `~`. + +If a tool goes missing, the fix is to delete its name from `remove:` — the config +hot-reloads, so it comes back without a restart. ## 4. Point Claude Code at it diff --git a/authbridge/docs/plugin-reference.md b/authbridge/docs/plugin-reference.md index 33d485044..abd6a4c6f 100644 --- a/authbridge/docs/plugin-reference.md +++ b/authbridge/docs/plugin-reference.md @@ -790,7 +790,7 @@ nothing about how the response may be relayed. > mutator would see rewritten bytes. Not enforced, because the check would reject > chains that validate today; closing it needs direction-specific *read* > capabilities. - +> > **Declaring is a contract, not an enforcement.** `SetBody` flips > `bodyMutated` unconditionally outside observe mode and the listeners gate > purely on that flag, so a plugin that calls `SetBody` *without* declaring the diff --git a/authbridge/docs/tool-prune-plugin.md b/authbridge/docs/tool-prune-plugin.md index 0cf4e2e13..72f0fab2c 100644 --- a/authbridge/docs/tool-prune-plugin.md +++ b/authbridge/docs/tool-prune-plugin.md @@ -75,7 +75,7 @@ Two occasions worth it: `abctl`'s plugin detail pane shows a `Metrics:` section (source: `GET /v1/pipeline`): -``` +```text Metrics: requests seen 2 count requests pruned 2 count @@ -98,12 +98,18 @@ The events pane's `TOKENS / SAVED` column splits the two halves across the rows they belong to — the saving on the request that was rewritten, the billed total on the response: -``` +```text # PHASE ACTION PLUGIN TOKENS / SAVED CODE 12 req modify tool-prune −24.7k $0.117 12 resp observe inference-parser 34,702 200 ``` +Under `on_error: observe` the saving is **projected**: the plugin measured what it +would remove but sent the request unchanged, so nothing was actually saved. Those +figures render with a leading `~` and no `−`, and the aggregate counts them +separately — an observe-mode run must not be added up as money not spent. + + The saving is not shown on the response row: nothing about the response was reduced, and putting it there reads as though it had been. The two rows share a `#` so they are read together anyway. @@ -289,7 +295,7 @@ change the plugin. ## Where the list comes from -``` +```sh abctl tools scan [--days 30] [--keep Name,Name] [--dir PATH] [--write CONFIG] ``` @@ -322,7 +328,7 @@ rather than a silent no-op. Every error path forwards the original bytes unmodified: the plugin fails open on a malformed or truncated body, an unparseable manifest, a rewrite that does not shrink the body, a rewrite that produces invalid JSON, an unexpected tool count -afterwards, and any panic. +afterward, and any panic. **What that does and does not promise.** It means the plugin's own failure modes cannot break a request — a bug or a surprising input forwards the original bytes diff --git a/docs/proposals/tool-prune.md b/docs/proposals/tool-prune.md index 019e84728..7c8659593 100644 --- a/docs/proposals/tool-prune.md +++ b/docs/proposals/tool-prune.md @@ -293,7 +293,7 @@ its 814 lines) into Go: launches the terminal UI. The change checks for a non-flag first argument before `flag.Parse()` and dispatches, falling through to the UI otherwise. -``` +```sh abctl tools scan [--days 30] [--keep Name,Name] [--write ] ``` @@ -464,7 +464,7 @@ Roughly 20 lines in the pane, 3 in `describePipeline`, 2 struct fields, and abou The operator reads something like: -``` +```text Metrics: requests seen 1284 count requests pruned 1284 count