diff --git a/crates/goose/src/providers/databricks_v2.rs b/crates/goose/src/providers/databricks_v2.rs index 56a20e24b94c..9d0bc276659b 100644 --- a/crates/goose/src/providers/databricks_v2.rs +++ b/crates/goose/src/providers/databricks_v2.rs @@ -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, ProviderError> { + fn parse_list_endpoints_response( + json: &Value, + ) -> Result<(Vec, Option), ProviderError> { let endpoints = json .get("endpoints") .and_then(|v| v.as_array()) @@ -200,7 +203,7 @@ impl DatabricksV2Provider { ) })?; - let mut models: Vec = endpoints + let models: Vec = 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, 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 = 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); - let json: Value = response.json().await.map_err(|e| { - ProviderError::RequestFailed(format!( - "Failed to parse Databricks AI Gateway endpoints response: {e}" - )) - })?; + if next_page_token.is_none() || next_page_token == page_token { + break; + } + page_token = next_page_token; + } - 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]