fix(gateway): use Vertex model garden catalog endpoint

This commit is contained in:
MMEXA
2026-05-18 16:51:12 +00:00
parent b480f3aaff
commit 5ed8325592
2 changed files with 37 additions and 56 deletions

View File

@@ -464,8 +464,6 @@ async fn fetch_vertex_service_account_models(
});
};
let token = exchange_vertex_service_account_token(runtime, &transports[0], auth_config).await?;
let project_id = json_string(auth_config.get("project_id"))
.ok_or_else(|| "vertex_ai(service_account): missing project_id".to_string())?;
let gemini_transport =
select_transport_for_api_format(transports, "gemini:").unwrap_or(&transports[0]);
let claude_transport =
@@ -476,18 +474,12 @@ async fn fetch_vertex_service_account_models(
let mut soft_errors = Vec::new();
let mut has_success = false;
for region in vertex_regions(auth_config) {
let base = if region == "global" {
VERTEX_API_BASE_URL.to_string()
} else {
format!("https://{region}-aiplatform.googleapis.com")
};
for base in iter_vertex_base_urls(transports) {
for (publisher, transport, api_format) in [
("google", gemini_transport, "gemini:generate_content"),
("anthropic", claude_transport, "claude:messages"),
] {
let url =
build_vertex_service_account_list_url(&base, &project_id, &region, publisher, None);
let url = build_vertex_service_account_list_url(&base, publisher, None);
let outcome = fetch_vertex_models_from_url(
runtime,
transport,
@@ -929,14 +921,7 @@ fn iter_vertex_base_urls(transports: &[GatewayProviderTransportSnapshot]) -> Vec
}
fn build_vertex_google_list_url(base_url: &str, api_key: &str, page_token: Option<&str>) -> String {
let path = if base_url.trim_end_matches('/').ends_with("/v1")
|| base_url.trim_end_matches('/').ends_with("/v1beta")
{
"/publishers/google/models"
} else {
"/v1/publishers/google/models"
};
let url = build_simple_path_url(base_url, path);
let url = build_vertex_publisher_models_list_base_url(base_url, "google");
let mut url = append_query_param(url, "key", api_key);
url = append_query_param(url, "pageSize", VERTEX_PAGE_SIZE);
if let Some(page_token) = page_token {
@@ -947,14 +932,10 @@ fn build_vertex_google_list_url(base_url: &str, api_key: &str, page_token: Optio
fn build_vertex_service_account_list_url(
base_url: &str,
project_id: &str,
region: &str,
publisher: &str,
page_token: Option<&str>,
) -> String {
let path =
format!("/v1/projects/{project_id}/locations/{region}/publishers/{publisher}/models");
let mut url = build_simple_path_url(base_url, &path);
let mut url = build_vertex_publisher_models_list_base_url(base_url, publisher);
url = append_query_param(url, "pageSize", VERTEX_PAGE_SIZE);
if let Some(page_token) = page_token {
url = append_query_param(url, "pageToken", page_token);
@@ -962,6 +943,16 @@ fn build_vertex_service_account_list_url(
url
}
fn build_vertex_publisher_models_list_base_url(base_url: &str, publisher: &str) -> String {
let trimmed_base = base_url.trim().trim_end_matches('/');
let path = if trimmed_base.ends_with("/v1") || trimmed_base.ends_with("/v1beta1") {
format!("/publishers/{publisher}/models")
} else {
format!("/v1beta1/publishers/{publisher}/models")
};
build_simple_path_url(trimmed_base, &path)
}
fn build_simple_path_url(base_url: &str, path: &str) -> String {
format!("{}{}", base_url.trim().trim_end_matches('/'), path.trim())
}
@@ -1087,38 +1078,6 @@ fn vertex_effective_format(model_id: &str, auth_config: Option<&Value>) -> Strin
}
}
fn vertex_regions(auth_config: &Value) -> Vec<String> {
let mut seen = BTreeSet::new();
let mut regions = Vec::new();
let auth_config = auth_config.as_object();
let mut push_region = |region: Option<&str>| {
let Some(region) = region.map(str::trim).filter(|value| !value.is_empty()) else {
return;
};
if seen.insert(region.to_string()) {
regions.push(region.to_string());
}
};
push_region(
auth_config
.and_then(|value| value.get("region"))
.and_then(Value::as_str),
);
if let Some(model_regions) = auth_config
.and_then(|value| value.get("model_regions"))
.and_then(Value::as_object)
{
for value in model_regions.values() {
push_region(value.as_str());
}
}
push_region(Some("global"));
push_region(Some("us-central1"));
regions
}
fn is_soft_not_found(error: &str) -> bool {
error.trim().starts_with("HTTP 404:")
}
@@ -1498,7 +1457,7 @@ mod tests {
let urls = executed_urls.lock().expect("executed_urls lock");
assert_eq!(
urls.as_slice(),
&["https://aiplatform.googleapis.com/v1/publishers/google/models?key=vertex-secret&pageSize=100"]
&["https://aiplatform.googleapis.com/v1beta1/publishers/google/models?key=vertex-secret&pageSize=100"]
);
assert_eq!(outcome.fetched_model_ids, vec!["gemini-3.1-pro-preview"]);
assert_eq!(outcome.cached_models.len(), 1);
@@ -1508,6 +1467,24 @@ mod tests {
);
}
#[test]
fn vertex_publisher_models_list_url_uses_model_garden_resource_not_runtime_resource() {
let url = super::build_vertex_service_account_list_url(
"https://aiplatform.googleapis.com",
"google",
None,
);
assert_eq!(
url,
"https://aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=100"
);
assert!(
!url.contains("/projects/") && !url.contains("/locations/"),
"Model Garden publisher list must not use Vertex runtime project/location path"
);
}
#[tokio::test]
async fn codex_transport_fetches_upstream_models_instead_of_preset_catalog() {
let executed_urls = Arc::new(Mutex::new(Vec::new()));