fix(router): support model aliases end to end

A worker registers under a canonical model ID and may declare aliases
next to it. Routing only ever matched the canonical ID, so a request that
named an alias found no workers and got a 404.

The registry now keeps an alias index next to the model index. It is
separate so aliases stay out of `/v1/models` and out of the statistics,
and `get_by_model` falls back to it. Every other per-model map — hash
rings, retry overrides, load balancing policies — stays keyed by the
canonical ID, so each request entry point resolves the alias exactly once
and passes the canonical ID from there on:

- the gRPC pipeline resolves in `RequestContext::new` and rewrites both
  `input.model_id` and the request's own `model` field, so worker
  selection, tokenizer lookup, parser selection and tool call ID format
  all read the canonical ID without resolving anything themselves;
- both HTTP routers resolve at the top of their request path and rewrite
  the `model` field of the body they forward. The backend was registered
  under the canonical ID and has never heard of the alias.

One visible consequence: the response reports the canonical model ID, not
the alias the client sent, matching how the OpenAI API answers with the
model it actually ran.

`/v1/responses` gates on the client-supplied name before the pipeline
canonicalizes it, so it gets `WorkerRegistry::contains_model`, which
accepts both spellings. The `unknown` wildcard is not a registered name
and stays rejected.

Alias conflicts are resolved so that a name is either a canonical model
ID or an alias, never both:

- an alias naming a registered model is refused, and an alias already
  recorded is dropped when a model later claims that name — otherwise the
  name would start resolving to a different model once the real one's
  last worker left;
- when two models declare the same alias the first registration wins, and
  losing that worker hands the alias to a remaining model that declares
  it rather than stranding it;
- the alias entry is held across that whole decision, so a concurrent
  registration cannot be erased between the check and the delete.

Per-model retry overrides are keyed by the canonical ID and are read
before the pipeline canonicalizes, so `GrpcRouter::resolve_retry_config`
resolves the alias itself. Without it a request naming an alias silently
took the router default instead of the worker's override, on chat,
generate, messages and completion alike.

Tests cover the registry rules, both HTTP request paths end to end
against the body the worker actually received, the multipart form, the
rerank response, the gRPC pipeline, the Responses canonicalization, the
retry override under an alias, and the Responses gate.

Signed-off-by: Jun Liu <jun.c.liu@rakuten.com>

Author: StarDuster

model_gateway/src/routers/grpc/common/responses/utils.rs

@@ -98,14 +98,17 @@ pub(crate) async fn ensure_mcp_connection(
     Ok((false, Vec::new()))
 }
 
