[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.

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

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 }
 

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

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