refactor(rust): add gateway and router hook seams

litellm-rust/Cargo.lock

@@ -3781,6 +3781,7 @@ dependencies = [
  "litellm-gateway-inference",
  "litellm-gateway-mcp",
  "litellm-gateway-ui",
+ "litellm-host",
  "litellm-http",
  "litellm-inference",
  "litellm-llms",
@@ -3797,6 +3798,7 @@ dependencies = [
  "tower-sessions-moka-store",
  "tracing",
  "uuid",
+ "wiremock",
 ]
 
 [[package]]
@@ -4312,6 +4314,7 @@ dependencies = [
  "litellm-inference",
  "litellm-inference-messages",
  "rstest",
+ "tokio",
 ]
 
 [[package]]

litellm-rust/crates/gateway-inference/src/audio_transcription.rs

@@ -1,4 +1,7 @@
 use litellm_gateway_auth::AuthenticatedRequest;
+use litellm_host::interceptors::Interceptors;
+use litellm_inference::RouteError;
+use litellm_router::RouterHooks;
 use std::{path::Path, sync::Arc};
 
 use axum::{Json, extract::State, response::IntoResponse};
@@ -11,23 +14,26 @@ use crate::{
     request::{self, InferenceBody},
 };
 
-pub(crate) async fn create(
-    State(gateway): State<Arc<Gateway>>,
+pub(crate) async fn create<
+    R: RouterHooks + 'static,
+    I: Interceptors<RouteError> + Clone + 'static,
+>(
+    State(gateway): State<Arc<Gateway<R, I>>>,
     identity: AuthenticatedRequest,
     body: InferenceBody,
 ) -> Result<impl IntoResponse, Error> {
     handle(&gateway, &identity, body).await.map(Json)
 }
 
-async fn handle(
-    gateway: &Gateway,
+async fn handle<R: RouterHooks + 'static, I: Interceptors<RouteError> + Clone + 'static>(
+    gateway: &Gateway<R, I>,
     identity: &AuthenticatedRequest,
     InferenceBody {
         fields: body,
         upload,
     }: InferenceBody,
 ) -> Result<Value, Error> {
-    let deployment = request::resolve_deployment(gateway, &body)?;
+    let deployment = request::resolve_deployment(gateway, &body).await?;
     request::authorize_model(identity, deployment, &body).await?;
     let audio = match upload {
         Some(upload) => {

litellm-rust/crates/gateway-inference/src/caching.rs

@@ -69,39 +69,52 @@ fn duration(seconds: f64) -> Result<Duration, Error> {
         .ok_or_else(|| Error::InvalidBody("cache durations must be finite and positive".into()))
 }
 
-#[derive(Clone, Default)]
-pub(crate) struct CacheHeaders(std::sync::Arc<std::sync::OnceLock<String>>);
+#[derive(Clone)]
+pub(crate) struct CacheHeaders<H> {
+    key: std::sync::Arc<std::sync::OnceLock<String>>,
+    hooks: H,
+}
 
-impl litellm_host::interceptors::Interceptors<litellm_inference::RouteError> for CacheHeaders {
+impl<H: litellm_host::interceptors::Interceptors<litellm_inference::RouteError>>
+    litellm_host::interceptors::Interceptors<litellm_inference::RouteError> for CacheHeaders<H>
+{
     async fn before_provider_request(
         &self,
         wire: litellm_host::interceptors::WireRequest,
-        _: litellm_host::interceptors::RequestContext,
+        context: litellm_host::interceptors::RequestContext,
     ) -> Result<litellm_host::interceptors::WireRequest, litellm_inference::RouteError> {
-        Ok(wire)
+        self.hooks.before_provider_request(wire, context).await
     }
 
     async fn after_provider_response(
         &self,
-        _: litellm_host::interceptors::RawResponse,
+        raw: litellm_host::interceptors::RawResponse,
     ) -> Result<(), litellm_inference::RouteError> {
-        Ok(())
+        self.hooks.after_provider_response(raw).await
     }
 
     async fn result_ready(
         &self,
         facts: litellm_host::interceptors::ExecutionFacts,
     ) -> Result<(), litellm_inference::RouteError> {
+        self.hooks.result_ready(facts.clone()).await?;
         if let litellm_host::interceptors::ResultSource::Cache { key } = facts.source {
-            let _ = self.0.set(key);
+            let _ = self.key.set(key);
         }
         Ok(())
     }
 }
 
-impl CacheHeaders {
+impl<H> CacheHeaders<H> {
+    pub(crate) fn new(hooks: H) -> Self {
+        Self {
+            key: Default::default(),
+            hooks,
+        }
+    }
+
     pub(crate) fn apply(&self, mut response: axum::response::Response) -> axum::response::Response {
-        if let Some(key) = self.0.get()
+        if let Some(key) = self.key.get()
             && let Ok(value) = axum::http::HeaderValue::from_str(key)
         {
             response.headers_mut().insert("x-litellm-cache-key", value);

litellm-rust/crates/gateway-inference/src/chat_completions.rs

@@ -1,4 +1,7 @@
 use litellm_gateway_auth::AuthenticatedRequest;
+use litellm_host::interceptors::Interceptors;
+use litellm_inference::RouteError;
+use litellm_router::RouterHooks;
 use std::sync::Arc;
 
 use axum::{
@@ -11,16 +14,22 @@ use serde_json::{Map, Value};
 
 use crate::{Error, Gateway, JsonObject, request};
 
-pub(crate) async fn create(
-    State(gateway): State<Arc<Gateway>>,
+pub(crate) async fn create<
+    R: RouterHooks + 'static,
+    I: Interceptors<RouteError> + Clone + 'static,
+>(
+    State(gateway): State<Arc<Gateway<R, I>>>,
     identity: AuthenticatedRequest,
     JsonObject(body): JsonObject,
 ) -> Result<impl IntoResponse, Error> {
     handle(&gateway, &identity, body).await
 }
 
-pub(crate) async fn create_from_model_path(
-    State(gateway): State<Arc<Gateway>>,
+pub(crate) async fn create_from_model_path<
+    R: RouterHooks + 'static,
+    I: Interceptors<RouteError> + Clone + 'static,
+>(
+    State(gateway): State<Arc<Gateway<R, I>>>,
     identity: AuthenticatedRequest,
     Path(model): Path<String>,
     JsonObject(body): JsonObject,
@@ -35,12 +44,12 @@ pub(crate) async fn create_from_model_path(
     handle(&gateway, &identity, body).await
 }
 
-async fn handle(
-    gateway: &Gateway,
+async fn handle<R: RouterHooks + 'static, I: Interceptors<RouteError> + Clone + 'static>(
+    gateway: &Gateway<R, I>,
     identity: &AuthenticatedRequest,
     body: Map<String, Value>,
 ) -> Result<Response, Error> {
-    let deployment = request::resolve_deployment(gateway, &body)?;
+    let deployment = request::resolve_deployment(gateway, &body).await?;
     request::authorize_model(identity, deployment, &body).await?;
     let (body, cache_options) = crate::caching::prepare(identity, body)?;
     let route = gateway.chat_completions.clone();
@@ -53,7 +62,7 @@ async fn handle(
     };
 
     let messages = body.get("messages").cloned().unwrap_or_default();
-    let headers = crate::caching::CacheHeaders::default();
+    let headers = crate::caching::CacheHeaders::new(gateway.interceptors.clone());
     let response = litellm_host_http::serve_unary(
         route.machine(
             ChatCompletionsCall {

litellm-rust/crates/gateway-inference/src/lib.rs

@@ -26,30 +26,26 @@ use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings};
 use litellm_secrets::source::SecretSource;
 
 pub use error::Error;
-pub use litellm_router::{Deployment, Router as ModelRouter};
+use litellm_host::interceptors::Interceptors;
+use litellm_inference::RouteError;
+pub use litellm_router::{Deployment, Router as ModelRouter, RouterHooks};
 pub use request::{JsonObject, RequestId};
 
-pub struct Gateway {
+pub struct Gateway<R = (), I = ()> {
     cache: Option<Arc<dyn litellm_cache_response::ResponseCacheService>>,
     pub audio_transcription: AudioTranscriptionRoute,
     pub chat_completions: ChatCompletionsRoute,
     pub messages: MessagesRoute,
     pub ocr: OcrRoute,
     pub responses: ResponsesRoute,
-    pub models: ModelRouter,
+    pub models: ModelRouter<R>,
+    pub interceptors: I,
     pub secrets: Arc<dyn SecretSource>,
     pub resources: CoreResources,
     pub http: HttpClientConfig,
 }
 
 impl Gateway {
-    pub fn with_cache(self, cache: Arc<dyn litellm_cache_response::ResponseCacheService>) -> Self {
-        Self {
-            cache: Some(cache),
-            ..self
-        }
-    }
-
     pub fn new(
         resources: CoreResources,
         http: HttpClientConfig,
@@ -60,6 +56,7 @@ impl Gateway {
         let auth = resources.auth.clone();
         Ok(Self {
             cache: None,
+            interceptors: (),
             audio_transcription: AudioTranscriptionRoute::new(
                 provider.clone(),
                 auth.clone(),
@@ -88,22 +85,57 @@ impl Gateway {
     }
 }
 
-pub fn router(gateway: Arc<Gateway>) -> Router {
+impl<R, I> Gateway<R, I> {
+    pub fn with_cache(self, cache: Arc<dyn litellm_cache_response::ResponseCacheService>) -> Self {
+        Self {
+            cache: Some(cache),
+            ..self
+        }
+    }
+
+    pub fn with_hooks<T, H>(self, router_hooks: T, interceptors: H) -> Gateway<T, H> {
+        Gateway {
+            cache: self.cache,
+            audio_transcription: self.audio_transcription,
+            chat_completions: self.chat_completions,
+            messages: self.messages,
+            ocr: self.ocr,
+            responses: self.responses,
+            models: self.models.with_hooks(router_hooks),
+            interceptors,
+            secrets: self.secrets,
+            resources: self.resources,
+            http: self.http,
+        }
+    }
+}
+
+pub fn router<R, I>(gateway: Arc<Gateway<R, I>>) -> Router
+where
+    R: RouterHooks + 'static,
+    I: Interceptors<RouteError> + Clone + 'static,
+{
     Router::new()
-        .route("/v1/messages", post(messages::create))
-        .route("/ocr", post(ocr::create))
-        .route("/v1/ocr", post(ocr::create))
-        .route("/chat/completions", post(chat_completions::create))
-        .route("/v1/chat/completions", post(chat_completions::create))
-        .nest("/engines/{model}", model_routes())
-        .nest("/openai/deployments/{model}", model_routes())
-        .route("/audio/transcriptions", post(audio_transcription::create))
+        .route("/v1/messages", post(messages::create::<R, I>))
+        .route("/ocr", post(ocr::create::<R, I>))
+        .route("/v1/ocr", post(ocr::create::<R, I>))
+        .route("/chat/completions", post(chat_completions::create::<R, I>))
+        .route(
+            "/v1/chat/completions",
+            post(chat_completions::create::<R, I>),
+        )
+        .nest("/engines/{model}", model_routes::<R, I>())
+        .nest("/openai/deployments/{model}", model_routes::<R, I>())
+        .route(
+            "/audio/transcriptions",
+            post(audio_transcription::create::<R, I>),
+        )
         .route(
             "/v1/audio/transcriptions",
-            post(audio_transcription::create),
+            post(audio_transcription::create::<R, I>),
         )
-        .route("/responses", post(responses::create))
-        .route("/v1/responses", post(responses::create))
+        .route("/responses", post(responses::create::<R, I>))
+        .route("/v1/responses", post(responses::create::<R, I>))
         .route("/embeddings", post(request::unsupported))
         .route("/v1/embeddings", post(request::unsupported))
         .route("/completions", post(request::unsupported))
@@ -114,11 +146,15 @@ pub fn router(gateway: Arc<Gateway>) -> Router {
         .with_state(gateway)
 }
 
-fn model_routes() -> Router<Arc<Gateway>> {
+fn model_routes<R, I>() -> Router<Arc<Gateway<R, I>>>
+where
+    R: RouterHooks + 'static,
+    I: Interceptors<RouteError> + Clone + 'static,
+{
     Router::new()
         .route(
             "/chat/completions",
-            post(chat_completions::create_from_model_path),
+            post(chat_completions::create_from_model_path::<R, I>),
         )
         .route("/embeddings", post(request::unsupported))
         .route("/completions", post(request::unsupported))

litellm-rust/crates/gateway-inference/src/messages.rs

@@ -1,6 +1,9 @@
 //! `POST /v1/messages`, as the Python proxy's `anthropic_response` serves it.
 
 use litellm_gateway_auth::AuthenticatedRequest;
+use litellm_host::interceptors::Interceptors;
+use litellm_inference::RouteError;
+use litellm_router::RouterHooks;
 use std::sync::Arc;
 
 use axum::{
@@ -21,8 +24,8 @@ use crate::{Deployment, Error, Gateway, JsonObject, RequestId, request};
 const ANTHROPIC_API_HEADERS: [&str; 2] = ["anthropic-version", "anthropic-beta"];
 const ANTHROPIC_API_HEADER_PROVIDERS: &str = "anthropic,bedrock,bedrock_mantle,vertex_ai";
 
-pub async fn create(
-    State(gateway): State<Arc<Gateway>>,
+pub async fn create<R: RouterHooks + 'static, I: Interceptors<RouteError> + Clone + 'static>(
+    State(gateway): State<Arc<Gateway<R, I>>>,
     identity: AuthenticatedRequest,
     RequestId(request_id): RequestId,
     headers: HeaderMap,
@@ -35,13 +38,13 @@ pub async fn create(
     result.map_err(|error| (error.status(), Json(error.body(request_id.as_deref()))))
 }
 
-async fn handle(
-    gateway: &Gateway,
+async fn handle<R: RouterHooks + 'static, I: Interceptors<RouteError> + Clone + 'static>(
+    gateway: &Gateway<R, I>,
     identity: &AuthenticatedRequest,
     headers: &HeaderMap,
     body: Map<String, Value>,
 ) -> Result<Response, Error> {
-    let deployment = request::resolve_deployment(gateway, &body)?;
+    let deployment = request::resolve_deployment(gateway, &body).await?;
     request::authorize_model(identity, deployment, &body).await?;
     let (body, cache_options) = crate::caching::prepare(identity, body)?;
     let route = gateway.messages.clone();
@@ -57,7 +60,7 @@ async fn handle(
     let machine = route.machine(call, cache_options.policy);
     let stream =
         Sse::<Messages, _, _>::new(Json, |error| Bytes::from(Error::from(error).sse_frame()));
-    let headers = crate::caching::CacheHeaders::default();
+    let headers = crate::caching::CacheHeaders::new(gateway.interceptors.clone());
     let response = litellm_host_http::serve(machine, (), headers.clone(), stream, None).await?;
     Ok(headers.apply(response))
 }

litellm-rust/crates/gateway-inference/src/ocr.rs

@@ -1,4 +1,7 @@
 use litellm_gateway_auth::AuthenticatedRequest;
+use litellm_host::interceptors::Interceptors;
+use litellm_inference::RouteError;
+use litellm_router::RouterHooks;
 use std::sync::Arc;
 
 use axum::{Json, extract::State, http::HeaderMap, response::IntoResponse};
@@ -13,8 +16,11 @@ use crate::{
     request::{self, InferenceBody},
 };
 
-pub(crate) async fn create(
-    State(gateway): State<Arc<Gateway>>,
+pub(crate) async fn create<
+    R: RouterHooks + 'static,
+    I: Interceptors<RouteError> + Clone + 'static,
+>(
+    State(gateway): State<Arc<Gateway<R, I>>>,
     identity: AuthenticatedRequest,
     headers: HeaderMap,
     body: InferenceBody,
@@ -22,8 +28,8 @@ pub(crate) async fn create(
     handle(&gateway, &identity, &headers, body).await.map(Json)
 }
 
-async fn handle(
-    gateway: &Gateway,
+async fn handle<R: RouterHooks + 'static, I: Interceptors<RouteError> + Clone + 'static>(
+    gateway: &Gateway<R, I>,
     identity: &AuthenticatedRequest,
     headers: &HeaderMap,
     InferenceBody {
@@ -35,7 +41,7 @@ async fn handle(
         .get("x-req-format")
         .and_then(|value| value.to_str().ok())
         .map(str::to_owned);
-    let deployment = request::resolve_deployment(gateway, &body)?;
+    let deployment = request::resolve_deployment(gateway, &body).await?;
     request::authorize_model(identity, deployment, &body).await?;
     let document = match upload {
         Some(upload) => OcrDocumentInput::Bytes {

litellm-rust/crates/gateway-inference/src/request.rs

@@ -110,8 +110,8 @@ pub(crate) async fn authorize_model(
     Ok(())
 }
 
-pub(crate) fn resolve_deployment<'a>(
-    gateway: &'a Gateway,
+pub(crate) async fn resolve_deployment<'a, R: crate::RouterHooks, I>(
+    gateway: &'a Gateway<R, I>,
     body: &Map<String, Value>,
 ) -> Result<&'a Deployment, Error> {
     let model = body
@@ -120,7 +120,8 @@ pub(crate) fn resolve_deployment<'a>(
         .ok_or_else(|| Error::InvalidBody("model is required".into()))?;
     gateway
         .models
-        .get(model)
+        .select(model)
+        .await
         .ok_or_else(|| Error::UnknownModel(model.to_owned()))
 }
 

litellm-rust/crates/gateway-inference/src/responses.rs

@@ -1,3 +1,6 @@
+use litellm_host::interceptors::Interceptors;
+use litellm_inference::RouteError;
+use litellm_router::RouterHooks;
 use std::sync::Arc;
 
 use axum::{Json, body::Bytes, extract::State, response::Response};
@@ -8,12 +11,15 @@ use serde_json::json;
 
 use crate::{Error, Gateway, JsonObject, request};
 
-pub(crate) async fn create(
-    State(gateway): State<Arc<Gateway>>,
+pub(crate) async fn create<
+    R: RouterHooks + 'static,
+    I: Interceptors<RouteError> + Clone + 'static,
+>(
+    State(gateway): State<Arc<Gateway<R, I>>>,
     identity: AuthenticatedRequest,
     JsonObject(body): JsonObject,
 ) -> Result<Response, Error> {
-    let deployment = request::resolve_deployment(&gateway, &body)?;
+    let deployment = request::resolve_deployment(&gateway, &body).await?;
     request::authorize_model(&identity, deployment, &body).await?;
     let (body, cache_options) = crate::caching::prepare(&identity, body)?;
     let route = gateway.responses.clone();
@@ -46,7 +52,7 @@ pub(crate) async fn create(
             json!({"type": "error", "code": error.status().as_u16().to_string(), "message": error.to_string(), "param": null})
         ))
     });
-    let headers = crate::caching::CacheHeaders::default();
+    let headers = crate::caching::CacheHeaders::new(gateway.interceptors.clone());
     let response = litellm_host_http::serve(machine, (), headers.clone(), stream, None).await?;
     Ok(headers.apply(response))
 }

litellm-rust/crates/gateway/Cargo.toml

@@ -10,6 +10,7 @@ axum.workspace = true
 envy = "0.4.2"
 http-body-util = "0.1"
 litellm-inference.workspace = true
+litellm-host.workspace = true
 litellm-gateway-inference.workspace = true
 litellm-gateway-mcp.workspace = true
 tokio-util = "0.7"
@@ -31,6 +32,7 @@ uuid.workspace = true
 tower-sessions-moka-store = "0.15.0"
 
 [dev-dependencies]
+wiremock.workspace = true
 base64.workspace = true
 rstest.workspace = true
 tempfile.workspace = true

litellm-rust/crates/gateway/src/hooks.rs

@@ -0,0 +1,23 @@
+use std::future::Future;
+
+use axum::{extract::Request, middleware::Next, response::Response};
+
+pub trait GatewayHooks: Send + Sync {
+    fn pre_call(&self, request: Request) -> impl Future<Output = Result<Request, Response>> + Send {
+        async { Ok(request) }
+    }
+
+    fn post_call(&self, response: Response) -> impl Future<Output = Response> + Send {
+        async { response }
+    }
+}
+
+impl GatewayHooks for () {}
+
+pub(crate) async fn run<H: GatewayHooks>(hooks: &H, request: Request, next: Next) -> Response {
+    let request = match hooks.pre_call(request).await {
+        Ok(request) => request,
+        Err(response) => return response,
+    };
+    hooks.post_call(next.run(request).await).await
+}

litellm-rust/crates/gateway/src/lib.rs

@@ -1,5 +1,6 @@
 mod auth;
 mod error;
+pub mod hooks;
 mod secrets;
 pub use error::Error;
 
@@ -113,14 +114,40 @@ pub async fn build_mcp(
     ))
 }
 
-pub fn router(
-    inference: Arc<Gateway>,
+pub fn router<R, I>(
+    inference: Arc<Gateway<R, I>>,
     config: &Config,
     ui: Option<Router>,
     mcp: Option<Router>,
-) -> Router {
+) -> Router
+where
+    R: litellm_gateway_inference::RouterHooks + 'static,
+    I: litellm_host::interceptors::Interceptors<litellm_inference::RouteError> + Clone + 'static,
+{
+    router_with_hooks(inference, config, ui, mcp, ())
+}
+
+pub fn router_with_hooks<R, I, H>(
+    inference: Arc<Gateway<R, I>>,
+    config: &Config,
+    ui: Option<Router>,
+    mcp: Option<Router>,
+    hooks: H,
+) -> Router
+where
+    R: litellm_gateway_inference::RouterHooks + 'static,
+    I: litellm_host::interceptors::Interceptors<litellm_inference::RouteError> + Clone + 'static,
+    H: hooks::GatewayHooks + 'static,
+{
+    let hooks = Arc::new(hooks);
     let auth = Auth::from_config(config, inference.secrets.clone());
     let inference = litellm_gateway_inference::router(inference)
+        .route_layer(axum::middleware::from_fn(
+            move |request: Request, next: Next| {
+                let hooks = hooks.clone();
+                async move { hooks::run(hooks.as_ref(), request, next).await }
+            },
+        ))
         .merge(mcp.unwrap_or_default())
         .route_layer(axum::middleware::from_fn(auth::bind_session_owner))
         .route_layer(axum::middleware::from_fn_with_state(

litellm-rust/crates/gateway/tests/hooks.rs

@@ -0,0 +1,183 @@
+use std::sync::{Arc, mpsc};
+
+use axum::{
+    body::{Body, to_bytes},
+    extract::Request,
+    http::StatusCode,
+    response::{IntoResponse, Response},
+};
+use litellm_config::Config;
+use litellm_gateway::hooks::GatewayHooks;
+use litellm_gateway_inference::{Deployment, RouterHooks};
+use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest};
+use litellm_inference::RouteError;
+use rstest::{fixture, rstest};
+use serde_json::{Value, json};
+use tower::ServiceExt;
+use wiremock::{
+    Mock, MockServer, ResponseTemplate,
+    matchers::{header, method, path},
+};
+
+#[derive(Clone)]
+struct Recorder {
+    phases: mpsc::Sender<&'static str>,
+    reject_gateway: bool,
+    reject_router: bool,
+}
+
+impl GatewayHooks for Recorder {
+    async fn pre_call(&self, request: Request) -> Result<Request, Response> {
+        self.phases.send("gateway.pre_call").unwrap();
+        if self.reject_gateway {
+            return Err((StatusCode::FORBIDDEN, "rejected by gateway hook").into_response());
+        }
+        Ok(request)
+    }
+
+    async fn post_call(&self, response: Response) -> Response {
+        self.phases.send("gateway.post_call").unwrap();
+        ([("x-hook-completed", "true")], response).into_response()
+    }
+}
+
+impl RouterHooks for Recorder {
+    async fn filter_deployments<'a>(
+        &self,
+        model: &str,
+        deployments: Vec<&'a Deployment>,
+    ) -> Vec<&'a Deployment> {
+        self.phases.send("router.filter_deployments").unwrap();
+        assert_eq!(model, "public-model");
+        assert_eq!(deployments.len(), 1);
+        if self.reject_router {
+            return Vec::new();
+        }
+        deployments
+    }
+}
+
+impl Interceptors<RouteError> for Recorder {
+    async fn before_provider_request(
+        &self,
+        wire: WireRequest,
+        context: RequestContext,
+    ) -> Result<WireRequest, RouteError> {
+        self.phases
+            .send("inference.before_provider_request")
+            .unwrap();
+        assert_eq!(context.model, wire.body["model"]);
+        Ok(WireRequest {
+            headers: wire
+                .headers
+                .into_iter()
+                .chain([("x-inference-hook".into(), "called".into())])
+                .collect(),
+            ..wire
+        })
+    }
+
+    async fn after_provider_response(&self, raw: RawResponse) -> Result<(), RouteError> {
+        self.phases
+            .send("inference.after_provider_response")
+            .unwrap();
+        let response: Value = serde_json::from_str(&raw.body).unwrap();
+        assert_eq!(response["content"][0]["text"], "provider reply");
+        Ok(())
+    }
+}
+
+#[fixture]
+fn phases() -> (mpsc::Sender<&'static str>, mpsc::Receiver<&'static str>) {
+    mpsc::channel()
+}
+
+#[rstest]
+#[case::all_layers(false, false, 200, &[
+    "gateway.pre_call",
+    "router.filter_deployments",
+    "inference.before_provider_request",
+    "inference.after_provider_response",
+    "gateway.post_call",
+])]
+#[case::gateway_rejection(true, false, 403, &["gateway.pre_call"])]
+#[case::router_filters_every_candidate(false, true, 400, &[
+    "gateway.pre_call",
+    "router.filter_deployments",
+    "gateway.post_call",
+])]
+#[tokio::test]
+async fn request_hooks_compose_and_rejections_stop_inner_layers(
+    phases: (mpsc::Sender<&'static str>, mpsc::Receiver<&'static str>),
+    #[case] reject_gateway: bool,
+    #[case] reject_router: bool,
+    #[case] status: u16,
+    #[case] expected_phases: &[&str],
+) {
+    let upstream = MockServer::start().await;
+    let message = json!({
+        "id": "msg_test", "type": "message", "role": "assistant", "model": "test-model",
+        "content": [{"type": "text", "text": "provider reply"}],
+        "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1},
+    });
+    Mock::given(method("POST"))
+        .and(path("/v1/messages"))
+        .and(header("x-inference-hook", "called"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(&message))
+        .expect(if reject_gateway || reject_router {
+            0
+        } else {
+            1
+        })
+        .mount(&upstream)
+        .await;
+    let config = Config::from_yaml(&format!(
+        "model_list:\n  - model_name: public-model\n    litellm_params:\n      model: anthropic/test-model\n      api_key: test-key\n      api_base: {}\ngeneral_settings:\n  master_key: gateway-key\n",
+        upstream.uri(),
+    )).unwrap();
+    let (sender, receiver) = phases;
+    let recorder = Recorder {
+        phases: sender,
+        reject_gateway,
+        reject_router,
+    };
+    let inference = Arc::into_inner(litellm_gateway::build_inference(&config).unwrap())
+        .unwrap()
+        .with_hooks(recorder.clone(), recorder.clone());
+    let app =
+        litellm_gateway::router_with_hooks(Arc::new(inference), &config, None, None, recorder);
+    let response = app
+        .oneshot(
+            Request::post("/v1/messages")
+                .header("content-type", "application/json")
+                .header("authorization", "Bearer gateway-key")
+                .body(Body::from(
+                    json!({
+                        "model": "public-model", "messages": [{"role": "user", "content": "hello"}],
+                        "max_tokens": 16,
+                    })
+                    .to_string(),
+                ))
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+    assert_eq!(response.status().as_u16(), status);
+    assert_eq!(
+        response.headers().contains_key("x-hook-completed"),
+        !reject_gateway
+    );
+    let body = to_bytes(response.into_body(), 4096).await.unwrap();
+    if reject_gateway {
+        assert_eq!(body, "rejected by gateway hook");
+    } else if !reject_router {
+        let response: Value = serde_json::from_slice(&body).unwrap();
+        assert_eq!(response["content"], message["content"]);
+        assert_eq!(response["usage"], message["usage"]);
+    }
+    assert_eq!(receiver.try_iter().collect::<Vec<_>>(), expected_phases);
+    assert_eq!(
+        upstream.received_requests().await.unwrap().len(),
+        usize::from(!reject_gateway && !reject_router)
+    );
+}

litellm-rust/crates/router/Cargo.toml

@@ -11,4 +11,5 @@ litellm-inference.workspace = true
 litellm-inference-messages.workspace = true
 
 [dev-dependencies]
+tokio.workspace = true
 rstest.workspace = true

litellm-rust/crates/router/src/hooks.rs

@@ -0,0 +1,15 @@
+use std::future::Future;
+
+use crate::Deployment;
+
+pub trait RouterHooks: Send + Sync {
+    fn filter_deployments<'a>(
+        &self,
+        _model: &str,
+        deployments: Vec<&'a Deployment>,
+    ) -> impl Future<Output = Vec<&'a Deployment>> + Send {
+        async { deployments }
+    }
+}
+
+impl RouterHooks for () {}

litellm-rust/crates/router/src/lib.rs

@@ -1,4 +1,6 @@
 mod deployment;
+mod hooks;
+pub use hooks::RouterHooks;
 
 use std::collections::HashMap;
 
@@ -7,7 +9,10 @@ use litellm_config::Model;
 pub use deployment::Deployment;
 
 #[derive(Clone, Debug, Default)]
-pub struct Router(HashMap<String, Deployment>);
+pub struct Router<H = ()> {
+    deployments: HashMap<String, Deployment>,
+    hooks: H,
+}
 
 impl Router {
     pub fn from_model_list(model_list: &[Model]) -> Self {
@@ -33,12 +38,35 @@ impl Router {
     }
 
     pub fn get(&self, model_name: &str) -> Option<&Deployment> {
-        self.0.get(model_name)
+        self.deployments.get(model_name)
+    }
+}
+
+impl<H> Router<H> {
+    pub fn with_hooks<T>(self, hooks: T) -> Router<T> {
+        Router {
+            deployments: self.deployments,
+            hooks,
+        }
+    }
+}
+
+impl<H: RouterHooks> Router<H> {
+    pub async fn select(&self, model_name: &str) -> Option<&Deployment> {
+        let candidates = self.deployments.get(model_name).into_iter().collect();
+        self.hooks
+            .filter_deployments(model_name, candidates)
+            .await
+            .into_iter()
+            .next()
     }
 }
 
 impl FromIterator<(String, Deployment)> for Router {
     fn from_iter<I: IntoIterator<Item = (String, Deployment)>>(entries: I) -> Self {
-        Self(entries.into_iter().collect())
+        Self {
+            deployments: entries.into_iter().collect(),
+            hooks: (),
+        }
     }
 }

litellm-rust/crates/router/tests/hooks.rs

@@ -0,0 +1,54 @@
+use litellm_router::{Deployment, Router, RouterHooks};
+use rstest::{fixture, rstest};
+
+#[fixture]
+fn router() -> Router {
+    Router::from_iter([(
+        "public-model".into(),
+        Deployment {
+            model: "provider/model".into(),
+            ..Default::default()
+        },
+    )])
+}
+
+struct Reject;
+
+impl RouterHooks for Reject {
+    async fn filter_deployments<'a>(
+        &self,
+        model: &str,
+        deployments: Vec<&'a Deployment>,
+    ) -> Vec<&'a Deployment> {
+        assert_eq!(model, "public-model");
+        assert_eq!(deployments.len(), 1);
+        assert_eq!(deployments[0].model, "provider/model");
+        Vec::new()
+    }
+}
+
+#[rstest]
+#[case::known("public-model")]
+#[case::unknown("missing")]
+#[tokio::test]
+async fn default_hooks_preserve_model_lookup(router: Router, #[case] model: &str) {
+    assert_eq!(
+        router
+            .select(model)
+            .await
+            .map(|deployment| &deployment.model),
+        router.get(model).map(|deployment| &deployment.model),
+    );
+}
+
+#[rstest]
+#[tokio::test]
+async fn selection_uses_filtered_candidates(router: Router) {
+    assert!(
+        router
+            .with_hooks(Reject)
+            .select("public-model")
+            .await
+            .is_none()
+    );
+}