fix: page through all Databricks AI Gateway v2 endpoints when listing models

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<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);
 
-        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]