Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 86 additions & 2 deletions crates/openshell-router/src/backend.rs
Original file line number Diff line number Diff line change
Expand Up @@ -982,8 +982,8 @@ fn is_vertex_anthropic_rawpredict_route(route: &ResolvedRoute) -> bool {
mod tests {
use super::{
ValidationFailure, ValidationFailureKind, build_backend_url, build_provider_url,
parse_bedrock_invocation_path, prepare_backend_request, rewrite_bedrock_path,
route_is_bedrock, verify_backend_endpoint,
parse_bedrock_invocation_path, prepare_backend_request, proxy_to_backend,
rewrite_bedrock_path, route_is_bedrock, verify_backend_endpoint,
};
use crate::RouterError;
use crate::config::{DEFAULT_ROUTE_TIMEOUT, ResolvedRoute};
Expand Down Expand Up @@ -2597,6 +2597,90 @@ mod tests {
);
}

/// Vertex AI's OpenAI-compatible endpoint requires the body `model` field to
/// carry a publisher prefix (e.g. `google/gemini-2.5-flash`). This test
/// simulates the fix: `resolve_vertex_ai_route` sets `route.model` to the
/// prefixed form, and the body rewrite here forwards that value to Vertex.
///
/// The mock server only accepts the prefixed form — matching Vertex's
/// behaviour — and returns 400 "Malformed publisher model" for the bare name.
#[tokio::test]
async fn vertex_openai_compat_rewrites_body_model_to_publisher_prefixed_form() {
let mock_server = MockServer::start().await;

// Simulate Vertex accepting only the publisher-prefixed model name.
Mock::given(method("POST"))
.and(path("/chat/completions"))
.and(body_partial_json(
serde_json::json!({"model": "google/gemini-2.5-flash"}),
))
.respond_with(
ResponseTemplate::new(200).set_body_json(
serde_json::json!({"choices": [{"message": {"content": "hi"}}]}),
),
)
.expect(1)
.mount(&mock_server)
.await;

// Simulate Vertex rejecting the bare model name — the pre-fix failure.
Mock::given(method("POST"))
.and(path("/chat/completions"))
.and(body_partial_json(
serde_json::json!({"model": "gemini-2.5-flash"}),
))
.respond_with(ResponseTemplate::new(400).set_body_json(
serde_json::json!({"error": {"message": "Malformed publisher model"}}),
))
.expect(0) // must never be reached after the fix
.mount(&mock_server)
.await;

// Route as produced by resolve_vertex_ai_route after the fix:
// route.model carries the publisher prefix.
let route = ResolvedRoute {
name: "vertex-gemini".to_string(),
endpoint: mock_server.uri(),
model: "google/gemini-2.5-flash".to_string(),
api_key: "ya29.token".to_string(),
protocols: vec!["openai_chat_completions".to_string()],
auth: AuthHeader::Bearer,
default_headers: vec![],
passthrough_headers: vec![],
timeout: DEFAULT_ROUTE_TIMEOUT,
model_in_path: false,
request_path_override: Some("/chat/completions".to_string()),
};

// The client sends the bare model name; the body rewrite must replace it
// with route.model (the publisher-prefixed form) before forwarding.
let client_body = serde_json::to_vec(&serde_json::json!({
"model": "gemini-2.5-flash",
"messages": [{"role": "user", "content": "hello"}]
}))
.unwrap();

let client = reqwest::Client::new();
let result = proxy_to_backend(
&client,
&route,
"openai_chat_completions",
"POST",
"/chat/completions",
vec![("content-type".to_string(), "application/json".to_string())],
bytes::Bytes::from(client_body),
)
.await
.expect("proxy should succeed");

assert_eq!(
result.status, 200,
"Vertex mock must accept the publisher-prefixed model; \
got {}: body rewrite did not apply the prefix",
result.status
);
}

/// Defense-in-depth: a Bedrock route receiving a non-Bedrock path
/// is rejected rather than forwarded. The L7 pattern detector
/// upstream of the router should never produce this combination,
Expand Down
128 changes: 124 additions & 4 deletions crates/openshell-server/src/inference.rs
Original file line number Diff line number Diff line change
Expand Up @@ -698,6 +698,22 @@ fn resolve_vertex_ai_route(
|p| p.eq_ignore_ascii_case("anthropic"),
);

// Vertex's OpenAI-compatible endpoint requires the request body's `model`
// field to carry a publisher prefix: `<publisher>/<model_id>` (e.g.
// `google/gemini-2.5-flash`). The publisher is taken from the explicit
// VERTEX_AI_PUBLISHER config value (when set to a non-Anthropic value) or
// inferred from the model name. For unrecognised models with no explicit
// publisher, the bare model ID is forwarded unchanged; Vertex will return
// a 400 in that case, which is the correct observable signal to the caller.
// Anthropic rawPredict routes encode the model in the URL path, not the
// body, so they are unaffected.
let body_model_id: String = if is_anthropic {
model_id.to_string()
} else {
let publisher = explicit_publisher.or_else(|| infer_vertex_publisher(model_id));
publisher.map_or_else(|| model_id.to_string(), |p| format!("{p}/{model_id}"))
};

// Escape hatch: caller-supplied full base URL still uses the model-derived
// protocol and path contract, but only for the OpenAI-compatible Vertex surface.
// Anthropic-on-Vertex needs model-path shaping and body adaptation that a fully
Expand All @@ -721,7 +737,7 @@ fn resolve_vertex_ai_route(
return Ok(build_vertex_route(
route_name,
base_url,
model_id,
&body_model_id,
Comment thread
johntmyers marked this conversation as resolved.
api_key,
vec!["openai_chat_completions".to_string()],
profile,
Expand Down Expand Up @@ -772,7 +788,7 @@ fn resolve_vertex_ai_route(
Ok(build_vertex_route(
route_name,
endpoint,
model_id,
&body_model_id,
api_key,
protocols,
profile,
Expand Down Expand Up @@ -1140,7 +1156,7 @@ async fn resolve_route_by_name_with_credentials(
Ok(Some(ResolvedRoute {
name: route_name.to_string(),
base_url: resolved.route.endpoint,
model_id: config.model_id.clone(),
model_id: resolved.route.model.clone(),
api_key: resolved.route.api_key,
protocols: resolved.route.protocols,
provider_type: resolved.provider_type,
Expand Down Expand Up @@ -1730,13 +1746,51 @@ mod tests {
route.request_path_override,
Some("/chat/completions".to_string())
);
assert_eq!(route.model_id, "gemini-2.0-flash-001");
assert_eq!(route.model_id, "google/gemini-2.0-flash-001");
assert_eq!(
route.base_url,
"https://us-central1-aiplatform.googleapis.com/v1beta1/projects/my-gcp-project/locations/us-central1/endpoints/openapi"
);
}

#[tokio::test]
async fn bundle_vertex_ai_non_anthropic_model_id_carries_publisher_prefix() {
// Regression test: the bundle's model_id must carry the publisher prefix
// so the router sends e.g. "google/gemini-2.5-flash" in the request body,
// not the bare "gemini-2.5-flash" that Vertex AI rejects with HTTP 400.
let store = test_store().await;
let config = [
(
"VERTEX_AI_PROJECT_ID".to_string(),
"my-gcp-project".to_string(),
),
("VERTEX_AI_REGION".to_string(), "us-central1".to_string()),
]
.into_iter()
.collect();
let provider = make_vertex_provider_with_config("vertex-dev", config);
store
.put_message(&provider)
.await
.expect("persist provider");
let route = make_route(
CLUSTER_INFERENCE_ROUTE_NAME,
"vertex-dev",
"gemini-2.5-flash",
);
store.put_message(&route).await.expect("persist route");

let resp = resolve_inference_bundle(&store, "default")
.await
.expect("bundle should resolve");

assert_eq!(resp.routes.len(), 1);
assert_eq!(
resp.routes[0].model_id, "google/gemini-2.5-flash",
"bundle model_id must carry publisher prefix for non-Anthropic Vertex routes"
);
}

#[tokio::test]
async fn bundle_without_cluster_route_returns_empty_routes() {
let store = test_store().await;
Expand Down Expand Up @@ -2613,6 +2667,11 @@ mod tests {
.contains(&"anthropic_messages".to_string()),
"must not have anthropic_messages protocol for gemini"
);
// Vertex OpenAI-compatible endpoint requires publisher prefix in body model field
assert_eq!(
resolved.route.model, "google/gemini-pro",
"Vertex non-Anthropic body model must carry publisher prefix"
);
}

#[test]
Expand Down Expand Up @@ -2651,6 +2710,67 @@ mod tests {
.contains(&"anthropic_messages".to_string()),
"must not have anthropic_messages for unknown model"
);
// Unknown models have no inferred publisher; body model ID is unchanged
assert_eq!(resolved.route.model, "some-unknown-model");
}

#[test]
fn resolve_vertex_ai_route_non_anthropic_publisher_prefix_gemini() {
// Gemini models must get `google/<model>` in route.model so the
// OpenAI-compatible Vertex endpoint accepts the request body.
let config =
std::iter::once(("VERTEX_AI_PROJECT_ID".to_string(), "proj-abc".to_string())).collect();
let provider = make_vertex_provider_with_config("vertex-gemini-flash", config);

let resolved =
resolve_provider_route(&provider, "gemini-2.5-flash").expect("should resolve");

assert_eq!(resolved.route.model, "google/gemini-2.5-flash");
}

#[test]
fn resolve_vertex_ai_route_non_anthropic_publisher_prefix_llama() {
let config =
std::iter::once(("VERTEX_AI_PROJECT_ID".to_string(), "proj-abc".to_string())).collect();
let provider = make_vertex_provider_with_config("vertex-llama", config);

let resolved = resolve_provider_route(&provider, "llama-3-70b").expect("should resolve");

assert_eq!(resolved.route.model, "meta/llama-3-70b");
}

#[test]
fn resolve_vertex_ai_route_explicit_publisher_overrides_inference() {
// VERTEX_AI_PUBLISHER takes precedence over infer_vertex_publisher for
// unknown model names.
let config = [
("VERTEX_AI_PROJECT_ID".to_string(), "proj-abc".to_string()),
("VERTEX_AI_PUBLISHER".to_string(), "acme".to_string()),
]
.into_iter()
.collect();
let provider = make_vertex_provider_with_config("vertex-explicit", config);

let resolved =
resolve_provider_route(&provider, "some-acme-model").expect("should resolve");

assert_eq!(resolved.route.model, "acme/some-acme-model");
}

#[test]
fn resolve_vertex_ai_route_base_url_override_gemini_gets_publisher_prefix() {
// Publisher prefix must also be applied when a base URL override is used.
let config = std::iter::once((
"VERTEX_AI_BASE_URL".to_string(),
"https://us-central1-aiplatform.googleapis.com/v1beta1/projects/my-project/locations/us-central1/endpoints/openapi".to_string(),
))
.collect();
let provider = make_vertex_provider_with_config("vertex-base-url-gemini", config);

let resolved =
resolve_provider_route(&provider, "gemini-2.0-flash").expect("should resolve");

assert_eq!(resolved.route.model, "google/gemini-2.0-flash");
}

#[test]
Expand Down
Loading