fix: page through all Databricks AI Gateway v2 endpoints when listing models (#9753)

Signed-off-by: Kalvin Chau <kalvin@block.xyz>
This commit is contained in:
Kalvin C
2026-06-11 16:53:02 -07:00
committed by GitHub
parent 86e2f04264
commit f23347fc30
+60 -26
View File
@@ -37,6 +37,7 @@ use rmcp::model::Tool;
const DATABRICKS_V2_PROVIDER_NAME: &str = "databricks_v2";
const DATABRICKS_V2_LIST_ENDPOINTS_PATH: &str = "api/ai-gateway/v2/endpoints";
const DATABRICKS_V2_LIST_ENDPOINTS_PAGE_SIZE: usize = 100;
pub const DATABRICKS_V2_DEFAULT_MODEL: &str = "databricks-gpt-5-5";
pub const DATABRICKS_V2_KNOWN_MODELS: &[&str] =
&["databricks-gpt-5-5", "databricks-claude-opus-4-7"];
@@ -189,7 +190,9 @@ impl DatabricksV2Provider {
model_name.contains("claude")
}
fn parse_list_endpoints_response(json: &Value) -> Result<Vec<String>, ProviderError> {
fn parse_list_endpoints_response(
json: &Value,
) -> Result<(Vec<String>, Option<String>), ProviderError> {
let endpoints = json
.get("endpoints")
.and_then(|v| v.as_array())
@@ -200,7 +203,7 @@ impl DatabricksV2Provider {
)
})?;
let mut models: Vec<String> = endpoints
let models: Vec<String> = endpoints
.iter()
.filter_map(|endpoint| {
endpoint
@@ -209,8 +212,14 @@ impl DatabricksV2Provider {
.map(str::to_string)
})
.collect();
models.sort();
Ok(models)
let next_page_token = json
.get("next_page_token")
.and_then(|v| v.as_str())
.filter(|token| !token.is_empty())
.map(str::to_string);
Ok((models, next_page_token))
}
async fn stream_openai_responses(
@@ -417,31 +426,53 @@ impl Provider for DatabricksV2Provider {
}
async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
let response = self
.api_client
.response_get(None, DATABRICKS_V2_LIST_ENDPOINTS_PATH)
.await
.map_err(|e| {
let mut models = Vec::new();
let mut page_token: Option<String> = None;
loop {
let mut path = format!(
"{}?page_size={}",
DATABRICKS_V2_LIST_ENDPOINTS_PATH, DATABRICKS_V2_LIST_ENDPOINTS_PAGE_SIZE
);
if let Some(token) = &page_token {
path.push_str(&format!("&page_token={}", urlencoding::encode(token)));
}
let response = self
.api_client
.response_get(None, &path)
.await
.map_err(|e| {
ProviderError::RequestFailed(format!(
"Failed to fetch Databricks AI Gateway endpoints: {e}"
))
})?;
if !response.status().is_success() {
let status = response.status();
let detail = response.text().await.unwrap_or_default();
return Err(ProviderError::RequestFailed(format!(
"Failed to fetch Databricks AI Gateway endpoints: {status} {detail}"
)));
}
let json: Value = response.json().await.map_err(|e| {
ProviderError::RequestFailed(format!(
"Failed to fetch Databricks AI Gateway endpoints: {e}"
"Failed to parse Databricks AI Gateway endpoints response: {e}"
))
})?;
if !response.status().is_success() {
let status = response.status();
let detail = response.text().await.unwrap_or_default();
return Err(ProviderError::RequestFailed(format!(
"Failed to fetch Databricks AI Gateway endpoints: {status} {detail}"
)));
let (page_models, next_page_token) = Self::parse_list_endpoints_response(&json)?;
models.extend(page_models);
if next_page_token.is_none() || next_page_token == page_token {
break;
}
page_token = next_page_token;
}
let json: Value = response.json().await.map_err(|e| {
ProviderError::RequestFailed(format!(
"Failed to parse Databricks AI Gateway endpoints response: {e}"
))
})?;
Self::parse_list_endpoints_response(&json)
models.sort();
Ok(models)
}
}
@@ -480,19 +511,22 @@ mod tests {
{"name": "databricks-claude-opus-4-7"},
{"name": "databricks-gpt-5-5"},
{"name": "custom-model"}
]
],
"next_page_token": "tok"
});
let models = DatabricksV2Provider::parse_list_endpoints_response(&json).unwrap();
let (models, next_page_token) =
DatabricksV2Provider::parse_list_endpoints_response(&json).unwrap();
assert_eq!(
models,
vec![
"custom-model".to_string(),
"databricks-claude-opus-4-7".to_string(),
"databricks-gpt-5-5".to_string(),
"custom-model".to_string(),
]
);
assert_eq!(next_page_token.as_deref(), Some("tok"));
}
#[test]