diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 226f1577..ad10d12b 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -521,9 +521,7 @@ async fn dispatch( let provider = crate::dispatch::require_provider(model).map_err(with_model)?; let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model).map_err(with_model)?; - let bridge = state - .hub - .get(provider) + let bridge = crate::dispatch::resolve_bridge(&state.hub, &pk_entry.value, provider) .ok_or_else(|| with_model(ProxyError::ProviderUnavailable))?; let model_arc = Arc::new(model.clone()); let pk_arc = Arc::new(pk_entry.value.clone()); @@ -814,7 +812,12 @@ async fn dispatch( continue; } }; - let Some(bridge) = state.hub.get(provider) else { + // Phase D cutover join point: try two-tier (specialized vendor + // → adapter family) first; fall back to the legacy + // Provider-keyed registry when the new fields aren't filled in + // on this PK yet. See crate::dispatch::resolve_bridge. + let Some(bridge) = crate::dispatch::resolve_bridge(&state.hub, &pk_entry.value, provider) + else { last_err = Some(BridgeError::Config( "no bridge registered for provider".into(), )); diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index 11c4f690..d621d858 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -116,9 +116,7 @@ async fn dispatch( let provider = crate::dispatch::require_provider(model)?; let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model)?; - let bridge = state - .hub - .get(provider) + let bridge = crate::dispatch::resolve_bridge(&state.hub, &pk_entry.value, provider) .ok_or(ProxyError::ProviderUnavailable)?; let model_arc = Arc::new(model.clone()); diff --git a/crates/aisix-proxy/src/dispatch.rs b/crates/aisix-proxy/src/dispatch.rs index 5614be62..308f69f8 100644 --- a/crates/aisix-proxy/src/dispatch.rs +++ b/crates/aisix-proxy/src/dispatch.rs @@ -22,9 +22,49 @@ use std::sync::Arc; use aisix_core::models::Provider; use aisix_core::resource::ResourceEntry; use aisix_core::{AisixSnapshot, Model, ProviderKey}; +use aisix_gateway::{Bridge, Hub}; use crate::error::ProxyError; +/// Resolve the Bridge to dispatch this request through. +/// +/// Tries the two-tier dispatch first (specialized vendor → adapter +/// family), and falls back to the legacy `Provider`-keyed registry if +/// neither tier matches. This is the Phase D cutover join point: new +/// `ProviderKey` payloads carrying `adapter` and `provider` hit the +/// two-tier path; pre-cutover payloads (`adapter: None`, empty +/// `provider`) fall through to the legacy registry unchanged. +/// +/// Returns `None` only when both layers miss — i.e. the operator has no +/// bridge wired for this request at all. +pub(crate) fn resolve_bridge( + hub: &Hub, + provider_key: &ProviderKey, + provider: Provider, +) -> Option> { + if let Some(bridge) = hub.dispatch_two_tier(provider_key) { + return Some(bridge); + } + // Legacy fallback. Emit a tracing::debug! when the PK carries the + // new-shape fields (`provider` or `adapter`) but the two-tier path + // still missed — post-cutover this is the early signal that a + // specialized Bridge name was misregistered (typo, runtime + // unregister) or that the adapter map missed an entry. Pre-cutover + // PKs (empty provider + adapter: None) take the silent path because + // they're the dominant case. + let pk_carries_new_fields = !provider_key.provider.is_empty() || provider_key.adapter.is_some(); + if pk_carries_new_fields { + tracing::debug!( + target: "aisix_proxy::dispatch", + pk_provider = %provider_key.provider, + pk_adapter = ?provider_key.adapter, + legacy_provider = ?provider, + "two-tier dispatch missed for new-shape ProviderKey; falling back to legacy hub.get" + ); + } + hub.get(provider) +} + /// Look up the `ProviderKey` a given `Model` references. Returns a /// 400 if the Model is a virtual router (those don't dispatch /// directly — caller should walk `routing.targets` first), or if the @@ -466,4 +506,137 @@ mod tests { // Misuse — handlers should always pass a `/`-prefixed path. let _ = build_v1_url("https://api.openai.com", "responses"); } + + // --- resolve_bridge tests ------------------------------------- + // + // Cover the three reachable outcomes of resolve_bridge: + // 1. specialized hit — pk.provider matches a specialized entry + // 2. family hit — pk.adapter matches a family entry, + // specialized misses + // 3. legacy fallback — both new-tier maps miss, legacy hub.get + // serves the bridge + // 4. none miss — nothing registered at all + // + // A minimal Bridge stub is used so the test doesn't need reqwest + // or a real upstream. + + mod resolve_bridge_tests { + use super::*; + use aisix_core::models::Adapter; + use aisix_gateway::{ + Bridge, BridgeContext, BridgeError, ChatChunkStream, ChatFormat, ChatMessage, + ChatResponse, EmbeddingRequest, EmbeddingResponse, FinishReason, Hub, UsageStats, + }; + use async_trait::async_trait; + use futures::stream; + + /// Minimal Bridge that records its identity via `name()`. Lets + /// resolve_bridge tests verify which Bridge was returned without + /// dragging in reqwest. + struct StubBridge { + name: &'static str, + } + + #[async_trait] + impl Bridge for StubBridge { + fn name(&self) -> &'static str { + self.name + } + + async fn chat( + &self, + req: &ChatFormat, + _ctx: &BridgeContext, + ) -> Result { + Ok(ChatResponse { + id: "stub".into(), + model: req.model.clone(), + message: ChatMessage::assistant("stub"), + finish_reason: FinishReason::Stop, + usage: UsageStats::new(0, 0), + }) + } + + async fn chat_stream( + &self, + _req: &ChatFormat, + _ctx: &BridgeContext, + ) -> Result { + Ok(Box::pin(stream::iter(Vec::new()))) + } + + async fn embed( + &self, + _req: &EmbeddingRequest, + _ctx: &BridgeContext, + ) -> Result { + Err(BridgeError::Config("stub".into())) + } + } + + /// Build a ProviderKey JSON with the new-shape fields. `adapter` + /// is passed as the kebab-case wire string (`"openai"` / + /// `"azure-openai"` etc.) rather than the enum, to keep the + /// helper independent of any `as_str()` method on `Adapter`. + fn pk_with_provider_and_adapter(provider: &str, adapter: Option<&str>) -> ProviderKey { + let adapter_json = match adapter { + Some(a) => format!(", \"adapter\":\"{a}\""), + None => String::new(), + }; + let cfg = format!( + r#"{{"display_name":"x","secret":"k","provider":"{provider}"{adapter_json}}}"# + ); + serde_json::from_str(&cfg).unwrap() + } + + #[test] + fn specialized_hit_wins_over_family_and_legacy() { + let hub = Hub::new(); + hub.register_specialized( + "deepseek", + Arc::new(StubBridge { + name: "specialized", + }), + ); + hub.register_family(Adapter::Openai, Arc::new(StubBridge { name: "family" })); + hub.register(Provider::Openai, Arc::new(StubBridge { name: "legacy" })); + + let pk = pk_with_provider_and_adapter("deepseek", Some("openai")); + let bridge = resolve_bridge(&hub, &pk, Provider::Openai).unwrap(); + assert_eq!(bridge.name(), "specialized"); + } + + #[test] + fn family_hit_when_specialized_misses() { + let hub = Hub::new(); + hub.register_family(Adapter::Openai, Arc::new(StubBridge { name: "family" })); + hub.register(Provider::Openai, Arc::new(StubBridge { name: "legacy" })); + + // pk.provider = "unknown-vendor" → no specialized; pk.adapter + // = Openai → family hit. + let pk = pk_with_provider_and_adapter("unknown-vendor", Some("openai")); + let bridge = resolve_bridge(&hub, &pk, Provider::Openai).unwrap(); + assert_eq!(bridge.name(), "family"); + } + + #[test] + fn legacy_fallback_when_both_new_tiers_miss() { + let hub = Hub::new(); + // No family / specialized registered — only legacy. + hub.register(Provider::Openai, Arc::new(StubBridge { name: "legacy" })); + + // Today's on-disk PK shape: empty provider + adapter: None. + // The two-tier path returns None on both layers; legacy serves. + let pk = pk_with_provider_and_adapter("", None); + let bridge = resolve_bridge(&hub, &pk, Provider::Openai).unwrap(); + assert_eq!(bridge.name(), "legacy"); + } + + #[test] + fn none_when_nothing_registered() { + let hub = Hub::new(); + let pk = pk_with_provider_and_adapter("", None); + assert!(resolve_bridge(&hub, &pk, Provider::Openai).is_none()); + } + } } diff --git a/crates/aisix-proxy/src/embeddings.rs b/crates/aisix-proxy/src/embeddings.rs index 13b438cc..ac3ecf15 100644 --- a/crates/aisix-proxy/src/embeddings.rs +++ b/crates/aisix-proxy/src/embeddings.rs @@ -136,9 +136,7 @@ async fn dispatch( let provider = crate::dispatch::require_provider(model)?; let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model)?; - let bridge = state - .hub - .get(provider) + let bridge = crate::dispatch::resolve_bridge(&state.hub, &pk_entry.value, provider) .ok_or(ProxyError::ProviderUnavailable)?; let model_rl = diff --git a/crates/aisix-proxy/src/images.rs b/crates/aisix-proxy/src/images.rs index c77eefde..6cabc6bb 100644 --- a/crates/aisix-proxy/src/images.rs +++ b/crates/aisix-proxy/src/images.rs @@ -131,9 +131,7 @@ async fn dispatch( let provider = crate::dispatch::require_provider(model)?; let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model)?; - let bridge = state - .hub - .get(provider) + let bridge = crate::dispatch::resolve_bridge(&state.hub, &pk_entry.value, provider) .ok_or(ProxyError::ProviderUnavailable)?; let model_arc = Arc::new(model.clone()); diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 9b101359..99c43469 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -422,10 +422,9 @@ async fn cross_provider_dispatch( let provider = model.provider.ok_or_else(|| { ProxyError::InvalidRequest(format!("model `{model_name}` has no provider prefix")) })?; - let bridge: Arc = state - .hub - .get(provider) - .ok_or(ProxyError::ProviderUnavailable)?; + let bridge: Arc = + crate::dispatch::resolve_bridge(&state.hub, provider_key, provider) + .ok_or(ProxyError::ProviderUnavailable)?; // Parse the Anthropic-shape body into the gateway's normalised // ChatFormat. Errors here are 400 — the request is malformed