refactor: migrate governance routing from HTTPTransportPreHook to PreRequestHook

[Pratham-Mishra04]

Summary

Governance routing logic has been migrated from HTTPTransportPreHook into PreRequestHook, operating directly on BifrostRequest structs rather than raw HTTP bodies. This eliminates the need to unmarshal/marshal JSON or multipart bodies in the transport hook, removes integration-specific path sniffing (Gemini /genai, Bedrock /bedrock), and makes routing decisions available earlier in the pipeline where provider/model fields are already normalized.

Changes

• HTTPTransportPreHook is now a no-op stub retained only to satisfy the HTTPTransportPlugin interface. All routing (virtual key load balancing, routing rules, MCP tool injection) flows through PreRequestHook.
• loadBalanceProvider and applyRoutingRules now accept *schemas.BifrostRequest instead of map[string]any + *schemas.HTTPRequest, mutating Provider, Model, and Fallbacks directly on the request struct. Fallbacks are now typed []schemas.Fallback rather than []string.
• governLargePayload and governRealtimeQueryParam are removed. Large-payload routing is handled in PreRequestHook by reading LargePayloadMetadata.Model from context and writing back the routed model. Realtime WebSocket upgrades now invoke RunPreRequestHooks explicitly in wsrealtime.go before the upgrade completes, with provider/model mutations read back into local vars and mirrored to fasthttp user values.
• addMCPIncludeTools (which wrote an HTTP header) is replaced by computeMCPIncludeTools (which returns a []string), stored via ctx.SetValue(schemas.MCPContextKeyIncludeTools, ...) instead of mutating request headers.
• stampGovernanceCtxFromVK is extracted as a standalone helper in utils.go to copy team/customer identifiers from a virtual key onto the context, including the team's customer relationship which was previously missed.
• validateRequiredHeaders is moved from main.go to utils.go.
• parseVirtualKeyFromHTTPRequest is removed; virtual key resolution now reads BifrostContextKeyVirtualKey set upstream by the transport middleware.
• BifrostContextKeyRequestQuery is added to the schema and populated in ConvertToBifrostContext so governance CEL routing rules can evaluate query parameters. Keys are lowercased for case-insensitive lookup.
• github.com/bytedance/sonic is demoted from a direct dependency to an indirect dependency in the governance module since it is no longer used for body unmarshaling in this plugin.
• Existing HTTPTransportPreHook tests are skipped with a note to rewrite them as PreRequestHook tests in Phase 3 of the refactor.

Type of change

• Bug fix
• Feature
• Refactor
• Documentation
• Chore/CI

Affected areas

• Core (Go)
• Transports (HTTP)
• Providers/Integrations
• Plugins
• UI (React)
• Docs

How to test

go test ./plugins/governance/...
go test ./transports/bifrost-http/...
go build ./...

Note: several HTTPTransportPreHook tests are currently skipped pending Phase 3 rewrites. Realtime routing can be validated end-to-end by connecting a WebSocket client to /v1/realtime?model=<model> with a virtual key that has weighted provider configs and confirming the selected provider is reflected in the upstream connection.

Screenshots/Recordings

N/A

Breaking changes

• Yes
• No

loadBalanceProvider and applyRoutingRules signatures have changed. Any internal callers outside the governance plugin that referenced these methods directly will need to be updated to pass *schemas.BifrostRequest instead of map[string]any. Fallbacks previously written as []string in the request body are now []schemas.Fallback structs on the request object.

Related issues

N/A

Security considerations

Virtual key resolution no longer parses raw Authorization/x-api-key headers inside the governance plugin; it relies on the value already extracted and stored in BifrostContextKeyVirtualKey by the transport layer. Ensure the transport middleware correctly populates this context key before governance runs.

Checklist

• I read docs/contributing/README.md and followed the guidelines
• I added/updated tests where appropriate
• I updated documentation where needed
• I verified builds succeed (Go and UI)
• I verified the CI pipeline passes locally if applicable

Summary by CodeRabbit

• New Features

  • Governance routing now runs earlier (pre-request) and applies to WebSocket realtime and large-payload requests, affecting provider/model selection and include-tool injection.
  • HTTP query parameters are captured (lowercased) for routing decisions.
  • Virtual-key extraction supports multiple header forms and stamps governance context.
  • Required-header validation returns a clear 400 error listing missing headers.

• Tests

  • Several governance transport tests are marked to skip pending migrated behavior.

• Chores

  • Minor dependency update.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[coderabbitai]

Warning

Review limit reached

@Pratham-Mishra04, we couldn't start this review because you've reached your PR review rate limit.

More reviews will be available in 27 minutes and 36 seconds. Learn how PR review limits work.

Your organization has run out of usage credits. Purchase more in the billing tab.

⌛ How to resolve this issue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available.

Please see our Fair Usage Limits Policy for further information.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 10e6e068-0185-49fb-8a9f-ace0e3b54ea3

📥 Commits

Reviewing files that changed from the base of the PR and between 4b9952a and 1f4b826.

📒 Files selected for processing (7)

