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/README.md b/README.md index 27136e11c..3c2093f66 100644 --- a/README.md +++ b/README.md @@ -37,6 +37,15 @@ 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. 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 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/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/listener/extproc/server.go b/authbridge/authlib/listener/extproc/server.go index 8de88d384..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), @@ -657,7 +663,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 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 5100c4b13..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{WritesBody: 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/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/bridgehealth_test.go b/authbridge/authlib/listener/forwardproxy/bridgehealth_test.go new file mode 100644 index 000000000..de79503c7 --- /dev/null +++ b/authbridge/authlib/listener/forwardproxy/bridgehealth_test.go @@ -0,0 +1,151 @@ +package forwardproxy + +import ( + "strings" + "sync" + "testing" + + "github.com/rossoctl/cortex/authbridge/authlib/tlsbridge" +) + +// 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 + 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 attempts are normal (startup races)", + bridge: &tlsbridge.Engine{}, + tunnels: tunnelWarnThreshold - 1, + wantWarns: 0, + }, + { + name: "attempts but something was decrypted: bridge is working", + bridge: &tlsbridge.Engine{}, + tunnels: 50, + bridged: 1, + wantWarns: 0, + }, + { + name: "many attempts, 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.noteBridgeAttempt() + 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) + } +} + +// 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/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..f32eb97b5 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,22 @@ type Server struct { SkipHosts *skiphost.Matcher 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 + // 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 + bridgeAttempts atomic.Uint64 + bridgedRequests atomic.Uint64 + bridgeWarnOnce sync.Once + bridgeWarned atomic.Bool } // MTLSOptions configures outbound mTLS for the forward proxy. When @@ -211,6 +228,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, @@ -252,7 +272,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 { @@ -305,6 +325,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), @@ -352,7 +373,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 +421,21 @@ 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() { - // 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) + 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. + // 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. @@ -419,18 +447,20 @@ 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 - // 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) @@ -438,7 +468,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) @@ -531,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 } @@ -568,6 +603,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), @@ -880,6 +916,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, @@ -911,6 +948,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 @@ -1032,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 @@ -1220,3 +1259,68 @@ 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. 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) 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", + countKey, count, + "tunnels_opened", s.tunnelsOpened.Load(), + "bridged_requests", 0, + "likely_cause", cause, + "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/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_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 47d1c276f..8c7c84177 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,16 @@ func TestForwardProxy_Streaming_FramesFlowThrough(t *testing.T) { } } -// TestForwardProxy_Streaming_WritesBodyFallsBackToBuffered asserts the -// safety guard: a pipeline with a WritesBody 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_WritesBodyFallsBackToBuffered(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_WritesBodyFallsBackToBuffered(t *testing.T) { })) defer upstream.Close() - probe := newStreamingProbe(true) // WritesBody=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_WritesBodyFallsBackToBuffered(t *testing.T) { 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/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 15e108b5f..a58f7519e 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() { @@ -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, @@ -440,14 +441,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 + // 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")) && 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.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 @@ -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/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_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/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/bodydirection_test.go b/authbridge/authlib/pipeline/bodydirection_test.go new file mode 100644 index 000000000..3ede431c0 --- /dev/null +++ b/authbridge/authlib/pipeline/bodydirection_test.go @@ -0,0 +1,223 @@ +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) + } + }) + } +} + +// 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") + } +} + +// 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/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/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/context.go b/authbridge/authlib/pipeline/context.go index 789c9ec18..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 @@ -387,12 +391,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 WritesBody: 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 -// 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 +444,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/errorkind_test.go b/authbridge/authlib/pipeline/errorkind_test.go new file mode 100644 index 000000000..f7031c607 --- /dev/null +++ b/authbridge/authlib/pipeline/errorkind_test.go @@ -0,0 +1,113 @@ +package pipeline + +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 +// 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) + } +} + +// 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/holder.go b/authbridge/authlib/pipeline/holder.go index b5f8e018c..b27bab81e 100644 --- a/authbridge/authlib/pipeline/holder.go +++ b/authbridge/authlib/pipeline/holder.go @@ -81,11 +81,23 @@ 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 -// 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() } +// 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. +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/metrics.go b/authbridge/authlib/pipeline/metrics.go new file mode 100644 index 000000000..1cab473e6 --- /dev/null +++ b/authbridge/authlib/pipeline/metrics.go @@ -0,0 +1,42 @@ +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)". +// +// 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 +// 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/pipeline/pipeline.go b/authbridge/authlib/pipeline/pipeline.go index f46cf6de4..3f246e599 100644 --- a/authbridge/authlib/pipeline/pipeline.go +++ b/authbridge/authlib/pipeline/pipeline.go @@ -364,25 +364,77 @@ 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 { + 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.WritesBody { + // 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 } } 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 + } + } + 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 } } @@ -547,27 +599,97 @@ 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 + // 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.WritesBody { - 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()) + if caps.WritesRequestBody { + 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()) } - mutatorName = plugin.Name() - } else if caps.ReadsBody && mutatorName != "" && readerAfterMutator == "" { + 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()) + } + 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. + // + // 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() } } + 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, 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 } + +// 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/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..a924c6896 100644 --- a/authbridge/authlib/pipeline/plugin.go +++ b/authbridge/authlib/pipeline/plugin.go @@ -24,18 +24,35 @@ type PluginCapabilities struct { // a read silently sees "no body." ReadsBody bool - // WritesBody: 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 WritesBody - // plugin per direction — mutation ordering would be ambiguous. - // Waypoint mode (ext_authz) cannot support WritesBody at all: - // ext_authz has no body-mutation field. main.go enforces this at - // process boot. - WritesBody bool + // 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 @@ -71,12 +88,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.WritesResponseBody { c.ReadsBody = true } return c diff --git a/authbridge/authlib/pipeline/requestid.go b/authbridge/authlib/pipeline/requestid.go new file mode 100644 index 000000000..e23b82f39 --- /dev/null +++ b/authbridge/authlib/pipeline/requestid.go @@ -0,0 +1,59 @@ +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 { + 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 strconv.FormatUint(n, 36) + "-" + 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/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/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/authlib/pipeline/snapshot.go b/authbridge/authlib/pipeline/snapshot.go index 079cc36d0..e3d05bd81 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,75 @@ 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. +// +// 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 +// 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 "" + } + // 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 = stringOrNumber(gjson.GetBytes(body, "error.code")) + } + if t == "" { + return "" + } + if len(t) > 64 { + t = t[:64] + } + 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/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..1f4987d2b 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,11 @@ 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, + 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/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..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, - WritesBody: caps.WritesBody, - 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..05ba61a72 --- /dev/null +++ b/authbridge/authlib/plugins/registry_capsclone_test.go @@ -0,0 +1,117 @@ +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_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", + 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 f0a4ebd20..f99ab0852 100644 --- a/authbridge/authlib/plugins/sparc/plugin.go +++ b/authbridge/authlib/plugins/sparc/plugin.go @@ -211,8 +211,13 @@ func (p *SPARC) Capabilities() pipeline.PluginCapabilities { // 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.", + // 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 51c6cc02a..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.WritesBody || !caps.ReadsBody { - t.Error("expected ReadsBody+WritesBody") + 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/event.go b/authbridge/authlib/plugins/toolprune/event.go new file mode 100644 index 000000000..b531f919f --- /dev/null +++ b/authbridge/authlib/plugins/toolprune/event.go @@ -0,0 +1,85 @@ +package toolprune + +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. +// +// 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 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. + 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 +} + +// 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 new file mode 100644 index 000000000..a24ba73ef --- /dev/null +++ b/authbridge/authlib/plugins/toolprune/metrics.go @@ -0,0 +1,262 @@ +package toolprune + +import ( + "fmt" + "sort" + "strings" + "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 + + // 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 + + // 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 + + // 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() { + 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) observeSaving(tokens float64, t tier, usd float64, src rateSource, model string) { + m.mu.Lock() + switch t { + case tierCacheWrite: + m.savedCacheWrite += tokens + case tierCacheRead: + m.savedCacheRead += tokens + default: + m.savedInput += tokens + } + m.requestsCosted++ + if src != rateNone { + m.usdSaved += usd + if src == rateDefault { + m.usedDefaultRates = true + } + } else { + m.unpriced++ + if m.unpricedModels == nil { + m.unpricedModels = make(map[string]uint64) + } + if model == "" { + model = "(unknown)" + } + m.unpricedModels[model]++ + } + 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 { + out = append(out, pipeline.Metric{ + Name: "bytes removed / request", Value: float64(m.bytesRemoved) / float64(acted), Unit: "bytes", + }) + } + + // 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) + } + for _, t := range []struct { + name string + val float64 + }{ + {"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}) + } + } + if m.requestsCosted == 0 && acted > 0 { + out = append(out, pipeline.Metric{ + Name: "tokens saved", Value: 0, Unit: "tokens", + Note: "no response usage seen yet", + }) + } + + // Dollars, accumulated per request at that request's model rate. + 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. + // 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 + } + } + // 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: grossNote, + }) + } + } + + // 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, ", "), + }) + } + + // 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..9ce62d407 --- /dev/null +++ b/authbridge/authlib/plugins/toolprune/plugin.go @@ -0,0 +1,786 @@ +// 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, 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 ( + "bytes" + "context" + "encoding/json" + "fmt" + "log/slog" + "sort" + "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)."` + + // Pricing gives per-token rates per model. Rates are per model because they + // 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:"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:"-"` + // 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 + // cost is reported rather than a price being assumed. + // + // There is deliberately no output rate: pruning only ever shrinks the + // prompt, so attributing output cost to it would be false. + 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 { + // 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 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 + 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. +// +// 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, true + } + case tierCacheRead: + if r.CacheReadCostPerToken > 0 { + return r.CacheReadCostPerToken, true + } + } + return r.InputCostPerToken, r.InputCostPerToken > 0 +} + +func (r modelRates) set() bool { + return r.InputCostPerToken > 0 || r.CacheWriteCostPerToken > 0 || r.CacheReadCostPerToken > 0 +} + +// 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) + // 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 + } + // 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 + } + if c.flat.set() { + return c.flat, rateConfigured + } + return modelRates{}, rateNone +} + +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)) + globs := map[string]modelRates{} + // 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 + } + 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; +// everything else is read-only after Configure. +type ToolPrune struct { + cfg config + raw json.RawMessage + remove map[string]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{} } + +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() + if c.pricingErr != nil { + return fmt.Errorf("tool-prune config: %w", c.pricingErr) + } + + 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. +// +// 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 + } + } + 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 { + 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() +} + +// 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". +// +// 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 "", true // "auto" / "none" / "any" / "required" + } + 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 "", false +} + +// 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) { + // 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 + // 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. + 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 + } + // 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)) + 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. + 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) + } + } + if len(victims) == 0 { + // 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 + } + + 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 { + // 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 + } + // 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() { + 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 { + 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) + // 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)) + rateInput, _ := rates.rateFor(tierInput) + rateWrite, okW := rates.rateFor(tierCacheWrite) + rateRead, okR := rates.rateFor(tierCacheRead) + if !okW { + rateWrite = 0 + } + 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: bodySent, + Projected: !applied, + Model: inferenceModel(pctx), + RateInput: rateInput, + RateCacheWrite: rateWrite, + RateCacheRead: rateRead, + 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. + pipeline.SetState(pctx, p.Name(), &requestState{bytesRemoved: removedBytes}) + if applied { + 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} +} + +// 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) { + 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 + } + tokens := float64(st.bytesRemoved) * float64(promptTotal) / float64(len(pctx.Body)) + if tokens <= 0 { + return + } + t := tierOf(inf) + rates, src := p.cfg.ratesFor(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. +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 +// 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() { + p.driftChecked = true + 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..c3de5bbf8 --- /dev/null +++ b/authbridge/authlib/plugins/toolprune/plugin_test.go @@ -0,0 +1,1312 @@ +package toolprune + +import ( + "context" + "encoding/json" + "fmt" + "math" + "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()) + } +} + +// 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 +} + +// 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) + } +} + +// 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) + } + } + }) + } +} + +// 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, 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. + // 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) + } + } +} + +// 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, "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) + } + } +} + +// 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) { + cfg := func(t *testing.T) *ToolPrune { + p := New() + raw := []byte(`{"remove":["NotebookEdit"],` + + `"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 + } + + write := cfg(t) + finish(t, write, pruneOnce(t, write), 0, 0, 24701) + read := cfg(t) + finish(t, read, pruneOnce(t, read), 0, 24701, 0) + + 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 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) + } + // 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) + } +} + +// 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{} +} + +// 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) + } + }) + } +} + +// 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) + } +} + +// 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":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"} { + 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":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") + } + 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":1e-05}}}`) + pruneWithModel(t, p, "claude-opus-5") + if findMetric(t, p.Metrics(), "$ saved").Value <= 0 { + 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, "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) + } +} + +// 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) + } + } +} + +// 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.pricingErr != nil { + t.Fatalf("compile: %v", c.pricingErr) + } + 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.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 new file mode 100644 index 000000000..734abcb4a --- /dev/null +++ b/authbridge/authlib/plugins/toolprune/pricing.go @@ -0,0 +1,116 @@ +package toolprune + +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. +// +// 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. +// - Rates change. Nothing here refreshes them. +// +// Any `pricing` entry in config overrides the matching model. +var defaultPatterns = mustCompilePatterns(map[string]modelRates{ + // 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: 3.80 / tokensPerMillion, + CacheWriteCostPerToken: 4.75 / tokensPerMillion, // 1.25x input + CacheReadCostPerToken: 0.38 / tokensPerMillion, // 0.10x input + }, + "*claude-sonnet-*": { + InputCostPerToken: 1.52 / tokensPerMillion, + CacheWriteCostPerToken: 1.90 / tokensPerMillion, + CacheReadCostPerToken: 0.152 / tokensPerMillion, + }, + "*claude-haiku-*": { + InputCostPerToken: 0.76 / tokensPerMillion, + CacheWriteCostPerToken: 0.95 / tokensPerMillion, + CacheReadCostPerToken: 0.076 / tokensPerMillion, + }, +}) + +// 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 { + switch s { + case rateConfigured: + return "configured" + case rateDefault: + return "default" + default: + return "none" + } +} 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..90a4c71b7 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,13 @@ 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 { + // 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 } return out @@ -335,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/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/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/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..7aca6e019 --- /dev/null +++ b/authbridge/cmd/abctl/toolscan/patch.go @@ -0,0 +1,165 @@ +package toolscan + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "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 { + 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") + + 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 + } + // 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 + } + lines[i] = replacement + if err := writeFileAtomic(path, []byte(strings.Join(lines, "\n"))); 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) +} + +// 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) +} + +// 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 new file mode 100644 index 000000000..7389ea5a0 --- /dev/null +++ b/authbridge/cmd/abctl/toolscan/patch_test.go @@ -0,0 +1,247 @@ +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) + } +} + +// 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) + } +} + +// 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 new file mode 100644 index 000000000..488194779 --- /dev/null +++ b/authbridge/cmd/abctl/toolscan/scan.go @@ -0,0 +1,199 @@ +package toolscan + +import ( + "bufio" + "bytes" + "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++ + // Propagate: a partial scan silently proposes more tools. + 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. + // 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++ + + 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]++ + } + } + // 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 +} + +// 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 + // 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(" 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..2f8971ca5 --- /dev/null +++ b/authbridge/cmd/abctl/toolscan/scan_test.go @@ -0,0 +1,210 @@ +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) + } + // 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: []") { + t.Errorf("no candidates should render an empty list:\n%s", empty) + } +} diff --git a/authbridge/cmd/abctl/tui/app.go b/authbridge/cmd/abctl/tui/app.go index edfe1fa14..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 @@ -440,13 +444,18 @@ 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) 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) } @@ -583,11 +592,35 @@ 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. + // 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 + // 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/events_pane.go b/authbridge/cmd/abctl/tui/events_pane.go index 1fe21e487..c083f2fce 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" @@ -28,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), @@ -90,12 +91,10 @@ 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. + // 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, partner := computeEventPairs(eventRows) - glyphs := computeSpanGlyphs(partner, len(eventRows)) rows := make([]table.Row, 0, len(eventRows)) m.visibleRows = m.visibleRows[:0] @@ -121,15 +120,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"), @@ -140,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) @@ -477,21 +476,70 @@ 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 + } + 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 @@ -530,112 +578,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. @@ -812,3 +754,50 @@ func truncateScopes(scopes []string, n int) string { } return strings.Join(scopes[:n], ", ") + fmt.Sprintf(" +%d more", len(scopes)-n) } + +// tokensCellWithSaving renders the TOKENS / SAVED cell, splitting the two halves +// across the rows they actually belong to: +// +// - 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 { + 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 "" // no response yet: the ratio and tier are not known + } + resp := rows[j].event + 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 "" + } + return formatSavedOnly(tokens, usd, ps.RateSource, ps.Projected) +} diff --git a/authbridge/cmd/abctl/tui/events_pane_test.go b/authbridge/cmd/abctl/tui/events_pane_test.go index c8d6b281b..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) } } @@ -822,3 +709,141 @@ 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) + } +} + +// 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) + } + } + } +} 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 3bed7b03e..c775337ae 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)", @@ -80,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.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..9e18a323c --- /dev/null +++ b/authbridge/cmd/abctl/tui/plugin_metrics_test.go @@ -0,0 +1,123 @@ +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) + } +} + +// 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") + } +} diff --git a/authbridge/cmd/abctl/tui/prune_saving.go b/authbridge/cmd/abctl/tui/prune_saving.go new file mode 100644 index 000000000..490a0b1c4 --- /dev/null +++ b/authbridge/cmd/abctl/tui/prune_saving.go @@ -0,0 +1,125 @@ +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"` + // 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. +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 +} + +// 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. +// +// 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)) + } + 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..125c6d377 --- /dev/null +++ b/authbridge/cmd/abctl/tui/prune_saving_test.go @@ -0,0 +1,211 @@ +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") + } +} + +// 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", false) + for _, want := range []string{"−10.6k", "$0.050"} { + if !strings.Contains(got, want) { + t.Errorf("cell %q missing %q", got, want) + } + } + 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", 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", false) + 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) + } + } +} + +// 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") + } +} + +// 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/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..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" ) @@ -45,20 +47,53 @@ 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. + # 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: enforce + config: + remove: [] ` } -// 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 bb1a6faf1..31eae9df0 100644 --- a/authbridge/cmd/authbridge-proxy/demo_test.go +++ b/authbridge/cmd/authbridge-proxy/demo_test.go @@ -1,8 +1,10 @@ package main import ( + "os" "path/filepath" "slices" + "strings" "testing" "github.com/rossoctl/cortex/authbridge/authlib/config" @@ -72,8 +74,62 @@ 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. 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" { + tp = &cfg.Pipeline.Outbound.Plugins[i] + } + } + if tp == nil { + t.Fatal("tool-prune entry not found") + } + 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) + } +} + +// 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/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 +} 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/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..6265d3773 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 - WritesBody 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`. -`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 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 `WritesBody: 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,20 +584,31 @@ 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 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.** 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 | -| `WritesBody` | 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 | +| `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** | + +**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 `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 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: @@ -802,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` / `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` 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). @@ -821,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`, `WritesBody`. +- `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` / `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` / `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/laptop-token-savings.md b/authbridge/docs/laptop-token-savings.md new file mode 100644 index 000000000..b88f59a49 --- /dev/null +++ b/authbridge/docs/laptop-token-savings.md @@ -0,0 +1,154 @@ +# 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 +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 +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 + +```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. 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 + +```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`. + +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'`. + +## Seeing the saving in money + +`$ 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. + +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 + +`/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. diff --git a/authbridge/docs/plugin-catalog.md b/authbridge/docs/plugin-catalog.md index 450bbad7d..9fd977b63 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,25 @@ 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`. +- `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 +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 bd68a476e..abd6a4c6f 100644 --- a/authbridge/docs/plugin-reference.md +++ b/authbridge/docs/plugin-reference.md @@ -217,7 +217,8 @@ fail loud before serving traffic. ```go type PluginCapabilities struct { ReadsBody bool - WritesBody 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) @@ -616,6 +617,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`. @@ -684,7 +724,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 @@ -699,26 +739,65 @@ 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 + 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. -- `WritesBody`: implies `ReadsBody`. Listener propagates `pctx.SetBody` - rewrites to the upstream (and `pctx.SetResponseBody` to the - downstream client). +- `WritesRequestBody`: implies `ReadsBody`. Listener propagates `pctx.SetBody` + 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 | +| response mutator (`sparc`) | `WritesResponseBody` | no — buffered | +| both (`cpex`) | both | no — buffered | +| pure reader (parsers) | `ReadsBody` | yes | ### Build-time validation (enforced by `pipeline.New`) -- At most **one** `WritesBody` 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 - 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 - `WritesBody: true` plugins with `mode: waypoint`. + 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 +> 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 201e4923a..d51ca5fe9 100644 --- a/authbridge/docs/plugin-tutorial.md +++ b/authbridge/docs/plugin-tutorial.md @@ -209,12 +209,19 @@ 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 -`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 { - 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,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 `WritesBody` plugin per pipeline. Two mutators = ambiguous - ordering → build fails at startup. -- A `WritesBody` 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..72f0fab2c --- /dev/null +++ b/authbridge/docs/tool-prune-plugin.md @@ -0,0 +1,379 @@ +# `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 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] +``` + +| 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). + +## Turning it on + +**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 +``` + +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 + +`abctl`'s plugin detail pane shows a `Metrics:` section (source: +`GET /v1/pipeline`): + +```text +Metrics: + 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. + +### Per request, in the events timeline + +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. + +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 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 + +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. + +### 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 +rossoctl LiteLLM gateway, so `$ saved` appears with no configuration: + +| pattern | input | cache write (1.25x) | cache read (0.10x) | +|---|---|---|---| +| `*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 +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. + +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), +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: + "*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_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. + +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_million` 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 +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. 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 +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, 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. + +## Where the list comes from + +```sh +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: 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 +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 +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 + 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..3c87862b5 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: 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" +if [ -f "${demo_cfg}" ]; then + info " Cut tool-manifest waste (fills the remove: list; hot-reloaded, no restart):" + info " ${abctl_cmd} tools scan --write ${demo_cfg}" + 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}" info " Send traffic through it (e.g. Claude Code):" info " HTTPS_PROXY=http://localhost:${DEMO_FORWARD_PORT} \\" diff --git a/docs/proposals/tool-prune.md b/docs/proposals/tool-prune.md new file mode 100644 index 000000000..7c8659593 --- /dev/null +++ b/docs/proposals/tool-prune.md @@ -0,0 +1,567 @@ +# Directional Body Capabilities and the `tool-prune` Plugin + +**Status**: Draft +**Date**: September 2026 + +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 + 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. +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 + +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` | **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 | + +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, 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 + +```yaml +pipeline: + outbound: + plugins: + - inference-parser + - mcp-parser + - a2a-parser + - name: tool-prune + # on_error defaults to enforce; the empty remove list is the gate. + 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, 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` + +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. + +```sh +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. + +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 +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: + +```text +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 + +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. + 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. **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 + +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. + +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 | +|---|---| +| 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 | +| `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 + +None blocking. Two items deliberately deferred: + +- Adding the missing enforcement so `SetBody` matches its documented contract. + Needs its own compatibility review. +- 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.