Jazzcort add timeout (#6772)

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

crates/goose-cli/src/cli.rs

@@ -112,6 +112,38 @@ pub struct SessionOptions {
     pub container: Option<String>,
 }
 
+#[derive(Debug, Clone)]
+pub struct StreamableHttpOptions {
+    pub url: String,
+    pub timeout: u64,
+}
+
+fn parse_streamable_http_extension(input: &str) -> Result<StreamableHttpOptions, String> {
+    let mut input_iter = input.split_whitespace();
+    let (mut url, mut timeout) = (String::new(), goose::config::DEFAULT_EXTENSION_TIMEOUT);
+
+    if let Some(url_str) = input_iter.next() {
+        url.push_str(url_str);
+    }
+
+    for kv_pair in input_iter {
+        if !kv_pair.contains('=') {
+            continue;
+        }
+
+        let (key, value) = kv_pair.split_once('=').unwrap();
+
+        // We Can have more keys here for setting other properties
+        if key == "timeout" {
+            if let Ok(seconds) = value.parse::<u64>() {
+                timeout = seconds;
+            }
+        }
+    }
+
+    Ok(StreamableHttpOptions { url, timeout })
+}
+
 /// Extension configuration options shared between Session and Run commands
 #[derive(Args, Debug, Clone, Default)]
 pub struct ExtensionOptions {
@@ -128,10 +160,11 @@ pub struct ExtensionOptions {
         long = "with-streamable-http-extension",
         value_name = "URL",
         help = "Add streamable HTTP extensions (can be specified multiple times)",
-        long_help = "Add streamable HTTP extensions from a URL. Can be specified multiple times. Format: 'url...'",
-        action = clap::ArgAction::Append
+        long_help = "Add streamable HTTP extensions from a URL. Can be specified multiple times. Format: 'url...' or 'url... timeout=100' to set up timeout other than default",
+        action = clap::ArgAction::Append,
+        value_parser = parse_streamable_http_extension
     )]
-    pub streamable_http_extensions: Vec<String>,
+    pub streamable_http_extensions: Vec<StreamableHttpOptions>,
 
     #[arg(
         long = "with-builtin",

crates/goose-cli/src/commands/bench.rs

@@ -1,3 +1,4 @@
+use crate::cli::StreamableHttpOptions;
 use crate::session::build_session;
 use crate::session::SessionBuilderConfig;
 use crate::{logging, CliSession};
@@ -34,13 +35,21 @@ pub async fn agent_generator(
     requirements: ExtensionRequirements,
     session_id: String,
 ) -> BenchAgent {
+    let streamable_http_extensions: Vec<StreamableHttpOptions> = requirements
+        .streamable_http
+        .iter()
+        .map(|s| StreamableHttpOptions {
+            url: s.clone(),
+            timeout: goose::config::DEFAULT_EXTENSION_TIMEOUT,
+        })
+        .collect();
     let base_session = build_session(SessionBuilderConfig {
         session_id: Some(session_id),
         resume: false,
         fork: false,
         no_session: false,
         extensions: requirements.external,
-        streamable_http_extensions: requirements.streamable_http,
+        streamable_http_extensions,
         builtins: requirements.builtin,
         recipe: None,
         additional_system_prompt: None,

crates/goose-cli/src/session/builder.rs

@@ -1,3 +1,5 @@
+use crate::cli::StreamableHttpOptions;
+
 use super::output;
 use super::CliSession;
 use console::style;
@@ -87,7 +89,7 @@ pub struct SessionBuilderConfig {
     /// List of stdio extension commands to add
     pub extensions: Vec<String>,
     /// List of streamable HTTP extension commands to add
-    pub streamable_http_extensions: Vec<String>,
+    pub streamable_http_extensions: Vec<StreamableHttpOptions>,
     /// List of builtin extension commands to add
     pub builtins: Vec<String>,
     /// Recipe for the session
@@ -578,6 +580,7 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
     )
     .await;
 
+
     if let Err(e) = session
         .agent
         .persist_extension_state(&session_id.clone())
@@ -629,7 +632,10 @@ mod tests {
             fork: false,
             no_session: false,
             extensions: vec!["echo test".to_string()],
-            streamable_http_extensions: vec!["http://localhost:8080/mcp".to_string()],
+            streamable_http_extensions: vec![StreamableHttpOptions {
+                url: "http://localhost:8080/mcp".to_string(),
+                timeout: goose::config::DEFAULT_EXTENSION_TIMEOUT,
+            }],
             builtins: vec!["developer".to_string()],
             recipe: None,
             additional_system_prompt: Some("Test prompt".to_string()),

crates/goose-cli/src/session/mod.rs

@@ -9,6 +9,7 @@ mod prompt;
 mod task_execution_display;
 mod thinking;
 
+use crate::cli::StreamableHttpOptions;
 use crate::session::task_execution_display::{
     format_task_execution_notification, TASK_EXECUTION_NOTIFICATION_TYPE,
 };

ui/desktop/src/App.tsx

@@ -418,9 +418,7 @@ export function AppInner() {
       window.electron.reactReady();
     } catch (error) {
       console.error('Error sending reactReady:', error);
-      setFatalError(
-        `React ready notification failed: ${errorMessage(error, 'Unknown error')}`
-      );
+      setFatalError(`React ready notification failed: ${errorMessage(error, 'Unknown error')}`);
     }
   }, []);
 

ui/desktop/src/components/alerts/AlertBox.tsx

@@ -90,9 +90,7 @@ export const AlertBox = ({ alert, className }: AlertBoxProps) => {
       }
     } catch (error) {
       console.error('Error saving threshold:', error);
-      window.alert(
-        `Failed to save threshold: ${errorMessage(error, 'Unknown error')}`
-      );
+      window.alert(`Failed to save threshold: ${errorMessage(error, 'Unknown error')}`);
     } finally {
       setIsSaving(false);
     }

ui/desktop/src/components/schedule/SchedulesView.tsx

@@ -206,9 +206,7 @@ const SchedulesView: React.FC<SchedulesViewProps> = ({ onClose: _onClose }) => {
       setSchedules(fetchedSchedules);
     } catch (error) {
       console.error('Failed to fetch schedules:', error);
-      setApiError(
-        errorMessage(error, 'An unknown error occurred while fetching schedules.')
-      );
+      setApiError(errorMessage(error, 'An unknown error occurred while fetching schedules.'));
     } finally {
       setIsLoading(false);
     }
@@ -375,8 +373,7 @@ const SchedulesView: React.FC<SchedulesViewProps> = ({ onClose: _onClose }) => {
       await fetchSchedules();
     } catch (error) {
       console.error(`Failed to kill running job "${id}":`, error);
-      const errorMsg =
-        errorMessage(error, `Unknown error killing job "${id}".`);
+      const errorMsg = errorMessage(error, `Unknown error killing job "${id}".`);
       setApiError(errorMsg);
       toastError({
         title: 'Kill Job Error',
@@ -413,8 +410,7 @@ const SchedulesView: React.FC<SchedulesViewProps> = ({ onClose: _onClose }) => {
       }
     } catch (error) {
       console.error(`Failed to inspect running job "${id}":`, error);
-      const errorMsg =
-        errorMessage(error, `Unknown error inspecting job "${id}".`);
+      const errorMsg = errorMessage(error, `Unknown error inspecting job "${id}".`);
       setApiError(errorMsg);
       toastError({
         title: 'Inspect Job Error',

ui/desktop/src/components/sessions/SessionHistoryView.tsx

@@ -189,9 +189,7 @@ const SessionHistoryView: React.FC<SessionHistoryViewProps> = ({
       setIsShareModalOpen(true);
     } catch (error) {
       console.error('Error sharing session:', error);
-      toast.error(
-        `Failed to share session: ${errorMessage(error, 'Unknown error')}`
-      );
+      toast.error(`Failed to share session: ${errorMessage(error, 'Unknown error')}`);
     } finally {
       setIsSharing(false);
     }

ui/desktop/src/components/sessions/SessionListView.tsx

@@ -476,7 +476,9 @@ const SessionListView: React.FC<SessionListViewProps> = React.memo(
         );
       } catch (error) {
         console.error('Error deleting session:', error);
-        toast.error(`Failed to delete session "${sessionName}": ${errorMessage(error, 'Unknown error')}`);
+        toast.error(
+          `Failed to delete session "${sessionName}": ${errorMessage(error, 'Unknown error')}`
+        );
       }
       await loadSessions();
     }, [sessionToDelete, loadSessions]);

ui/desktop/src/utils/autoUpdater.ts

@@ -247,12 +247,7 @@ export function registerUpdateIpcHandlers() {
       log.error('Error downloading update:', error);
       const version = githubUpdateInfo.latestVersion || lastUpdateState?.latestVersion || 'unknown';
       const method = isUsingGitHubFallback ? 'github-fallback' : 'electron-updater';
-      trackUpdateDownloadCompleted(
-        false,
-        version,
-        method,
-        errorMessage(error, 'unknown')
-      );
+      trackUpdateDownloadCompleted(false, version, method, errorMessage(error, 'unknown'));
       return {
         success: false,
         error: errorMessage(error, 'Unknown error'),

ui/desktop/src/utils/sessionCache.ts

@@ -58,7 +58,9 @@ export async function loadSession(sessionId: string, forceRefresh = false): Prom
 
       return session;
     } catch (error) {
-      throw new Error(`Error loading session ${sessionId}: ${errorMessage(error, 'Unknown error')}`);
+      throw new Error(
+        `Error loading session ${sessionId}: ${errorMessage(error, 'Unknown error')}`
+      );
     } finally {
       inFlightRequests.delete(sessionId);
     }