Expand databricks agent error handling to fallback to token limit errors

crates/goose-llm/src/providers/databricks.rs

@@ -119,6 +119,38 @@ impl DatabricksProvider {
                 ProviderError::RequestFailed("Response body is not valid JSON".to_string())
             }),
             StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
+                let payload_str = serde_json::to_string(&payload)
+                    .unwrap_or_default()
+                    .to_lowercase();
+
+                // Check for potential context length issues in 403 errors too
+                let check_phrases = [
+                    "too long",
+                    "context length",
+                    "context_length_exceeded",
+                    "reduce the length",
+                    "token count",
+                    "exceeds",
+                    "exceed context limit",
+                    "input length",
+                    "max_tokens",
+                    "decrease input length",
+                    "context limit",
+                ];
+                if check_phrases.iter().any(|c| payload_str.contains(c)) {
+                    return Err(ProviderError::ContextLengthExceeded(payload_str));
+                }
+
+                // Also check for generic permission denied that might be context length
+                if payload_str.contains("external_model_provider")
+                    && (payload_str.contains("forbidden") || payload_str.contains("permission"))
+                {
+                    return Err(ProviderError::ContextLengthExceeded(format!(
+                        "Potential context length exceeded (Databricks permission error): {}",
+                        payload_str
+                    )));
+                }
+
                 Err(ProviderError::Authentication(format!(
                     "Authentication failed. Please ensure your API keys are valid and have the required permissions. \
                     Status: {}. Response: {:?}",

crates/goose-server/src/routes/session.rs

@@ -187,14 +187,20 @@ async fn get_session_insights(
 
         // Track tokens - only add positive values to prevent negative totals
         if let Some(tokens) = session.metadata.accumulated_total_tokens {
-            if tokens > 0 {
-                total_tokens += tokens as i64;
-            } else if tokens < 0 {
-                // Log negative token values for debugging
-                info!(
-                    "Warning: Session {} has negative accumulated_total_tokens: {}",
-                    session.id, tokens
-                );
+            match tokens.cmp(&0) {
+                std::cmp::Ordering::Greater => {
+                    total_tokens += tokens as i64;
+                }
+                std::cmp::Ordering::Less => {
+                    // Log negative token values for debugging
+                    info!(
+                        "Warning: Session {} has negative accumulated_total_tokens: {}",
+                        session.id, tokens
+                    );
+                }
+                std::cmp::Ordering::Equal => {
+                    // Do nothing for zero tokens
+                }
             }
         }
 

crates/goose/src/providers/databricks.rs

@@ -349,7 +349,9 @@ impl DatabricksProvider {
                     // We try to extract the error message from the payload and check for phrases that indicate context length exceeded
                     let bytes = response.bytes().await?;
                     let payload_str = String::from_utf8_lossy(&bytes).to_lowercase();
-                    let check_phrases = [
+
+                    // First check for explicit context length phrases
+                    let explicit_context_phrases = [
                         "too long",
                         "context length",
                         "context_length_exceeded",
@@ -362,10 +364,27 @@ impl DatabricksProvider {
                         "decrease input length",
                         "context limit",
                     ];
-                    if check_phrases.iter().any(|c| payload_str.contains(c)) {
+                    if explicit_context_phrases
+                        .iter()
+                        .any(|c| payload_str.contains(c))
+                    {
                         return Err(ProviderError::ContextLengthExceeded(payload_str));
                     }
 
+                    // For Databricks, also check for generic "Bad Request" errors that might be context length issues
+                    // This is a heuristic since Databricks sometimes returns generic errors for context length issues
+                    if payload_str.contains("external_model_provider")
+                        && payload_str.contains("bad_request")
+                    {
+                        // This might be a context length exceeded error disguised as a generic bad request
+                        // Let's treat it as such to trigger auto-summarization
+                        tracing::warn!("Databricks returned generic Bad Request error, treating as potential context length exceeded: {}", payload_str);
+                        return Err(ProviderError::ContextLengthExceeded(format!(
+                            "Potential context length exceeded (Databricks generic error): {}",
+                            payload_str
+                        )));
+                    }
+
                     let mut error_msg = "Unknown error".to_string();
                     if let Ok(response_json) = serde_json::from_slice::<Value>(&bytes) {
                         // try to convert message to string, if that fails use external_model_message

ui/desktop/src/api/types.gen.ts

@@ -273,6 +273,10 @@ export type ModelInfo = {
      * Cost per token for output (optional)
      */
     output_token_cost?: number | null;
+    /**
+     * Whether this model supports cache control
+     */
+    supports_cache_control?: boolean | null;
 };
 
 export type PermissionConfirmationRequest = {

ui/desktop/src/components/ProgressiveMessageList.tsx

@@ -77,9 +77,19 @@ export default function ProgressiveMessageList({
     getContextHandlerType = contextManager.getContextHandlerType;
   } catch (error) {
     // Context manager not available (e.g., in session history view)
-    // This is fine, we'll just skip context handler functionality
-    hasContextHandlerContent = undefined;
-    getContextHandlerType = undefined;
+    // Provide fallback functions to still handle context messages
+    hasContextHandlerContent = (message: Message) => {
+      return message.content.some(
+        (content) =>
+          content.type === 'contextLengthExceeded' || content.type === 'summarizationRequested'
+      );
+    };
+    getContextHandlerType = (message: Message) => {
+      const hasContextLengthExceeded = message.content.some(
+        (content) => content.type === 'contextLengthExceeded'
+      );
+      return hasContextLengthExceeded ? 'contextLengthExceeded' : 'summarizationRequested';
+    };
   }
 
   // Simple progressive loading - start immediately when component mounts if needed