-/// Validate that workers are available for the requested model
+/// Validate that workers are available for the requested model.
+///
+/// Runs on the client-supplied name, before the pipeline canonicalizes it, so
+/// it has to accept aliases as well as canonical model IDs. `contains_model`
+/// covers both; listing `get_models()` and testing membership would reject
+/// every alias here.
 pub(crate) fn validate_worker_availability(
     worker_registry: &Arc<WorkerRegistry>,
     model: &str,
 ) -> Option<Response> {
-    let available_models = worker_registry.get_models();
-
-    if !available_models.contains(&model.to_string()) {
+    if !worker_registry.contains_model(model) {
         return Some(error::model_not_found(model));
     }
 
@@ -175,3 +178,47 @@ pub(crate) async fn persist_response_if_needed(
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use openai_protocol::{model_card::ModelCard, worker::HealthCheckConfig};
+
+    use super::*;
+    use crate::worker::{BasicWorkerBuilder, UNKNOWN_MODEL_ID};
+
+    fn registry_with_aliased_worker() -> Arc<WorkerRegistry> {
+        let registry = Arc::new(WorkerRegistry::new());
+        let worker = BasicWorkerBuilder::new("http://worker:8080")
+            .model(ModelCard::new("canonical-model").with_alias("model-alias"))
+            .health_config(HealthCheckConfig {
+                disable_health_check: true,
+                ..Default::default()
+            })
+            .build();
+        registry.register_or_replace(Arc::new(worker));
+        registry
+    }
+
+    #[test]
+    fn worker_availability_accepts_alias_and_preserves_unknown_rejection() {
+        let registry = registry_with_aliased_worker();
+
+        assert!(validate_worker_availability(&registry, "canonical-model").is_none());
+        assert!(validate_worker_availability(&registry, "model-alias").is_none());
+
+        let response = validate_worker_availability(&registry, UNKNOWN_MODEL_ID)
+            .expect("unknown model should remain rejected for Responses");
+        assert_eq!(response.status(), http::StatusCode::NOT_FOUND);
+    }
+
+    #[test]
+    fn worker_availability_rejects_alias_once_its_worker_is_gone() {
+        let registry = registry_with_aliased_worker();
+        let worker_id = registry.get_id_by_url("http://worker:8080").unwrap();
+        assert!(registry.remove(&worker_id).is_some());
+
+        let response = validate_worker_availability(&registry, "model-alias")
+            .expect("alias must stop resolving with no workers behind it");
+        assert_eq!(response.status(), http::StatusCode::NOT_FOUND);
+    }
+}

model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs

@@ -27,14 +27,16 @@ impl PipelineStage for DispatchMetadataStage {
         })?;
 
         let request_id = execution_plan.request_id().to_string();
+        // The model the response reports. `RequestContext::new` already
+        // rewrote every one of these to the canonical model ID, so a request
+        // that arrived under an alias is answered under the canonical name.
         let model = match &ctx.input.request_type {
             RequestType::Chat(req) => req.model.clone(),
             RequestType::Completion(req) => req.model.clone(),
-            RequestType::Generate(_req) => {
-                // Generate requests don't have a model field
-                // Use model_id from input
-                ctx.input.model_id.clone()
-            }
+            // `GenerateRequest` carries a model field too, but callers of the
+            // native `/generate` route may leave it empty, so prefer the
+            // model the router resolved.
+            RequestType::Generate(_req) => ctx.input.model_id.clone(),
             RequestType::Responses(req) => req.model.clone(),
             RequestType::Embedding(req) => req.model.clone(),
             RequestType::Classify(req) => req.model.clone(),

model_gateway/src/routers/grpc/context.rs

@@ -32,7 +32,7 @@ use super::{
 };
 use crate::{
     middleware::TenantRequestMeta,
-    worker::{RuntimeType, Worker, WorkerLoadGuard},
+    worker::{RuntimeType, Worker, WorkerLoadGuard, WorkerRegistry},
 };
 
 /// Main request processing context
@@ -50,6 +50,7 @@ pub(crate) struct RequestContext {
 pub(crate) struct RequestInput {
     pub request_type: RequestType,
     pub headers: Option<HeaderMap>,
+    /// Canonical model ID used after aliases are resolved at request entry.
     pub model_id: String,
     pub tenant_request_meta: Option<TenantRequestMeta>,
 }
@@ -67,6 +68,29 @@ pub(crate) enum RequestType {
 }
 
 impl RequestType {
+    /// Overwrite the request's own `model` field.
+    ///
+    /// Callers hold the request behind an `Arc` that the retry loop also
+    /// holds, so `Arc::make_mut` copies the request here. That cost is paid
+    /// only on the alias path — [`RequestContext::new`] skips this call
+    /// entirely when the client already used the canonical model ID.
+    fn set_model(&mut self, model_id: &str) {
+        fn replace(model: &mut String, model_id: &str) {
+            model.clear();
+            model.push_str(model_id);
+        }
+
+        match self {
+            Self::Chat(request) => replace(&mut Arc::make_mut(request).model, model_id),
+            Self::Generate(request) => replace(&mut Arc::make_mut(request).model, model_id),
+            Self::Completion(request) => replace(&mut Arc::make_mut(request).model, model_id),
+            Self::Responses(request) => replace(&mut Arc::make_mut(request).model, model_id),
+            Self::Embedding(request) => replace(&mut Arc::make_mut(request).model, model_id),
+            Self::Classify(request) => replace(&mut Arc::make_mut(request).model, model_id),
+            Self::Messages(request) => replace(&mut Arc::make_mut(request).model, model_id),
+        }
+    }
+
     /// Client-supplied backend request id (`rid`), where the protocol carries
     /// one. Responses ids are storage-owned (`resp_*`) and never client-set.
     pub fn rid(&self) -> Option<&str> {
@@ -112,6 +136,7 @@ impl std::fmt::Display for FinalResponse {
 /// Shared components (injected once at creation)
 pub(crate) struct SharedComponents {
     pub tokenizer_registry: Arc<TokenizerRegistry>,
+    pub worker_registry: Arc<WorkerRegistry>,
     pub tool_parser_factory: ToolParserFactory,
     pub reasoning_parser_factory: ReasoningParserFactory,
     /// Configured tool parser name (from CLI `--tool-call-parser`)
@@ -448,16 +473,32 @@ pub(crate) struct ResponseState {
 }
 
 impl RequestContext {
-    /// Create context for chat completion request
-    pub fn for_chat(
-        request: Arc<ChatCompletionRequest>,
+    /// Build a context, resolving a model alias to its canonical model ID.
+    ///
+    /// This is the single place the gRPC pipeline canonicalizes. Both
+    /// `input.model_id` and the request's own `model` field are rewritten, so
+    /// every stage below — worker selection, tokenizer lookup, parser
+    /// selection, tool call ID format — reads the canonical ID without
+    /// resolving anything itself.
+    ///
+    /// One visible consequence: the response reports the canonical model ID,
+    /// not the alias the client sent. That matches how the OpenAI API answers
+    /// with the model it actually ran.
+    fn new(
+        mut request_type: RequestType,
         headers: Option<HeaderMap>,
-        model_id: String,
+        mut model_id: String,
         components: Arc<SharedComponents>,
     ) -> Self {
+        if let Some(canonical_model_id) = components.worker_registry.resolve_model_alias(&model_id)
+        {
+            model_id.clear();
+            model_id.push_str(&canonical_model_id);
+            request_type.set_model(&model_id);
+        }
         Self {
             input: RequestInput {
-                request_type: RequestType::Chat(request),
+                request_type,
                 headers,
                 model_id,
                 tenant_request_meta: None,
@@ -467,23 +508,29 @@ impl RequestContext {
         }
     }
 
+    /// Create context for chat completion request
+    pub fn for_chat(
+        request: Arc<ChatCompletionRequest>,
+        headers: Option<HeaderMap>,
+        model_id: String,
+        components: Arc<SharedComponents>,
+    ) -> Self {
+        Self::new(RequestType::Chat(request), headers, model_id, components)
+    }
+
     /// Create context for generate request
     pub fn for_generate(
         request: Arc<GenerateRequest>,
         headers: Option<HeaderMap>,
         model_id: String,
         components: Arc<SharedComponents>,
     ) -> Self {
-        Self {
-            input: RequestInput {
-                request_type: RequestType::Generate(request),
-                headers,
-                model_id,
-                tenant_request_meta: None,
-            },
+        Self::new(
+            RequestType::Generate(request),
+            headers,
+            model_id,
             components,
-            state: ProcessingState::default(),
-        }
+        )
     }
 
     /// Create context for completion request
@@ -493,16 +540,12 @@ impl RequestContext {
         model_id: String,
         components: Arc<SharedComponents>,
     ) -> Self {
-        Self {
-            input: RequestInput {
-                request_type: RequestType::Completion(request),
-                headers,
-                model_id,
-                tenant_request_meta: None,
-            },
+        Self::new(
+            RequestType::Completion(request),
+            headers,
+            model_id,
             components,
-            state: ProcessingState::default(),
-        }
+        )
     }
 
     /// Create context for Responses API request
@@ -512,16 +555,12 @@ impl RequestContext {
         model_id: String,
         components: Arc<SharedComponents>,
     ) -> Self {
-        Self {
-            input: RequestInput {
-                request_type: RequestType::Responses(request),
-                headers,
-                model_id,
-                tenant_request_meta: None,
-            },
+        Self::new(
+            RequestType::Responses(request),
+            headers,
+            model_id,
             components,
-            state: ProcessingState::default(),
-        }
+        )
     }
 
     /// Create context for embedding request
@@ -531,16 +570,12 @@ impl RequestContext {
         model_id: String,
         components: Arc<SharedComponents>,
     ) -> Self {
-        Self {
-            input: RequestInput {
-                request_type: RequestType::Embedding(request),
-                headers,
-                model_id,
-                tenant_request_meta: None,
-            },
+        Self::new(
+            RequestType::Embedding(request),
+            headers,
+            model_id,
             components,
-            state: ProcessingState::default(),
-        }
+        )
     }
 
     /// Create context for classify request
@@ -550,16 +585,12 @@ impl RequestContext {
         model_id: String,
         components: Arc<SharedComponents>,
     ) -> Self {
-        Self {
-            input: RequestInput {
-                request_type: RequestType::Classify(request),
-                headers,
-                model_id,
-                tenant_request_meta: None,
-            },
+        Self::new(
+            RequestType::Classify(request),
+            headers,
+            model_id,
             components,
-            state: ProcessingState::default(),
-        }
+        )
     }
 
     /// Create context for messages request
@@ -569,16 +600,12 @@ impl RequestContext {
         model_id: String,
         components: Arc<SharedComponents>,
     ) -> Self {
-        Self {
-            input: RequestInput {
-                request_type: RequestType::Messages(request),
-                headers,
-                model_id,
-                tenant_request_meta: None,
-            },
+        Self::new(
+            RequestType::Messages(request),
+            headers,
+            model_id,
             components,
-            state: ProcessingState::default(),
-        }
+        )
     }
 
     /// Get chat request (panics if not chat)