From b763f924dfcd2d0670d7b8181c15ee88e61d560a Mon Sep 17 00:00:00 2001 From: Maycon Santos Date: Sun, 28 Jun 2026 14:06:47 +0200 Subject: [PATCH] [management,proxy] Reduce agent-network cognitive complexity (#6566) Address the SonarCloud quality-gate findings in new agent-network code by extracting focused helpers. No behavior change. - synthesizer.go: split buildIdentityInjectConfigJSON into per-shape rule builders; extract mergeGuardrail from mergeGuardrails to cut nesting depth. - llm_identity_inject: extract injectionEmitsAnything validation predicate from New. - llm_response_parser/streaming.go: extract applyOpenAIStreamUsage and applyAnthropicStreamUsage (via a named anthropicStreamUsage type) and simplify the OpenAI scanner loop. - reverseproxy.go: decompose ServeHTTP into serveRouteError, buildTargetContext, serveDirect, serveWithChain, captureRequestForChain, serveDeny, newResponseWriter, observeResponse, and forwardUpstream, preserving the defer ordering so response observation still reads the captured writer before it is released. --- .../modules/agentnetwork/synthesizer.go | 224 +++++++------ .../builtin/llm_identity_inject/middleware.go | 67 ++-- .../builtin/llm_response_parser/streaming.go | 125 +++---- proxy/internal/proxy/reverseproxy.go | 313 ++++++++++-------- 4 files changed, 402 insertions(+), 327 deletions(-) diff --git a/management/internals/modules/agentnetwork/synthesizer.go b/management/internals/modules/agentnetwork/synthesizer.go index 3e59c1146..9814d1a11 100644 --- a/management/internals/modules/agentnetwork/synthesizer.go +++ b/management/internals/modules/agentnetwork/synthesizer.go @@ -558,84 +558,8 @@ func buildIdentityInjectConfigJSON(providers []*types.Provider, groupIndex map[s if !ok { continue } - rule := identityInjectProvider{ProviderID: p.ID} - // Identity-stamping shape (one of HeaderPair / JSONMetadata). - // Skip the shape silently when the catalog entry doesn't - // declare one — extras can still apply, see below. - if entry.IdentityInjection != nil { - switch { - case entry.IdentityInjection.HeaderPair != nil: - hp := entry.IdentityInjection.HeaderPair - // For Customizable shapes (Bifrost today) the wire - // header names come from the provider record verbatim; - // the catalog values are placeholder defaults shown by - // the dashboard, not authoritative. Empty operator - // value disables stamping for that dimension — - // applyHeaderPair already no-ops on empty header - // names. The body-inject flags stay catalog-owned - // because Customizable=true today only applies to - // gateways that read identity from headers (the - // flags would be no-ops anyway). - userHeader := hp.EndUserIDHeader - tagsHeader := hp.TagsHeader - if hp.Customizable { - userHeader = p.IdentityHeaderUserID - tagsHeader = p.IdentityHeaderGroups - } - if userHeader != "" || tagsHeader != "" || hp.TagsInBody || hp.EndUserIDInBody { - rule.HeaderPair = &identityInjectHeaderPair{ - EndUserIDHeader: userHeader, - TagsHeader: tagsHeader, - TagsInBody: hp.TagsInBody, - EndUserIDInBody: hp.EndUserIDInBody, - } - } - case entry.IdentityInjection.JSONMetadata != nil: - jm := entry.IdentityInjection.JSONMetadata - // Customizable JSONMetadata reuses the same provider- - // record fields HeaderPair uses — IdentityHeaderUserID - // becomes the JSON key for the user dimension, and - // IdentityHeaderGroups becomes the JSON key for groups. - // Empty operator value is honored as "skip this key"; - // applyJSONMetadata already drops keys with empty - // names. Header itself is catalog-owned (e.g. - // cf-aig-metadata) — operators only override the keys - // inside the JSON, not the wire header that carries it. - userKey := jm.UserKey - groupsKey := jm.GroupsKey - if jm.Customizable { - userKey = p.IdentityHeaderUserID - groupsKey = p.IdentityHeaderGroups - } - if jm.Header != "" { - rule.JSONMetadata = &identityInjectJSONMetadata{ - Header: jm.Header, - UserKey: userKey, - GroupsKey: groupsKey, - MaxValueLength: jm.MaxValueLength, - } - } - } - } - // Extra catalog-declared static headers (e.g. Portkey config - // id). Only emit entries whose value the operator has filled - // in on the provider record; missing/empty values are no-ops. - for _, h := range entry.ExtraHeaders { - if h.Name == "" { - continue - } - v := strings.TrimSpace(p.ExtraValues[h.Name]) - if v == "" { - continue - } - rule.ExtraHeaders = append(rule.ExtraHeaders, identityInjectExtraHeader{ - Name: h.Name, - Value: v, - }) - } - // If this provider would emit nothing, skip it entirely so - // the middleware doesn't carry an inert rule for it. - if rule.HeaderPair == nil && rule.JSONMetadata == nil && len(rule.ExtraHeaders) == 0 { + rule, ok := buildIdentityInjectRule(p, entry) + if !ok { continue } cfg.Providers = append(cfg.Providers, rule) @@ -647,6 +571,102 @@ func buildIdentityInjectConfigJSON(providers []*types.Provider, groupIndex map[s return out, nil } +// buildIdentityInjectRule assembles the injection rule for one provider +// from its record and catalog entry. The second return is false when the +// provider would emit nothing, so the caller can skip it entirely rather +// than carry an inert rule for it. +func buildIdentityInjectRule(p *types.Provider, entry catalog.Provider) (identityInjectProvider, bool) { + rule := identityInjectProvider{ProviderID: p.ID} + // Identity-stamping shape (one of HeaderPair / JSONMetadata). Skip the + // shape silently when the catalog entry doesn't declare one — extras + // can still apply, see below. + if entry.IdentityInjection != nil { + switch { + case entry.IdentityInjection.HeaderPair != nil: + rule.HeaderPair = buildIdentityHeaderPair(p, entry.IdentityInjection.HeaderPair) + case entry.IdentityInjection.JSONMetadata != nil: + rule.JSONMetadata = buildIdentityJSONMetadata(p, entry.IdentityInjection.JSONMetadata) + } + } + rule.ExtraHeaders = buildIdentityExtraHeaders(p, entry.ExtraHeaders) + if rule.HeaderPair == nil && rule.JSONMetadata == nil && len(rule.ExtraHeaders) == 0 { + return identityInjectProvider{}, false + } + return rule, true +} + +// buildIdentityHeaderPair resolves the header-pair injection shape, +// returning nil when nothing would be stamped. For Customizable shapes +// (Bifrost today) the wire header names come from the provider record +// verbatim; the catalog values are placeholder defaults shown by the +// dashboard, not authoritative. Empty operator value disables stamping +// for that dimension — applyHeaderPair already no-ops on empty header +// names. The body-inject flags stay catalog-owned because Customizable +// today only applies to gateways that read identity from headers (the +// flags would be no-ops anyway). +func buildIdentityHeaderPair(p *types.Provider, hp *catalog.HeaderPairInjection) *identityInjectHeaderPair { + userHeader := hp.EndUserIDHeader + tagsHeader := hp.TagsHeader + if hp.Customizable { + userHeader = p.IdentityHeaderUserID + tagsHeader = p.IdentityHeaderGroups + } + if userHeader == "" && tagsHeader == "" && !hp.TagsInBody && !hp.EndUserIDInBody { + return nil + } + return &identityInjectHeaderPair{ + EndUserIDHeader: userHeader, + TagsHeader: tagsHeader, + TagsInBody: hp.TagsInBody, + EndUserIDInBody: hp.EndUserIDInBody, + } +} + +// buildIdentityJSONMetadata resolves the JSON-metadata injection shape, +// returning nil when the catalog entry carries no header. Customizable +// JSONMetadata reuses the same provider-record fields HeaderPair uses — +// IdentityHeaderUserID becomes the JSON key for the user dimension, and +// IdentityHeaderGroups becomes the JSON key for groups. Empty operator +// value is honored as "skip this key"; applyJSONMetadata already drops +// keys with empty names. Header itself is catalog-owned (e.g. +// cf-aig-metadata) — operators only override the keys inside the JSON, +// not the wire header that carries it. +func buildIdentityJSONMetadata(p *types.Provider, jm *catalog.JSONMetadataInjection) *identityInjectJSONMetadata { + if jm.Header == "" { + return nil + } + userKey := jm.UserKey + groupsKey := jm.GroupsKey + if jm.Customizable { + userKey = p.IdentityHeaderUserID + groupsKey = p.IdentityHeaderGroups + } + return &identityInjectJSONMetadata{ + Header: jm.Header, + UserKey: userKey, + GroupsKey: groupsKey, + MaxValueLength: jm.MaxValueLength, + } +} + +// buildIdentityExtraHeaders collects catalog-declared static headers (e.g. +// Portkey config id), emitting only entries whose value the operator has +// filled in on the provider record; missing/empty values are no-ops. +func buildIdentityExtraHeaders(p *types.Provider, extras []catalog.ExtraHeader) []identityInjectExtraHeader { + var out []identityInjectExtraHeader + for _, h := range extras { + if h.Name == "" { + continue + } + v := strings.TrimSpace(p.ExtraValues[h.Name]) + if v == "" { + continue + } + out = append(out, identityInjectExtraHeader{Name: h.Name, Value: v}) + } + return out +} + // buildMiddlewareChain assembles the per-target middleware chain that // implements the Agent Network behaviour at the proxy. Slot order on // the request leg is the slice order; on the response leg it runs in @@ -1023,28 +1043,7 @@ func mergeGuardrails(policies []*types.Policy, byID map[string]*types.Guardrail) if !ok || g == nil { continue } - if g.Checks.ModelAllowlist.Enabled { - allowlistEnabled = true - for _, m := range g.Checks.ModelAllowlist.Models { - if m != "" { - allowlist[m] = struct{}{} - } - } - } - if g.Checks.PromptCapture.Enabled { - merged.PromptCapture.Enabled = true - if g.Checks.PromptCapture.RedactPii { - merged.PromptCapture.RedactPii = true - } - } - // TokenLimits, Budget, and Retention have moved off - // guardrails. Token and budget caps now live on the - // Policy itself (Policy.Limits); retention will move to - // account-level Settings. The proxy still consumes - // MergedTokenLimits/MergedBudget/MergedRetention from - // this struct for backwards compatibility, but they - // remain at zero/disabled until the new Policy.Limits - // merge is wired in by the enforcement track. + mergeGuardrail(g, &merged, allowlist, &allowlistEnabled) } } @@ -1057,3 +1056,28 @@ func mergeGuardrails(policies []*types.Policy, byID map[string]*types.Guardrail) } return merged } + +// mergeGuardrail folds a single guardrail's enabled checks into the +// running merge: model-allowlist models join the shared set (and flip +// allowlistEnabled), and prompt-capture / redact-pii stick once any +// enabling guardrail turns them on. +// +// TokenLimits, Budget, and Retention have moved off guardrails — token +// and budget caps now live on the Policy itself (Policy.Limits) and +// retention moves to account-level Settings — so they are not merged here. +func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails, allowlist map[string]struct{}, allowlistEnabled *bool) { + if g.Checks.ModelAllowlist.Enabled { + *allowlistEnabled = true + for _, m := range g.Checks.ModelAllowlist.Models { + if m != "" { + allowlist[m] = struct{}{} + } + } + } + if g.Checks.PromptCapture.Enabled { + merged.PromptCapture.Enabled = true + if g.Checks.PromptCapture.RedactPii { + merged.PromptCapture.RedactPii = true + } + } +} diff --git a/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go b/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go index 96f686c0e..ee3f1c20d 100644 --- a/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go +++ b/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go @@ -40,8 +40,8 @@ const Version = "1.0.0" // Middleware stamps NetBird identity onto upstream requests for the // configured set of resolved providers. type Middleware struct { - cfg Config - byID map[string]ProviderInjection + cfg Config + byID map[string]ProviderInjection } // New constructs a Middleware from the supplied configuration. A nil @@ -49,45 +49,44 @@ type Middleware struct { func New(cfg Config) *Middleware { byID := make(map[string]ProviderInjection, len(cfg.Providers)) for _, p := range cfg.Providers { - if p.ProviderID == "" { + if p.ProviderID == "" || !injectionEmitsAnything(p) { continue } - // Drop entries that wouldn't inject anything — keeps the - // runtime check tight. Also drop entries that set both - // shapes (configuration error; refuse to guess which wins). - // Extras alone are enough to keep the rule alive even if - // neither identity shape is set. - hasExtras := false - for _, e := range p.ExtraHeaders { - if e.Name != "" && e.Value != "" { - hasExtras = true - break - } - } - switch { - case p.HeaderPair != nil && p.JSONMetadata != nil: - continue - case p.HeaderPair != nil: - if p.HeaderPair.EndUserIDHeader == "" && p.HeaderPair.TagsHeader == "" && !p.HeaderPair.TagsInBody && !p.HeaderPair.EndUserIDInBody && !hasExtras { - continue - } - case p.JSONMetadata != nil: - if p.JSONMetadata.Header == "" { - continue - } - if p.JSONMetadata.UserKey == "" && p.JSONMetadata.GroupsKey == "" && !hasExtras { - continue - } - default: - if !hasExtras { - continue - } - } byID[p.ProviderID] = p } return &Middleware{cfg: cfg, byID: byID} } +// injectionEmitsAnything reports whether a provider injection rule would +// stamp anything at runtime. Rules that set both identity shapes are a +// configuration error (we refuse to guess which wins), and rules that +// resolve to no headers are dropped to keep the runtime check tight. +// Non-empty extras alone keep a rule alive even when neither identity +// shape is set. +func injectionEmitsAnything(p ProviderInjection) bool { + hasExtras := false + for _, e := range p.ExtraHeaders { + if e.Name != "" && e.Value != "" { + hasExtras = true + break + } + } + switch { + case p.HeaderPair != nil && p.JSONMetadata != nil: + return false + case p.HeaderPair != nil: + return p.HeaderPair.EndUserIDHeader != "" || p.HeaderPair.TagsHeader != "" || + p.HeaderPair.TagsInBody || p.HeaderPair.EndUserIDInBody || hasExtras + case p.JSONMetadata != nil: + if p.JSONMetadata.Header == "" { + return false + } + return p.JSONMetadata.UserKey != "" || p.JSONMetadata.GroupsKey != "" || hasExtras + default: + return hasExtras + } +} + // ID returns the registry identifier. func (m *Middleware) ID() string { return ID } diff --git a/proxy/internal/middleware/builtin/llm_response_parser/streaming.go b/proxy/internal/middleware/builtin/llm_response_parser/streaming.go index 9d3e8a204..6462ab4ae 100644 --- a/proxy/internal/middleware/builtin/llm_response_parser/streaming.go +++ b/proxy/internal/middleware/builtin/llm_response_parser/streaming.go @@ -84,15 +84,12 @@ func accumulateOpenAIStream(body []byte) (llm.Usage, string) { for { ev, err := scanner.Next() if err != nil { - if errors.Is(err, io.EOF) { - break - } break } - if ev.Data == "" || ev.Data == openAIDoneSentinel { - if ev.Data == openAIDoneSentinel { - break - } + if ev.Data == openAIDoneSentinel { + break + } + if ev.Data == "" { continue } @@ -113,26 +110,35 @@ func accumulateOpenAIStream(body []byte) (llm.Usage, string) { if u == nil && chunk.Response != nil { u = chunk.Response.Usage } - if u != nil { - usage.InputTokens = pickInt64(u.InputTokens, u.PromptTokens) - usage.OutputTokens = pickInt64(u.OutputTokens, u.CompletionTokens) - usage.TotalTokens = derefInt64(u.TotalTokens) - if u.InputTokensDetails != nil { - if v := derefInt64(u.InputTokensDetails.CachedTokens); v > 0 { - usage.CachedInputTokens = v - } - } - if usage.CachedInputTokens == 0 && u.PromptTokensDetails != nil { - usage.CachedInputTokens = derefInt64(u.PromptTokensDetails.CachedTokens) - } - if usage.TotalTokens == 0 && (usage.InputTokens > 0 || usage.OutputTokens > 0) { - usage.TotalTokens = usage.InputTokens + usage.OutputTokens - } - } + applyOpenAIStreamUsage(u, &usage) } return usage, completion.String() } +// applyOpenAIStreamUsage lifts the token counts off a final-frame usage +// block into the running usage, normalising the chat.completions +// (prompt_/completion_) and Responses-API (input_/output_) names and +// backfilling total tokens when the provider omits them. +func applyOpenAIStreamUsage(u *openAIStreamUsage, usage *llm.Usage) { + if u == nil { + return + } + usage.InputTokens = pickInt64(u.InputTokens, u.PromptTokens) + usage.OutputTokens = pickInt64(u.OutputTokens, u.CompletionTokens) + usage.TotalTokens = derefInt64(u.TotalTokens) + if u.InputTokensDetails != nil { + if v := derefInt64(u.InputTokensDetails.CachedTokens); v > 0 { + usage.CachedInputTokens = v + } + } + if usage.CachedInputTokens == 0 && u.PromptTokensDetails != nil { + usage.CachedInputTokens = derefInt64(u.PromptTokensDetails.CachedTokens) + } + if usage.TotalTokens == 0 && (usage.InputTokens > 0 || usage.OutputTokens > 0) { + usage.TotalTokens = usage.InputTokens + usage.OutputTokens + } +} + // decodeJSONString unmarshals a JSON-encoded string value, returning // ok=false when the raw message is empty or not a string. func decodeJSONString(raw json.RawMessage) (string, bool) { @@ -149,26 +155,23 @@ func decodeJSONString(raw json.RawMessage) (string, bool) { // anthropicStreamEvent captures the union of Messages-API stream event // payloads we care about. Each named event on the wire fills only its // shape's fields; unknown keys are ignored. +type anthropicStreamUsage struct { + InputTokens *int64 `json:"input_tokens"` + OutputTokens *int64 `json:"output_tokens"` + CacheReadInputTokens *int64 `json:"cache_read_input_tokens"` + CacheCreationInputTokens *int64 `json:"cache_creation_input_tokens"` +} + type anthropicStreamEvent struct { Type string `json:"type"` Message *struct { - Usage *struct { - InputTokens *int64 `json:"input_tokens"` - OutputTokens *int64 `json:"output_tokens"` - CacheReadInputTokens *int64 `json:"cache_read_input_tokens"` - CacheCreationInputTokens *int64 `json:"cache_creation_input_tokens"` - } `json:"usage"` + Usage *anthropicStreamUsage `json:"usage"` } `json:"message"` Delta *struct { Type string `json:"type"` Text string `json:"text"` } `json:"delta"` - Usage *struct { - InputTokens *int64 `json:"input_tokens"` - OutputTokens *int64 `json:"output_tokens"` - CacheReadInputTokens *int64 `json:"cache_read_input_tokens"` - CacheCreationInputTokens *int64 `json:"cache_creation_input_tokens"` - } `json:"usage"` + Usage *anthropicStreamUsage `json:"usage"` } // accumulateAnthropicStream tracks input_tokens from message_start, @@ -216,44 +219,42 @@ func accumulateAnthropicStream(body []byte) (llm.Usage, string) { func applyAnthropicStreamEvent(eventType string, payload anthropicStreamEvent, usage *llm.Usage, completion *strings.Builder) { switch eventType { case "message_start": - if payload.Message != nil && payload.Message.Usage != nil { - if v := derefInt64(payload.Message.Usage.InputTokens); v > 0 { - usage.InputTokens = v - } - if v := derefInt64(payload.Message.Usage.OutputTokens); v > 0 { - usage.OutputTokens = v - } - if v := derefInt64(payload.Message.Usage.CacheReadInputTokens); v > 0 { - usage.CachedInputTokens = v - } - if v := derefInt64(payload.Message.Usage.CacheCreationInputTokens); v > 0 { - usage.CacheCreationTokens = v - } + if payload.Message != nil { + applyAnthropicStreamUsage(payload.Message.Usage, usage) } case "content_block_delta": if payload.Delta != nil && payload.Delta.Type == "text_delta" { completion.WriteString(payload.Delta.Text) } case "message_delta": - if payload.Usage != nil { - if v := derefInt64(payload.Usage.InputTokens); v > 0 { - usage.InputTokens = v - } - if v := derefInt64(payload.Usage.OutputTokens); v > 0 { - usage.OutputTokens = v - } - if v := derefInt64(payload.Usage.CacheReadInputTokens); v > 0 { - usage.CachedInputTokens = v - } - if v := derefInt64(payload.Usage.CacheCreationInputTokens); v > 0 { - usage.CacheCreationTokens = v - } - } + applyAnthropicStreamUsage(payload.Usage, usage) case "message_stop": // No-op; Anthropic does not emit usage here. } } +// applyAnthropicStreamUsage folds a non-nil Anthropic usage block into the +// running totals. Each field overwrites only when present and positive, so +// message_delta's post-completion counts supersede the message_start seed +// without zeroing dimensions a later event omits. +func applyAnthropicStreamUsage(u *anthropicStreamUsage, usage *llm.Usage) { + if u == nil { + return + } + if v := derefInt64(u.InputTokens); v > 0 { + usage.InputTokens = v + } + if v := derefInt64(u.OutputTokens); v > 0 { + usage.OutputTokens = v + } + if v := derefInt64(u.CacheReadInputTokens); v > 0 { + usage.CachedInputTokens = v + } + if v := derefInt64(u.CacheCreationInputTokens); v > 0 { + usage.CacheCreationTokens = v + } +} + func pickInt64(preferred, fallback *int64) int64 { if preferred != nil { return *preferred diff --git a/proxy/internal/proxy/reverseproxy.go b/proxy/internal/proxy/reverseproxy.go index c64d0fa26..2c0304ecd 100644 --- a/proxy/internal/proxy/reverseproxy.go +++ b/proxy/internal/proxy/reverseproxy.go @@ -83,13 +83,8 @@ func NewReverseProxy(transport http.RoundTripper, forwardedProto string, trusted func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { result, exists := p.findTargetForRequest(r) if !exists { - if cd := CapturedDataFromContext(r.Context()); cd != nil { - cd.SetOrigin(OriginNoRoute) - } - requestID := getRequestID(r) - web.ServeErrorPage(w, r, http.StatusNotFound, "Service Not Found", - "The requested service could not be found. Please check the URL, try refreshing, or check if the peer is running. If that doesn't work, see our documentation for help.", - requestID, web.ErrorStatus{Proxy: true, Destination: false}) + p.serveRouteError(w, r, http.StatusNotFound, "Service Not Found", + "The requested service could not be found. Please check the URL, try refreshing, or check if the peer is running. If that doesn't work, see our documentation for help.") return } @@ -99,19 +94,13 @@ func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { // with 421 (Misdirected Request) so the caller sees an explicit // error instead of silently doubling tunnel traffic. if p.isSelfTargetLoop(r, result.target.URL) { - if cd := CapturedDataFromContext(r.Context()); cd != nil { - cd.SetOrigin(OriginNoRoute) - } - requestID := getRequestID(r) - web.ServeErrorPage(w, r, http.StatusMisdirectedRequest, "Loop Detected", - "This peer is the target of the requested service. Reach the backend directly instead of dialing the public service URL from the same machine.", - requestID, web.ErrorStatus{Proxy: true, Destination: false}) + p.serveRouteError(w, r, http.StatusMisdirectedRequest, "Loop Detected", + "This peer is the target of the requested service. Reach the backend directly instead of dialing the public service URL from the same machine.") return } - ctx := r.Context() - // Set the account ID in the context for the roundtripper to use. - ctx = roundtrip.WithAccountID(ctx, result.accountID) + pt := result.target + ctx := p.buildTargetContext(r.Context(), result) // Populate captured data if it exists (allows middleware to read after handler completes). // This solves the problem of passing data UP the middleware chain: we put a mutable struct @@ -124,8 +113,34 @@ func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { capturedData.SetSuppressAccessLog(result.target != nil && result.target.DisableAccessLog) } - pt := result.target + rewriteMatchedPath := result.matchedPath + if pt.PathRewrite == PathRewritePreserve { + rewriteMatchedPath = "" + } + chain := p.resolveChain(result) + if chain == nil || chain.Empty() { + p.serveDirect(w, r, ctx, result, rewriteMatchedPath) + return + } + p.serveWithChain(w, r, ctx, result, chain, rewriteMatchedPath, capturedData) +} + +// serveRouteError marks the request as un-routed on any captured-data +// context and renders the proxy error page. +func (p *ReverseProxy) serveRouteError(w http.ResponseWriter, r *http.Request, status int, title, message string) { + if cd := CapturedDataFromContext(r.Context()); cd != nil { + cd.SetOrigin(OriginNoRoute) + } + web.ServeErrorPage(w, r, status, title, message, getRequestID(r), + web.ErrorStatus{Proxy: true, Destination: false}) +} + +// buildTargetContext layers the per-target roundtrip flags (account id, +// TLS-verify skip, direct upstream, dial timeout) onto the request context. +func (p *ReverseProxy) buildTargetContext(ctx context.Context, result targetResult) context.Context { + pt := result.target + ctx = roundtrip.WithAccountID(ctx, result.accountID) if pt.SkipTLSVerify { ctx = roundtrip.WithSkipTLSVerify(ctx) } @@ -135,32 +150,66 @@ func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { if pt.RequestTimeout > 0 { ctx = types.WithDialTimeout(ctx, pt.RequestTimeout) } + return ctx +} - rewriteMatchedPath := result.matchedPath - if pt.PathRewrite == PathRewritePreserve { - rewriteMatchedPath = "" +// serveDirect forwards the request without a middleware chain — the common +// path for plain reverse-proxy targets. +func (p *ReverseProxy) serveDirect(w http.ResponseWriter, r *http.Request, ctx context.Context, result targetResult, rewriteMatchedPath string) { + pt := result.target + rp := &httputil.ReverseProxy{ + Rewrite: p.rewriteFunc(pt.URL, rewriteMatchedPath, result.passHostHeader, pt.PathRewrite, pt.CustomHeaders, result.stripAuthHeaders), + Transport: p.transport, + FlushInterval: -1, + ErrorHandler: p.proxyErrorHandler, } - - chain := p.resolveChain(result) - if chain == nil || chain.Empty() { - rp := &httputil.ReverseProxy{ - Rewrite: p.rewriteFunc(pt.URL, rewriteMatchedPath, result.passHostHeader, pt.PathRewrite, pt.CustomHeaders, result.stripAuthHeaders), - Transport: p.transport, - FlushInterval: -1, - ErrorHandler: p.proxyErrorHandler, - } - if result.rewriteRedirects { - rp.ModifyResponse = p.rewriteLocationFunc(pt.URL, rewriteMatchedPath, r) //nolint:bodyclose - } - rp.ServeHTTP(w, r.WithContext(ctx)) - return + if result.rewriteRedirects { + rp.ModifyResponse = p.rewriteLocationFunc(pt.URL, rewriteMatchedPath, r) //nolint:bodyclose } + rp.ServeHTTP(w, r.WithContext(ctx)) +} +// serveWithChain runs the per-target middleware chain around the upstream +// request: request-leg capture and authorisation, then (on allow) the +// upstream forward with response/terminal observation deferred so it reads +// the captured response before the writer is released. +func (p *ReverseProxy) serveWithChain(w http.ResponseWriter, r *http.Request, ctx context.Context, result targetResult, chain *middleware.Chain, rewriteMatchedPath string, capturedData *CapturedData) { middlewareIDs := chain.IDs() p.logger.Debugf("middleware chain matched: service=%s path=%s middlewares=%v", result.serviceID, result.matchedPath, middlewareIDs) - capturedBody, truncated, originalSize, bypass, releaseBudget, captureErr := bodytap.CaptureRequest(r, pt.CaptureConfig, p.middlewareManager.Budget()) + capturedBody, truncated, originalSize, releaseBudget := p.captureRequestForChain(ctx, r, result, capturedData) defer releaseBudget() + + acc := middleware.NewAccumulator(middleware.MaxRequestMetadataBytes) + reqInput := buildRequestInput(r, result, capturedData, capturedBody, truncated, originalSize) + + denyOutput, requestMeta, upstreamRewrite, _ := chain.RunRequest(ctx, r, reqInput, acc) + if capturedData != nil { + for _, kv := range requestMeta { + capturedData.SetMetadata(kv.Key, kv.Value) + } + } + if denyOutput != nil { + p.serveDeny(w, denyOutput, result, middlewareIDs) + return + } + + respWriter, capturingWriter := p.newResponseWriter(ctx, w, result, capturedData) + if capturingWriter != nil { + defer capturingWriter.Release() + defer p.observeResponse(ctx, chain, acc, reqInput, requestMeta, capturingWriter, w, capturedData, result, middlewareIDs) + } + + p.forwardUpstream(respWriter, r, ctx, result, rewriteMatchedPath, upstreamRewrite) +} + +// captureRequestForChain copies the request body for inspection by the +// chain, records any capture bypass, and applies agent-network routing +// recovery for oversized bodies. The returned release frees the capture +// budget and must be deferred by the caller. +func (p *ReverseProxy) captureRequestForChain(ctx context.Context, r *http.Request, result targetResult, capturedData *CapturedData) ([]byte, bool, int64, func()) { + pt := result.target + capturedBody, truncated, originalSize, bypass, releaseBudget, captureErr := bodytap.CaptureRequest(r, pt.CaptureConfig, p.middlewareManager.Budget()) if captureErr != nil { p.logger.Debugf("middleware request body capture error: %v", captureErr) } @@ -184,112 +233,114 @@ func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { p.logger.Debugf("agent-network routing recovery: extracted model=%s stream=%t from oversized request body (service=%s)", model, stream, result.serviceID) } } + return capturedBody, truncated, originalSize, releaseBudget +} - acc := middleware.NewAccumulator(middleware.MaxRequestMetadataBytes) - reqInput := buildRequestInput(r, result, capturedData, capturedBody, truncated, originalSize) +// serveDeny renders the chain's deny response. Policy/budget/routing/guardrail +// denials are expected runtime outcomes and can be high-volume under +// misconfigured or hostile clients; per-request detail stays at Debug and +// metrics/access logs carry the signal at scale. +func (p *ReverseProxy) serveDeny(w http.ResponseWriter, denyOutput *middleware.Output, result targetResult, middlewareIDs []string) { + middlewareID := "middleware" + if denyOutput.DenyReason != nil && denyOutput.DenyReason.Code != "" { + middlewareID = denyOutput.DenyReason.Code + } + p.logger.Debugf("middleware chain denied request: service=%s path=%s middlewares=%v reason=%s status=%d", + result.serviceID, result.matchedPath, middlewareIDs, middlewareID, denyOutput.DenyStatus) + middleware.RenderDenyResponse(w, middlewareID, denyOutput.DenyReason, denyOutput.DenyStatus) +} - denyOutput, requestMeta, upstreamRewrite, _ := chain.RunRequest(ctx, r, reqInput, acc) +// newResponseWriter returns the writer the upstream forward should use. When +// response capture is enabled and not bypassed it wraps w in a capturing +// writer (also returned so the caller can release it and feed the response +// leg); otherwise the capturing writer is nil and w is used directly. +func (p *ReverseProxy) newResponseWriter(ctx context.Context, w http.ResponseWriter, result targetResult, capturedData *CapturedData) (http.ResponseWriter, *bodytap.CapturingResponseWriter) { + pt := result.target + if pt.CaptureConfig == nil || pt.CaptureConfig.MaxResponseBytes <= 0 { + return w, nil + } + capturingWriter := bodytap.NewCapturingResponseWriter(w, pt.CaptureConfig.MaxResponseBytes, p.middlewareManager.Budget()) + if capturingWriter.Bypassed() { + if capturedData != nil { + capturedData.SetMetadata("mw.capture.bypass_reason", capturingWriter.BypassReason()) + } + p.middlewareManager.Metrics().IncCaptureBypass(ctx, string(result.serviceID), capturingWriter.BypassReason()) + capturingWriter.Release() + return w, nil + } + return capturingWriter, capturingWriter +} + +// observeResponse runs the response and terminal middleware slots after the +// body has been forwarded. It is deferred by serveWithChain so it reads the +// captured response before the writer is released. +func (p *ReverseProxy) observeResponse(ctx context.Context, chain *middleware.Chain, acc *middleware.Accumulator, reqInput *middleware.Input, requestMeta []middleware.KV, capturingWriter *bodytap.CapturingResponseWriter, w http.ResponseWriter, capturedData *CapturedData, result targetResult, middlewareIDs []string) { + respInput := &middleware.Input{ + Slot: middleware.SlotOnResponse, + RequestID: reqInput.RequestID, + TargetID: reqInput.TargetID, + Method: reqInput.Method, + URL: reqInput.URL, + Headers: reqInput.Headers, + Status: capturingWriter.Status(), + RespHeaders: headerToKV(w.Header()), + RespBody: capturingWriter.Body(), + RespBodyTruncated: capturingWriter.Truncated(), + OriginalRespSize: capturingWriter.BytesWritten(), + ServiceID: reqInput.ServiceID, + AccountID: reqInput.AccountID, + UserID: reqInput.UserID, + // UserEmail / UserGroups / UserGroupNames must flow into the + // response leg too — llm_limit_record needs UserGroups to send + // group_ids on RecordLLMUsage so management's account-budget + // fan-out can match group-targeted rules; identity-stamping and + // any future response-side authorisation also depend on these. + UserEmail: reqInput.UserEmail, + UserGroups: reqInput.UserGroups, + UserGroupNames: reqInput.UserGroupNames, + AuthMethod: reqInput.AuthMethod, + SourceIP: reqInput.SourceIP, + Metadata: requestMeta, + AgentNetwork: reqInput.AgentNetwork, + } + // The response/terminal phase runs after the body is forwarded, so + // a streaming client (e.g. Codex) has usually disconnected by now, + // cancelling r.Context(). These middlewares only observe and record + // (token/cost metering, usage recording) and must still complete — + // otherwise the dispatcher short-circuits each to fail-mode and the + // usage is silently lost. Detach from client cancellation, keep ctx + // values, and bound the work. + obsCtx, obsCancel := context.WithTimeout(context.WithoutCancel(ctx), observabilityPhaseTimeout) + defer obsCancel() + + respMeta := chain.RunResponse(obsCtx, respInput, acc) if capturedData != nil { - for _, kv := range requestMeta { + for _, kv := range respMeta { capturedData.SetMetadata(kv.Key, kv.Value) } } - if denyOutput != nil { - middlewareID := "middleware" - if denyOutput.DenyReason != nil && denyOutput.DenyReason.Code != "" { - middlewareID = denyOutput.DenyReason.Code - } - // Policy/budget/routing/guardrail denials are expected runtime outcomes - // and can be high-volume under misconfigured or hostile clients; keep - // per-request detail at Debug and rely on metrics/access logs at scale. - p.logger.Debugf("middleware chain denied request: service=%s path=%s middlewares=%v reason=%s status=%d", - result.serviceID, result.matchedPath, middlewareIDs, middlewareID, denyOutput.DenyStatus) - middleware.RenderDenyResponse(w, middlewareID, denyOutput.DenyReason, denyOutput.DenyStatus) - return - } - respWriter := http.ResponseWriter(w) - var capturingWriter *bodytap.CapturingResponseWriter - if pt.CaptureConfig != nil && pt.CaptureConfig.MaxResponseBytes > 0 { - capturingWriter = bodytap.NewCapturingResponseWriter(w, pt.CaptureConfig.MaxResponseBytes, p.middlewareManager.Budget()) - defer capturingWriter.Release() - if capturingWriter.Bypassed() { - if capturedData != nil { - capturedData.SetMetadata("mw.capture.bypass_reason", capturingWriter.BypassReason()) - } - p.middlewareManager.Metrics().IncCaptureBypass(ctx, string(result.serviceID), capturingWriter.BypassReason()) - capturingWriter = nil - } else { - respWriter = capturingWriter + // Terminal slot sees the merged metadata bag from request and + // response phases. + mergedMeta := append(append([]middleware.KV(nil), requestMeta...), respMeta...) + termInput := *respInput + termInput.Slot = middleware.SlotTerminal + termInput.Metadata = mergedMeta + termMeta := chain.RunTerminal(obsCtx, &termInput, acc) + if capturedData != nil { + for _, kv := range termMeta { + capturedData.SetMetadata(kv.Key, kv.Value) } } - defer func() { - if capturingWriter == nil { - return - } - respInput := &middleware.Input{ - Slot: middleware.SlotOnResponse, - RequestID: reqInput.RequestID, - TargetID: reqInput.TargetID, - Method: reqInput.Method, - URL: reqInput.URL, - Headers: reqInput.Headers, - Status: capturingWriter.Status(), - RespHeaders: headerToKV(w.Header()), - RespBody: capturingWriter.Body(), - RespBodyTruncated: capturingWriter.Truncated(), - OriginalRespSize: capturingWriter.BytesWritten(), - ServiceID: reqInput.ServiceID, - AccountID: reqInput.AccountID, - UserID: reqInput.UserID, - // UserEmail / UserGroups / UserGroupNames must flow into the - // response leg too — llm_limit_record needs UserGroups to send - // group_ids on RecordLLMUsage so management's account-budget - // fan-out can match group-targeted rules; identity-stamping and - // any future response-side authorisation also depend on these. - UserEmail: reqInput.UserEmail, - UserGroups: reqInput.UserGroups, - UserGroupNames: reqInput.UserGroupNames, - AuthMethod: reqInput.AuthMethod, - SourceIP: reqInput.SourceIP, - Metadata: requestMeta, - AgentNetwork: reqInput.AgentNetwork, - } - // The response/terminal phase runs after the body is forwarded, so - // a streaming client (e.g. Codex) has usually disconnected by now, - // cancelling r.Context(). These middlewares only observe and record - // (token/cost metering, usage recording) and must still complete — - // otherwise the dispatcher short-circuits each to fail-mode and the - // usage is silently lost. Detach from client cancellation, keep ctx - // values, and bound the work. - obsCtx, obsCancel := context.WithTimeout(context.WithoutCancel(ctx), observabilityPhaseTimeout) - defer obsCancel() - - respMeta := chain.RunResponse(obsCtx, respInput, acc) - if capturedData != nil { - for _, kv := range respMeta { - capturedData.SetMetadata(kv.Key, kv.Value) - } - } - - // Terminal slot sees the merged metadata bag from request and - // response phases. - mergedMeta := append(append([]middleware.KV(nil), requestMeta...), respMeta...) - termInput := *respInput - termInput.Slot = middleware.SlotTerminal - termInput.Metadata = mergedMeta - termMeta := chain.RunTerminal(obsCtx, &termInput, acc) - if capturedData != nil { - for _, kv := range termMeta { - capturedData.SetMetadata(kv.Key, kv.Value) - } - } - - p.logger.Debugf("middleware chain ran: service=%s path=%s middlewares=%v status=%d req_meta=%d resp_meta=%d term_meta=%d", - result.serviceID, result.matchedPath, middlewareIDs, capturingWriter.Status(), len(requestMeta), len(respMeta), len(termMeta)) - }() + p.logger.Debugf("middleware chain ran: service=%s path=%s middlewares=%v status=%d req_meta=%d resp_meta=%d term_meta=%d", + result.serviceID, result.matchedPath, middlewareIDs, capturingWriter.Status(), len(requestMeta), len(respMeta), len(termMeta)) +} +// forwardUpstream applies any middleware-emitted upstream rewrite and proxies +// the request to the effective upstream URL. +func (p *ReverseProxy) forwardUpstream(respWriter http.ResponseWriter, r *http.Request, ctx context.Context, result targetResult, rewriteMatchedPath string, upstreamRewrite *middleware.UpstreamRewrite) { + pt := result.target effectiveURL := applyUpstreamRewrite(pt.URL, upstreamRewrite) if upstreamRewrite != nil { r.Host = effectiveURL.Host