From 422bb77c383ca0ae380427c04698583db5484be2 Mon Sep 17 00:00:00 2001
From: David Katz <dkatz@squareup.com>
Date: Thu, 4 Sep 2025 10:32:12 -0400
Subject: [PATCH] Fix databricks streaming errors  (#4506)

---
 crates/goose-cli/src/session/mod.rs      | 26 ++++++++++++++++++------
 crates/goose/src/providers/databricks.rs | 15 ++++++++------
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs
index 7e9ec62b5d..30de1d7640 100644
--- a/crates/goose-cli/src/session/mod.rs
+++ b/crates/goose-cli/src/session/mod.rs
@@ -1276,12 +1276,26 @@ impl Session {
                             if let Err(e) = self.handle_interrupted_messages(false).await {
                                 eprintln!("Error handling interruption: {}", e);
                             }
-                            output::render_error(
-                                "The error above was an exception we were not able to handle.\n\
-                                These errors are often related to connection or authentication\n\
-                                We've removed the conversation up to the most recent user message\n\
-                                - depending on the error you may be able to continue",
-                            );
+
+                            // Check if it's a ProviderError::ContextLengthExceeded
+                            if e.downcast_ref::<goose::providers::errors::ProviderError>()
+                                .map(|provider_error| matches!(provider_error, goose::providers::errors::ProviderError::ContextLengthExceeded(_)))
+                                .unwrap_or(false) {
+                                output::render_error(
+                                    "Context length exceeded error.\n\
+                                    The conversation is too long for the model's context window.\n\
+                                    Consider using /summarize to condense the conversation history\n\
+                                    or /clear to start fresh.\n\
+                                    We've removed the conversation up to the most recent user message.",
+                                );
+                            } else {
+                                output::render_error(
+                                    "The error above was an exception we were not able to handle.\n\
+                                    These errors are often related to connection or authentication\n\
+                                    We've removed the conversation up to the most recent user message\n\
+                                    - depending on the error you may be able to continue",
+                                );
+                            }
                             break;
                         }
                         None => break,
diff --git a/crates/goose/src/providers/databricks.rs b/crates/goose/src/providers/databricks.rs
index 8a09b522f2..3792b10ab6 100644
--- a/crates/goose/src/providers/databricks.rs
+++ b/crates/goose/src/providers/databricks.rs
@@ -16,7 +16,9 @@ use super::errors::ProviderError;
 use super::formats::databricks::{create_request, response_to_message};
 use super::oauth;
 use super::retry::ProviderRetry;
-use super::utils::{get_model, handle_response_openai_compat, ImageFormat};
+use super::utils::{
+    get_model, handle_response_openai_compat, map_http_error_to_provider_error, ImageFormat,
+};
 use crate::config::ConfigError;
 use crate::conversation::message::Message;
 use crate::impl_provider_default;
@@ -326,11 +328,12 @@ impl Provider for DatabricksProvider {
             .with_retry(|| async {
                 let resp = self.api_client.response_post(&path, &payload).await?;
                 if !resp.status().is_success() {
-                    return Err(ProviderError::RequestFailed(format!(
-                        "HTTP {}: {}",
-                        resp.status(),
-                        resp.text().await.unwrap_or_default()
-                    )));
+                    let status = resp.status();
+                    let error_text = resp.text().await.unwrap_or_default();
+
+                    // Parse as JSON if possible to pass to map_http_error_to_provider_error
+                    let json_payload = serde_json::from_str::<Value>(&error_text).ok();
+                    return Err(map_http_error_to_provider_error(status, json_payload));
                 }
                 Ok(resp)
             })