fix: commit routing-pinned key ID to reserved BifrostContextKeyAPIKeyID after PreRequestHook unblock

[Pratham-Mishra04]

Summary

Routing rules that pin a specific API key by ID were silently broken. During PreRequestHook execution, core blocks writes to reserved context keys (including BifrostContextKeyAPIKeyID) to prevent plugins from overriding caller-supplied values. The governance plugin was writing the routing-rule key pin directly to that reserved key, so the write was dropped and the pin never took effect.

Changes

• Introduced a new non-reserved context key BifrostContextKeyRoutingPinnedAPIKeyID that the governance plugin writes its routing-rule key pin to during the blocked PreRequestHook phase.
• After all PreRequestHooks complete and the restricted-write block is lifted, core's RunPreRequestHooks commits the routing pin from BifrostContextKeyRoutingPinnedAPIKeyID into the reserved BifrostContextKeyAPIKeyID. A non-empty routing pin overrides any caller-supplied pin, since the routing rule represents authoritative server-side policy.
• The defer ctx.UnblockRestrictedWrites() was replaced with an explicit call after the plugin loop, so the commit step runs after the block is lifted rather than after the function returns.
• Updated the governance plugin's applyRoutingRules to write to the new non-reserved key instead of the reserved one.
• Updated tests in both core/bifrost_test.go and plugins/governance/routing_test.go to exercise the real propagation path, including the restricted-write block and plugin scope, and to assert the pin lands on the correct context key.

Type of change

• Bug fix

Affected areas

• Core (Go)
• Plugins

How to test

go test ./core/... ./plugins/governance/...

The new TestRunPreRequestHooks_CommitsRoutingPinnedKey test covers three cases:

1. A routing pin is committed to the reserved BifrostContextKeyAPIKeyID.
2. A routing pin overrides a caller-supplied BifrostContextKeyAPIKeyID.
3. A caller-supplied BifrostContextKeyAPIKeyID is preserved when no routing pin is set.

The updated TestEvaluateRoutingRules_MultiTargetDeterministicWithPinnedKey test exercises the full propagation path through applyRoutingRules under the same restricted-write block that production uses.

Breaking changes

• Yes
• No

Security considerations

The routing pin is written to a dedicated non-reserved key and committed to the reserved key exclusively by core after the plugin phase completes. This preserves the invariant that only core is the authoritative writer of BifrostContextKeyAPIKeyID, preventing plugins from directly overriding key selection outside of the sanctioned routing-rule mechanism.

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

Release Notes

• New Features

  • Routing plugins can now pin specific API keys that take precedence over caller-supplied selections through dedicated routing-pinned key handling in the pre-request hook phase.

• Tests

  • Added tests to verify routing-pinned API key commitment, precedence handling, and proper context propagation through key selection logic.

[mintlify]

Preview deployment for your docs. Learn more about Mintlify Previews.

Project	Status	Preview	Updated (UTC)
Bifrost	🟡 Building	–	Jun 13, 2026, 7:59 AM

💡 Tip: Enable Workflows to automatically generate PRs for you.

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

[Pratham-Mishra04]

• fix: commit routing-pinned key ID to reserved BifrostContextKeyAPIKeyID after PreRequestHook unblock #4359 👈 (View in Graphite)
• dev

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

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces a routing-specific context key for API key pinning that the governance routing plugin populates, then commits into the canonical key-selection field during pre-request hook processing. This ensures routing-pinned keys take precedence over caller-supplied pins while maintaining separation of concerns between the plugin-scoped routing context and the reserved key-selection context.

Changes

Routing-pinned API key context commitment in pre-request hooks