• core/schemas/bifrost.go
• plugins/governance/go.mod
• plugins/governance/httptransportprehook_test.go
• plugins/governance/main.go
• plugins/governance/utils.go
• transports/bifrost-http/handlers/wsrealtime.go
• transports/bifrost-http/lib/ctx.go

📝 Walkthrough

Walkthrough

Migrates governance routing/load-balancing from HTTPTransportPreHook to PreRequestHook, adds a Bifrost context key for lowercased query params, updates VK parsing/stamping and header validation, applies pre-request hooks to WebSocket upgrades, updates tests, and adjusts an indirect dependency.

Changes

Governance Routing Migration to PreRequestHook

Layer / File(s)	Summary
Context schema and query parameter collection core/schemas/bifrost.go, transports/bifrost-http/lib/ctx.go	Introduces BifrostContextKeyRequestQuery and updates ConvertToBifrostContext to collect lowercased query parameters into the Bifrost context.
Governance utility helpers plugins/governance/utils.go	Exports ParseVirtualKeyFromFastHTTPRequest, adds stampGovernanceCtxFromVK to copy governance identifiers from TableVirtualKey, and implements validateRequiredHeaders returning a 400 missing_required_headers error when required headers are absent.
Governance routing refactor in main plugins/governance/main.go	Moves routing/load-balancing to operate on *schemas.BifrostRequest, adds runPreRequestRouting, refactors loadBalanceProvider and applyRoutingRules, implements PreRequestHook, and updates computeMCPIncludeTools.
WebSocket realtime pre-request governance transports/bifrost-http/handlers/wsrealtime.go	Runs PreRequestHooks on WebSocket upgrade requests: builds Bifrost context (headers+query), marks realtime metadata, executes hooks, applies routed provider/model, and mirrors user values back to fasthttp context.
Test migration for routing logic relocation plugins/governance/httptransportprehook_test.go	Adds t.Skip() to six HTTPTransportPreHook tests that exercised body-containing routing behavior now moved to PreRequestHook.
Dependency version adjustment plugins/governance/go.mod	Downgrades indirect dependency github.com/bytedance/sonic from v1.15.1 to v1.15.0.

Sequence Diagram(s)

sequenceDiagram
  participant Request
  participant PreRequestHook
  participant VirtualKeyStore
  participant RoutingEngine
  participant LoadBalancer
  participant ModelCatalog
  participant BifrostContext
  Request->>PreRequestHook: incoming request
  PreRequestHook->>VirtualKeyStore: load VK from context (if present)
  VirtualKeyStore-->>PreRequestHook: return VK
  PreRequestHook->>RoutingEngine: evaluate rules (headers + query from BifrostContext)
  RoutingEngine-->>PreRequestHook: decision (provider/model/fallbacks, KeyID)
  PreRequestHook->>LoadBalancer: perform weighted selection if needed
  LoadBalancer->>ModelCatalog: optional model refinement
  ModelCatalog-->>LoadBalancer: refined models
  LoadBalancer-->>PreRequestHook: selected provider/model and fallbacks
  PreRequestHook->>BifrostContext: stamp APIKeyID and MCP include-tools
  PreRequestHook-->>Request: return routed request (possibly mutated model/provider)

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

• maximhq/bifrost#3924: Related work touching provider/fallback selection and loadBalanceProvider adjustments.

Suggested reviewers

• akshaydeo
• danpiths

Poem

🐰 I hopped through headers, queries, too,
Stamped VK tales and routing through,
WebSocket doors now run the hook,
Providers balanced with one look,
A little rabbit cheers the queue!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely summarizes the main change: migrating governance routing from HTTPTransportPreHook to PreRequestHook. It directly reflects the core refactoring described in the PR objectives.
Description check	✅ Passed	The description comprehensively covers all required sections: summary, changes, type of change, affected areas, testing instructions, breaking changes, security considerations, and checklist. All major changes are documented with clear rationale.
Docstring Coverage	✅ Passed	Docstring coverage is 90.91% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 06-01-feat_governance_routing_moved_to_prerequesthook

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[Pratham-Mishra04]

• feat: enforce routing allowlist in core and propagate VK provider constraints via context key #3936 : 2 dependent PRs (#3711 , #4014 )
• docs: document PreRequestHook routing phase, per-request vs per-attempt semantics, and updated plugin sequencing #3935
• refactor: move model-catalog provider resolution into a modelcatalogresolver PreRequestHook plugin #3934
• refactor: migrate governance routing from HTTPTransportPreHook to PreRequestHook #3933 👈 (View in Graphite)
• feat: add PreRequestHook to LLMPlugin interface as a once-per-request routing phase before PreLLMHook #3932
• dev

This stack of pull requests is managed by Graphite. Learn more about stacking.

[greptile-apps]

Confidence Score: 4/5

Safe to merge with awareness that all governance routing tests are skipped and will remain so until Phase 3; any regression in weighted load balancing, routing-rule matching, or fallback generation will be invisible in CI until those tests are rewritten.

