fix(litellm): use context limit from /model/info for custom models (#9303)

Signed-off-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Douwe Osinga <douwe@squareup.com>

crates/goose/src/providers/litellm.rs

@@ -30,6 +30,8 @@ pub struct LiteLLMProvider {
     model: ModelConfig,
     #[serde(skip)]
     name: String,
+    #[serde(skip)]
+    cached_model_info: tokio::sync::OnceCell<Vec<ModelInfo>>,
 }
 
 impl LiteLLMProvider {
@@ -77,10 +79,18 @@ impl LiteLLMProvider {
             base_path,
             model,
             name: LITELLM_PROVIDER_NAME.to_string(),
+            cached_model_info: tokio::sync::OnceCell::new(),
         })
     }
 
-    async fn fetch_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
+    async fn get_or_fetch_models(&self) -> Result<&[ModelInfo], ProviderError> {
+        self.cached_model_info
+            .get_or_try_init(|| self.fetch_models_from_api())
+            .await
+            .map(|v| v.as_slice())
+    }
+
+    async fn fetch_models_from_api(&self) -> Result<Vec<ModelInfo>, ProviderError> {
         let response = self
             .api_client
             .request(None, "model/info")
@@ -184,7 +194,22 @@ impl Provider for LiteLLMProvider {
     }
 
     fn get_model_config(&self) -> ModelConfig {
-        self.model.clone()
+        let mut config = self.model.clone();
+        // The cache is populated lazily by the first stream() call (via
+        // supports_cache_control). On turn 1 this will be None and we fall
+        // back to DEFAULT_CONTEXT_LIMIT, which is fine — the conversation is
+        // too small to trigger compaction. From turn 2 onward the real limit
+        // from /model/info is used.
+        if config.context_limit.is_none() {
+            if let Some(models) = self.cached_model_info.get() {
+                if let Some(info) = models.iter().find(|m| m.name == config.model_name) {
+                    if info.context_limit > 0 {
+                        config.context_limit = Some(info.context_limit);
+                    }
+                }
+            }
+        }
+        config
     }
 
     async fn stream(
@@ -237,7 +262,7 @@ impl Provider for LiteLLMProvider {
     }
 
     async fn supports_cache_control(&self) -> bool {
-        if let Ok(models) = self.fetch_models().await {
+        if let Ok(models) = self.get_or_fetch_models().await {
             if let Some(model_info) = models.iter().find(|m| m.name == self.model.model_name) {
                 return model_info.supports_cache_control.unwrap_or(false);
             }
@@ -247,10 +272,8 @@ impl Provider for LiteLLMProvider {
     }
 
     async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
-        let models = self.fetch_models().await.map_err(|e| {
+        let models = self.get_or_fetch_models().await?;
-            ProviderError::RequestFailed(format!("Failed to fetch models from LiteLLM: {}", e))
+        Ok(models.iter().map(|m| m.name.clone()).collect())
-        })?;
-        Ok(models.into_iter().map(|m| m.name).collect())
     }
 }