auth: support model alias claims with body rewrite

The platform's MODEL_HF_PATH_OVERRIDES alias was previously one-way: the
JWT model claim and the orchestrator's TOML model.name both held the
canonical HF path, but envs that hard-code the user-facing name (e.g.
emilschmitz/curriculum-oversight's LLM judge sending
sprints/Llama-3.2-1B-Instruct) hit 403 because the canonical-only claim
required exact-string match.

RftClaims gains an optional model_aliases list. allows_model accepts
alias hits, and pin_and_check_model variants (typed, string, JSON) plus
the transparent-proxy fallback rewrite the request body's model field
to self.model after an alias-only match, so vLLM only ever sees the
canonical served name. Base-model and LoRA matches pass through
unchanged.

Author: JannikSt

src/auth.rs

@@ -17,6 +17,14 @@ pub struct RftClaims {
     /// Allowed LoRA adapter name
     #[serde(default)]
     pub lora: Option<String>,
+    /// Alternate names that should resolve to the base `model`. Used when
+    /// the platform exposes a user-facing model identifier (e.g.
+    /// `sprints/Llama-3.2-1B-Instruct`) that differs from the canonical
+    /// HF path vLLM serves (`meta-llama/Llama-3.2-1B-Instruct`). A
+    /// request hitting an alias is authorized and the request body's
+    /// `model` field is rewritten to `self.model` before dispatch.
+    #[serde(default)]
+    pub model_aliases: Vec<String>,
 }
 
 /// Verifier for RS256 JWTs signed by the platform.
@@ -73,6 +81,10 @@ impl RftClaims {
             }
         }
 