The refactor is architecturally sound and the new PreRequestHook path is logically equivalent to the removed HTTPTransportPreHook body-routing logic. However, the six tests that validated VK load balancing, routing-rule-preserves-target (with and without in-memory store), and Gemini/Bedrock variants are all skipped with no replacement, so the core routing contracts are unverified by automated tests.

plugins/governance/httptransportprehook_test.go — all routing tests are skipped. plugins/governance/main.go — the new PreRequestHook and runPreRequestRouting paths have no direct test coverage.

Important Files Changed

Filename	Overview
core/schemas/bifrost.go	Adds BifrostContextKeyRequestQuery constant for query-param map populated by ConvertToBifrostContext and the WS realtime upgrade handler; straightforward schema addition.
plugins/governance/go.mod	Demotes github.com/bytedance/sonic from direct to indirect dependency; consistent with removing its direct use from body unmarshaling in the governance plugin.
plugins/governance/httptransportprehook_test.go	Adds t.Skip to all six routing-coverage tests (VK load balancing, routing-rule-preserves-target, Gemini/Bedrock variants) with no replacement PreRequestHook tests, leaving the migrated routing logic entirely uncovered in CI.
plugins/governance/main.go	Major refactor: HTTPTransportPreHook becomes a no-op stub; PreRequestHook now handles all routing (load balancing, routing rules, MCP tool injection, large-payload), with a new runPreRequestRouting helper for the large-payload path. A stale comment incorrectly says realtime/streaming are handled by the no-op HTTPTransportPreHook.
plugins/governance/utils.go	Removes parseVirtualKeyFromHTTPRequest (VK parsing now upstream); adds stampGovernanceCtxFromVK helper and moves validateRequiredHeaders here. The new else-branch in stampGovernanceCtxFromVK changes customer attribution semantics for VKs with both a Team and a direct CustomerID (already flagged in a prior review comment).
transports/bifrost-http/handlers/wsrealtime.go	Adds an explicit RunPreRequestHooks call during WS upgrade to give governance a chance to route the realtime connection; correctly populates headers and query params on the pre-request context and mirrors mutations back to fasthttp user values for snapshotRealtimeMiddlewareValues.
transports/bifrost-http/lib/ctx.go	Adds query-param collection to ConvertToBifrostContext, lowercased and stored under BifrostContextKeyRequestQuery; mirrors the header-collection pattern already present.

_{Reviews (13): Last reviewed commit: "feat: governance routing moved to prereq..." | Re-trigger Greptile}

[coderabbitai]

♻️ Duplicate comments (1)

plugins/governance/go.mod (1)

58-58: ⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Duplicate sonic requires make the v1.15.0 change ineffective.

Line 8 still has a direct require for github.com/bytedance/sonic v1.15.1, and line 58 has an indirect require at v1.15.0. Go's module resolver will select v1.15.1 (the highest version), so this change to v1.15.0 has no effect on the actual dependency resolution.

If the intent is to demote sonic to an indirect dependency (per the PR summary), remove the direct require at line 8. If the intent is to downgrade to v1.15.0, update line 8 to match or remove the conflicting indirect entry.

Run the following script to confirm the effective selected version:

#!/bin/bash
cd plugins/governance

echo "=== All sonic require lines in go.mod ==="
rg -n '^\s*github\.com/bytedance/sonic\s+v' go.mod

echo ""
echo "=== Effective version selected by Go module resolver ==="
go list -m -f '{{.Version}}' github.com/bytedance/sonic

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@plugins/governance/go.mod` at line 58, go.mod currently has duplicate entries
for the module github.com/bytedance/sonic (direct v1.15.1 and indirect v1.15.0)
so the indirect change is ineffective; fix this by editing the go.mod require
entries for github.com/bytedance/sonic to reflect the intended state — either
remove the direct require (demote it to indirect) or change the direct require
version to v1.15.0 to match the indirect entry — then run module cleanup (go mod
tidy) to refresh go.sum and verify the effective selected version with the
module resolver.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Duplicate comments:
In `@plugins/governance/go.mod`:
- Line 58: go.mod currently has duplicate entries for the module
github.com/bytedance/sonic (direct v1.15.1 and indirect v1.15.0) so the indirect
change is ineffective; fix this by editing the go.mod require entries for
github.com/bytedance/sonic to reflect the intended state — either remove the
direct require (demote it to indirect) or change the direct require version to
v1.15.0 to match the indirect entry — then run module cleanup (go mod tidy) to
refresh go.sum and verify the effective selected version with the module
resolver.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: a914a6f4-8505-4213-aaff-d7a2306ce770

📥 Commits

Reviewing files that changed from the base of the PR and between 21ae88f and 4b9952a.

📒 Files selected for processing (7)

• core/schemas/bifrost.go
• plugins/governance/go.mod
• plugins/governance/httptransportprehook_test.go
• plugins/governance/main.go
• plugins/governance/utils.go
• transports/bifrost-http/handlers/wsrealtime.go
• transports/bifrost-http/lib/ctx.go

[akshaydeo]

Merge activity

• Jun 9, 5:17 AM UTC: A user started a stack merge that includes this pull request via Graphite.
• Jun 9, 5:18 AM UTC: @akshaydeo merged this pull request with Graphite.