Layer / File(s)	Summary
Routing pinned key context constant core/schemas/bifrost.go	Introduces BifrostContextKeyRoutingPinnedAPIKeyID constant for routing-specific API key pinning. Also updates BifrostResponseExtraFields field formatting.
Governance plugin writes routing-specific key pin plugins/governance/main.go	When applyRoutingRules resolves a KeyID, it writes to the routing-specific context key BifrostContextKeyRoutingPinnedAPIKeyID instead of the reserved key-selection key, allowing downstream pre-request hook logic to identify and commit the routing pin.
Pre-request hook commits routing pin to canonical key core/bifrost.go	RunPreRequestHooks explicitly unblocks restricted context writes after plugin hook processing, then reads the routing-specific pinned key and commits its trimmed value into BifrostContextKeyAPIKeyID when non-empty, overriding any caller-supplied pin.
Test coverage for routing pin commitment core/bifrost_test.go, plugins/governance/routing_test.go	Added test helpers fakeRoutingPlugin and newRoutingCommitPipeline, plus TestRunPreRequestHooks_CommitsRoutingPinnedKey to verify pre-request hook commits routing pins with correct precedence (routing pin overrides caller pin; caller pin preserved when no routing pin). Updated governance routing test to verify pinned key propagation through the real routing application path.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• maximhq/bifrost#3932: Introduces the PreRequestHook stage in PluginPipeline.RunPreRequestHooks that this PR extends to commit routing-pinned API keys.
• maximhq/bifrost#4175: Also modifies the PreRequestHook flow in core/bifrost.go and relies on the hook infrastructure that this PR builds upon.
• maximhq/bifrost#3933: Related changes to governance routing's key-pinning logic in plugins/governance/main.go that this PR refactors to use a routing-specific context key.

Suggested reviewers

• akshaydeo
• danpiths

Poem

🐰 Routing pins and keys align,
No more conflicts—just one design,
Governance whispers what to choose,
Pre-request hooks commit and smoothe,
Caller hints now bow to signs! ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main change: committing routing-pinned key ID after unblocking, directly addressing the core fix in this PR.
Description check	✅ Passed	The description comprehensively follows the template, covering all major sections including summary, detailed changes, type of change, affected areas, testing instructions, breaking changes, security considerations, and checklist.
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 docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 06-13-fix_routing_rule_key_pinning_fixes

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (2.12.2)

level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies"

---

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

[greptile-apps]

Confidence Score: 5/5

Safe to merge. The change is narrowly scoped to the routing-pin propagation path, all existing reserved-key invariants are preserved, and the new commit step is correctly ordered after the write block is lifted.

The defer-to-explicit-call refactor and the two-phase write (staging key to reserved key) are both correct. The blockRestrictedWrites atomic flag is only false when the commit runs, so the write to BifrostContextKeyAPIKeyID succeeds. Both call sites of applyRoutingRules are inside PreRequestHook, so the commit in RunPreRequestHooks always fires after the governance write. Tests cover all meaningful cases.

No files require special attention.

Important Files Changed

Filename	Overview
core/bifrost.go	Replaces deferred UnblockRestrictedWrites with an explicit call followed by a commit step that promotes BifrostContextKeyRoutingPinnedAPIKeyID to BifrostContextKeyAPIKeyID; the ordering is correct and the fix resolves the silent-drop bug.
core/schemas/bifrost.go	Adds BifrostContextKeyRoutingPinnedAPIKeyID as a non-reserved context key with appropriate DO NOT SET THIS MANUALLY annotation; cosmetic struct-tag alignment changes are inert.
plugins/governance/main.go	Switches the routing-rule key pin write from the reserved BifrostContextKeyAPIKeyID to the new non-reserved BifrostContextKeyRoutingPinnedAPIKeyID; both call sites of applyRoutingRules are within PreRequestHook so the core commit step will always fire.
core/bifrost_test.go	Adds TestRunPreRequestHooks_CommitsRoutingPinnedKey covering all three meaningful cases using a minimal fake plugin that mirrors the real governance write path.
plugins/governance/routing_test.go	Updated TestEvaluateRoutingRules_MultiTargetDeterministicWithPinnedKey to exercise the real propagation path under BlockRestrictedWrites and WithPluginScope.

_{Reviews (1): Last reviewed commit: "fix: routing rule key pinning fixes" | Re-trigger Greptile}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

core/bifrost.go (1)

6854-6890: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Clear the routing pin staging key per request.