+        if self.is_model_alias(requested) {
+            return true;
+        }
+
         if let Some(lora) = self.lora.as_deref() {
             // Empty lora claim must never authorize anything; an empty
             // string is a prefix of every other string, which would let
@@ -96,19 +108,52 @@ impl RftClaims {
 
         false
     }
+
+    /// Whether `requested` is one of the alternate names declared in
+    /// `model_aliases`. Empty entries are ignored so a misconfigured
+    /// claim never authorizes the empty model.
+    fn is_model_alias(&self, requested: &str) -> bool {
+        if requested.is_empty() {
+            return false;
+        }
+        self.model_aliases
+            .iter()
+            .any(|a| !a.is_empty() && a == requested)
+    }
+
+    /// If `requested` matched the JWT only via `model_aliases`, return
+    /// the canonical model name (`self.model`) so the request body can
+    /// be rewritten before forwarding to vLLM. Returns `None` if the
+    /// request already targets the base model or a LoRA — those need to
+    /// pass through unchanged.
+    pub fn canonical_for_alias(&self, requested: &str) -> Option<String> {
+        if !self.is_model_alias(requested) {
+            return None;
+        }
+        let base = self.model.as_deref()?;
+        if base.is_empty() || base == requested {
+            return None;
+        }
+        Some(base.to_string())
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
     fn claims(model: Option<&str>, lora: Option<&str>) -> RftClaims {
+        claims_with_aliases(model, lora, &[])
+    }
+
+    fn claims_with_aliases(model: Option<&str>, lora: Option<&str>, aliases: &[&str]) -> RftClaims {
         RftClaims {
             sub: "user".into(),
             run_id: "abc".into(),
             team_id: String::new(),
             model: model.map(String::from),
             lora: lora.map(String::from),
+            model_aliases: aliases.iter().map(|s| (*s).to_string()).collect(),
         }
     }
 
@@ -165,4 +210,51 @@ mod tests {
         let c = claims(Some(""), Some("rft-abc"));
         assert!(!c.allows_model(""));
     }
+
+    #[test]
+    fn allows_model_alias() {
+        let c = claims_with_aliases(
+            Some("meta-llama/Llama-3.2-1B-Instruct"),
+            Some("rft-abc"),
+            &["sprints/Llama-3.2-1B-Instruct"],
+        );
+        assert!(c.allows_model("sprints/Llama-3.2-1B-Instruct"));
+        assert!(c.allows_model("meta-llama/Llama-3.2-1B-Instruct"));
+        assert!(!c.allows_model("other/model"));
+    }
+
+    #[test]
+    fn canonical_for_alias_rewrites_alias_only() {
+        let c = claims_with_aliases(
+            Some("meta-llama/Llama-3.2-1B-Instruct"),
+            Some("rft-abc"),
+            &["sprints/Llama-3.2-1B-Instruct"],
+        );
+        assert_eq!(
+            c.canonical_for_alias("sprints/Llama-3.2-1B-Instruct").as_deref(),
+            Some("meta-llama/Llama-3.2-1B-Instruct"),
+        );
+        // Base model and lora must NOT be rewritten — lora dispatch
+        // depends on the original name reaching vLLM.
+        assert_eq!(c.canonical_for_alias("meta-llama/Llama-3.2-1B-Instruct"), None);
+        assert_eq!(c.canonical_for_alias("rft-abc"), None);
+        assert_eq!(c.canonical_for_alias("rft-abc-step-42"), None);
+        assert_eq!(c.canonical_for_alias("unrelated"), None);
+    }
+
+    #[test]
+    fn empty_alias_entry_authorizes_nothing() {
+        let c = claims_with_aliases(Some("meta-llama/Llama-3.2-1B-Instruct"), Some("rft-abc"), &[""]);
+        assert!(!c.allows_model(""));
+        assert_eq!(c.canonical_for_alias(""), None);
+    }
+
+    #[test]
+    fn alias_matching_base_does_not_rewrite() {
+        // Pathological config: alias equals the base. Should still
+        // authorize, but not produce a rewrite (would be a no-op anyway).
+        let c = claims_with_aliases(Some("meta-llama/Llama-3.2-1B-Instruct"), None, &["meta-llama/Llama-3.2-1B-Instruct"]);
+        assert!(c.allows_model("meta-llama/Llama-3.2-1B-Instruct"));
+        assert_eq!(c.canonical_for_alias("meta-llama/Llama-3.2-1B-Instruct"), None);
+    }
 }

src/server.rs

@@ -170,7 +170,7 @@ async fn transparent_proxy_handler(State(state): State<Arc<AppState>>, req: Requ
     };
 
     // Parse body as JSON
-    let body_json: serde_json::Value = if body_bytes.is_empty() {
+    let mut body_json: serde_json::Value = if body_bytes.is_empty() {
         serde_json::Value::Null
     } else {
         match serde_json::from_slice(&body_bytes) {
@@ -193,8 +193,9 @@ async fn transparent_proxy_handler(State(state): State<Arc<AppState>>, req: Requ
             .get("model")
             .and_then(|v| v.as_str())
             .filter(|s| !s.is_empty())
+            .map(str::to_string)
         {
-            if !claims_ref.allows_model(model) {
+            if !claims_ref.allows_model(&model) {
                 warn!(
                     run_id = %claims_ref.run_id,
                     requested_model = %model,
@@ -209,6 +210,14 @@ async fn transparent_proxy_handler(State(state): State<Arc<AppState>>, req: Requ
                 )
                     .into_response();
             }
+            if let Some(canonical) = claims_ref.canonical_for_alias(&model) {
+                if let Some(obj) = body_json.as_object_mut() {
+                    obj.insert(
+                        "model".to_string(),
+                        serde_json::Value::String(canonical),
+                    );
+                }
+            }
         }
         if let Err(response) = enforce_no_lora_path_override_json(&claims, &body_json) {
             return response;
@@ -445,6 +454,9 @@ fn pin_and_check_model(
 
     let resolved = model.as_deref().unwrap_or("");
     if claims.allows_model(resolved) {
+        if let Some(canonical) = claims.canonical_for_alias(resolved) {
+            *model = Some(canonical);
+        }
         return Ok(());
     }
 
@@ -500,6 +512,9 @@ fn pin_and_check_model_string(
     }
 
     if claims.allows_model(model) {
+        if let Some(canonical) = claims.canonical_for_alias(model) {
+            *model = canonical;
+        }
         return Ok(());
     }
 
@@ -558,6 +573,14 @@ fn pin_and_check_model_json(
     };
 
     if claims.allows_model(&resolved) {
+        if let Some(canonical) = claims.canonical_for_alias(&resolved) {
+            if let Some(obj) = body.as_object_mut() {
+                obj.insert(
+                    "model".to_string(),
+                    serde_json::Value::String(canonical),
+                );
+            }
+        }
         return Ok(());
     }