From b6e176432197a224ff0eb4ce6e555b8990dab6bf Mon Sep 17 00:00:00 2001
From: whit3rabbit <whit3rabbit@users.noreply.github.com>
Date: Sat, 12 Sep 2026 08:58:24 -0500
Subject: [PATCH] fix: enforce virtual key policy after auto routing

---
 crates/proxy/src/server/chat_completions/handler.rs |  9 +++++++++
 crates/proxy/src/server/routes/messages.rs          | 10 ++++++++++
 2 files changed, 19 insertions(+)

diff --git a/crates/proxy/src/server/chat_completions/handler.rs b/crates/proxy/src/server/chat_completions/handler.rs
index 48f1065..b881876 100644
--- a/crates/proxy/src/server/chat_completions/handler.rs
+++ b/crates/proxy/src/server/chat_completions/handler.rs
@@ -93,6 +93,15 @@ pub(crate) async fn chat_completions(
         let signals = crate::server::router_signals::openai_signals(&body);
         state.resolve_router_tier(&rc, &signals)
     });
+    if let (Some(ref ctx), Some((model, _, _))) = (&vk_ctx, &router_tier) {
+        if !crate::server::policy::is_model_allowed(model, &ctx.allowed_models) {
+            return openai_error_response(
+                &format!("Model '{model}' is not allowed for this API key."),
+                "permission_error",
+                StatusCode::FORBIDDEN,
+            );
+        }
+    }
     let (mapped_model, effective, deployment) = match router_tier {
         Some((model, effective, deployment)) => (
             mapped_model_for_backend(&original_model, model, &effective),
diff --git a/crates/proxy/src/server/routes/messages.rs b/crates/proxy/src/server/routes/messages.rs
index dc79166..dd87534 100644
--- a/crates/proxy/src/server/routes/messages.rs
+++ b/crates/proxy/src/server/routes/messages.rs
@@ -96,6 +96,16 @@ pub(crate) async fn messages(
         }
         None => (body, None),
     };
+    if let (Some(ref ctx), Some((model, _, _))) = (&vk_ctx, &router_tier) {
+        if !super::super::policy::is_model_allowed(model, &ctx.allowed_models) {
+            let err = mapping::errors_map::create_anthropic_error(
+                anthropic::ErrorType::PermissionError,
+                format!("Model '{model}' is not allowed for this API key."),
+                None,
+            );
+            return (StatusCode::FORBIDDEN, Json(err)).into_response();
+        }
+    }
     let (mapped_model, effective, deployment) = match router_tier {
         Some(v) => v,
         None => match state.resolve_model_and_state(&body.model) {