BifrostContextKeyRoutingPinnedAPIKeyID now lives on the shared mutable context, but this path never clears it. On a reused/long-lived BifrostContext (the realtime/WebSocket case called out above), a pin from request A will still be present when request B has no routing match, so Lines 6887-6889 re-pin request B to the old key and override the caller pin you intended to preserve.

💡 One way to make the commit request-scoped and keep unblock panic-safe

 func (p *PluginPipeline) RunPreRequestHooks(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) {
   // If the skip plugin pipeline flag is set, skip the plugin pipeline
   if skipPluginPipeline, ok := ctx.Value(schemas.BifrostContextKeySkipPluginPipeline).(bool); ok && skipPluginPipeline {
     return
   }
-  ctx.BlockRestrictedWrites()
-  for _, plugin := range p.llmPlugins {
-    ...
-  }
-  ctx.UnblockRestrictedWrites()
+  ctx.ClearValue(schemas.BifrostContextKeyRoutingPinnedAPIKeyID)
+  func() {
+    ctx.BlockRestrictedWrites()
+    defer ctx.UnblockRestrictedWrites()
+    for _, plugin := range p.llmPlugins {
+      ...
+    }
+  }()

   // Commit the routing-rule key pin...
   if pin, ok := ctx.Value(schemas.BifrostContextKeyRoutingPinnedAPIKeyID).(string); ok {
+    ctx.ClearValue(schemas.BifrostContextKeyRoutingPinnedAPIKeyID)
     if pin = strings.TrimSpace(pin); pin != "" {
       ctx.SetValue(schemas.BifrostContextKeyAPIKeyID, pin)
     }
   }
 }

Based on learnings, (*schemas.BifrostContext).SetValue mutates the shared context in place, so values persist on reused contexts unless they are explicitly cleared.

🤖 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 `@core/bifrost.go` around lines 6854 - 6890, The staging routing pin
BifrostContextKeyRoutingPinnedAPIKeyID is left on a shared BifrostContext and
can leak between requests; fix by explicitly clearing/resetting that key
per-request (e.g., call
ctx.SetValue(schemas.BifrostContextKeyRoutingPinnedAPIKeyID, "") or remove it)
at the start of request handling (before running plugin PreRequestHook and
before any blocked/unblocked phases) and ensure the commit logic that reads the
pin (the block that calls ctx.SetValue(schemas.BifrostContextKeyAPIKeyID, pin))
treats an empty string as “no pin” so no stale value is re-applied; reference
symbols: BifrostContextKeyRoutingPinnedAPIKeyID, ctx.SetValue, PreRequestHook
loop and the pin-commit block that reads the routing pin.

Source: Learnings

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

Outside diff comments:
In `@core/bifrost.go`:
- Around line 6854-6890: The staging routing pin
BifrostContextKeyRoutingPinnedAPIKeyID is left on a shared BifrostContext and
can leak between requests; fix by explicitly clearing/resetting that key
per-request (e.g., call
ctx.SetValue(schemas.BifrostContextKeyRoutingPinnedAPIKeyID, "") or remove it)
at the start of request handling (before running plugin PreRequestHook and
before any blocked/unblocked phases) and ensure the commit logic that reads the
pin (the block that calls ctx.SetValue(schemas.BifrostContextKeyAPIKeyID, pin))
treats an empty string as “no pin” so no stale value is re-applied; reference
symbols: BifrostContextKeyRoutingPinnedAPIKeyID, ctx.SetValue, PreRequestHook
loop and the pin-commit block that reads the routing pin.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 85e216d1-72d0-4c79-adb1-5e608a858f73

📥 Commits

Reviewing files that changed from the base of the PR and between 0bbef33 and ac70de6.

📒 Files selected for processing (5)

• core/bifrost.go
• core/bifrost_test.go
• core/schemas/bifrost.go
• plugins/governance/main.go
• plugins/governance/routing_test.go

[akshaydeo]

Merge activity

• Jun 13, 2:10 PM UTC: A user started a stack merge that includes this pull request via Graphite.
• Jun 13, 2:10 PM UTC: @akshaydeo merged this pull request with Graphite.