diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs index 314d381fb2..85ecc399b5 100644 --- a/crates/goose-server/src/routes/reply.rs +++ b/crates/goose-server/src/routes/reply.rs @@ -85,6 +85,8 @@ pub struct ChatRequest { session_id: String, recipe_name: Option, recipe_version: Option, + #[serde(default)] + background_mode: bool, } pub struct SseResponse { @@ -144,7 +146,12 @@ enum MessageEvent { Ping, } -async fn stream_event(event: MessageEvent, tx: &mpsc::Sender) -> bool { +async fn stream_event( + event: MessageEvent, + tx: &mpsc::Sender, + cancel_token: &CancellationToken, + background_mode: bool, +) -> bool { let json = serde_json::to_string(&event).unwrap_or_else(|e| { format!( r#"{{"type":"Error","error":"Failed to serialize event: {}"}}"#, @@ -152,10 +159,14 @@ async fn stream_event(event: MessageEvent, tx: &mpsc::Sender) -> bool { ) }); if tx.send(format!("data: {}\n\n", json)).await.is_err() { - tracing::info!("client hung up - continuing agent task"); - // Don't cancel the agent task when client disconnects - // Agent should continue running and the session will remain in_use=true - // so that when the client reconnects via the session stream, they can see progress + if background_mode { + tracing::info!( + "client disconnected - continuing agent task in background (ALPHA mode)" + ); + } else { + tracing::info!("client disconnected - cancelling agent task (backwards compatibility)"); + cancel_token.cancel(); + } return false; } true @@ -210,12 +221,19 @@ pub async fn reply( let cancel_token = CancellationToken::new(); let messages = Conversation::new_unvalidated(request.messages); + let background_mode = request.background_mode; + + if background_mode { + let mut tokens = state.session_cancellation_tokens.lock().await; + tokens.insert(session_id.clone(), cancel_token.clone()); + } let task_cancel = cancel_token.clone(); let task_tx = tx.clone(); + let task_state = state.clone(); drop(tokio::spawn(async move { - let agent = match state.get_agent(session_id.clone()).await { + let agent = match task_state.get_agent(session_id.clone()).await { Ok(agent) => agent, Err(e) => { tracing::error!("Failed to get session agent: {}", e); @@ -224,6 +242,8 @@ pub async fn reply( error: format!("Failed to get session agent: {}", e), }, &task_tx, + &task_cancel, + background_mode, ) .await; return; @@ -239,6 +259,8 @@ pub async fn reply( error: format!("Failed to read session: {}", e), }, &task_tx, + &task_cancel, + background_mode, ) .await; return; @@ -270,6 +292,8 @@ pub async fn reply( error: e.to_string(), }, &task_tx, + &task_cancel, + background_mode, ) .await; return; @@ -288,7 +312,7 @@ pub async fn reply( } _ = heartbeat_interval.tick() => { if client_connected { - client_connected = stream_event(MessageEvent::Ping, &tx).await; + client_connected = stream_event(MessageEvent::Ping, &tx, &task_cancel, background_mode).await; } } response = timeout(Duration::from_millis(500), stream.next()) => { @@ -302,7 +326,7 @@ pub async fn reply( // Only send message to client if it's user_visible and client is still connected if message.is_user_visible() && client_connected { - client_connected = stream_event(MessageEvent::Message { message }, &tx).await; + client_connected = stream_event(MessageEvent::Message { message }, &tx, &task_cancel, background_mode).await; } } Ok(Some(Ok(AgentEvent::HistoryReplaced(new_messages)))) => { @@ -313,7 +337,7 @@ pub async fn reply( } Ok(Some(Ok(AgentEvent::ModelChange { model, mode }))) => { if client_connected { - client_connected = stream_event(MessageEvent::ModelChange { model, mode }, &tx).await; + client_connected = stream_event(MessageEvent::ModelChange { model, mode }, &tx, &task_cancel, background_mode).await; } } Ok(Some(Ok(AgentEvent::McpNotification((request_id, n))))) => { @@ -321,7 +345,7 @@ pub async fn reply( client_connected = stream_event(MessageEvent::Notification{ request_id: request_id.clone(), message: n, - }, &tx).await; + }, &tx, &task_cancel, background_mode).await; } } @@ -333,6 +357,8 @@ pub async fn reply( error: e.to_string(), }, &tx, + &task_cancel, + background_mode, ).await; } break; @@ -406,8 +432,15 @@ pub async fn reply( reason: "stop".to_string(), }, &task_tx, + &task_cancel, + background_mode, ) .await; + + if background_mode { + let mut tokens = task_state.session_cancellation_tokens.lock().await; + tokens.remove(&session_id); + } })); Ok(SseResponse::new(stream)) } @@ -459,6 +492,49 @@ pub async fn confirm_permission( Ok(Json(Value::Object(serde_json::Map::new()))) } +#[derive(Debug, Deserialize, Serialize, ToSchema)] +pub struct CancelRequest { + session_id: String, +} + +#[utoipa::path( + post, + path = "/cancel", + request_body = CancelRequest, + responses( + (status = 200, description = "Agent task cancelled successfully"), + (status = 404, description = "Session not found"), + (status = 500, description = "Internal server error") + ) +)] +pub async fn cancel_agent( + State(state): State>, + Json(request): Json, +) -> Result, StatusCode> { + tracing::info!("Cancelling agent task for session: {}", request.session_id); + + // Get the cancellation token for this session and trigger it + let mut tokens = state.session_cancellation_tokens.lock().await; + if let Some(cancel_token) = tokens.remove(&request.session_id) { + cancel_token.cancel(); + tracing::info!("Agent task cancelled for session: {}", request.session_id); + Ok(Json(serde_json::json!({ + "status": "cancelled", + "session_id": request.session_id + }))) + } else { + tracing::warn!( + "No active cancellation token found for session: {}", + request.session_id + ); + // Still return success - the session may have already completed + Ok(Json(serde_json::json!({ + "status": "not_active", + "session_id": request.session_id + }))) + } +} + pub fn routes(state: Arc) -> Router { Router::new() .route( @@ -466,6 +542,7 @@ pub fn routes(state: Arc) -> Router { post(reply).layer(DefaultBodyLimit::max(50 * 1024 * 1024)), ) .route("/confirm", post(confirm_permission)) + .route("/cancel", post(cancel_agent)) .with_state(state) } @@ -496,6 +573,7 @@ mod tests { session_id: "test-session".to_string(), recipe_name: None, recipe_version: None, + background_mode: false, }) .unwrap(), )) diff --git a/crates/goose-server/src/state.rs b/crates/goose-server/src/state.rs index b8b916bf91..dd57675174 100644 --- a/crates/goose-server/src/state.rs +++ b/crates/goose-server/src/state.rs @@ -6,6 +6,8 @@ use std::path::PathBuf; use std::sync::atomic::AtomicUsize; use std::sync::Arc; use tokio::sync::Mutex; +use tokio_util::sync::CancellationToken; + #[derive(Clone)] pub struct AppState { pub(crate) agent_manager: Arc, @@ -13,6 +15,8 @@ pub struct AppState { pub session_counter: Arc, /// Tracks sessions that have already emitted recipe telemetry to prevent double counting. recipe_session_tracker: Arc>>, + /// Tracks cancellation tokens for active sessions + pub(crate) session_cancellation_tokens: Arc>>, } impl AppState { @@ -23,6 +27,7 @@ impl AppState { recipe_file_hash_map: Arc::new(Mutex::new(HashMap::new())), session_counter: Arc::new(AtomicUsize::new(0)), recipe_session_tracker: Arc::new(Mutex::new(HashSet::new())), + session_cancellation_tokens: Arc::new(Mutex::new(HashMap::new())), })) } diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index 602c110849..57783f762a 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -2052,6 +2052,9 @@ "session_id" ], "properties": { + "background_mode": { + "type": "boolean" + }, "messages": { "type": "array", "items": { diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index 8097206a59..2d0d56b805 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -26,6 +26,7 @@ export type AuthorRequest = { }; export type ChatRequest = { + background_mode?: boolean; messages: Array; recipe_name?: string | null; recipe_version?: string | null; diff --git a/ui/desktop/src/hooks/useChatEngine.ts b/ui/desktop/src/hooks/useChatEngine.ts index 6e99a519ad..27d4583d9a 100644 --- a/ui/desktop/src/hooks/useChatEngine.ts +++ b/ui/desktop/src/hooks/useChatEngine.ts @@ -282,6 +282,22 @@ export const useChatEngine = ({ setLastInteractionTime(Date.now()); stopPowerSaveBlocker(); + // In ALPHA mode, when user explicitly stops, we need to cancel the agent task + // This is different from a page refresh where we want the agent to continue + if (process.env.ALPHA === 'true' && chat.sessionId) { + fetch(getApiUrl('/cancel'), { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + session_id: chat.sessionId, + }), + }).catch((error) => { + console.error('Failed to cancel agent task:', error); + }); + } + // Handle stopping the message stream const lastMessage = messages[messages.length - 1]; @@ -361,7 +377,7 @@ export const useChatEngine = ({ setMessages([...messages, responseMessage]); } } - }, [stop, messages, _setInput, setMessages, stopPowerSaveBlocker]); + }, [stop, messages, _setInput, setMessages, stopPowerSaveBlocker, chat.sessionId]); // Since server now handles all filtering, we just use messages directly const filteredMessages = useMemo(() => { diff --git a/ui/desktop/src/hooks/useChatStream.ts b/ui/desktop/src/hooks/useChatStream.ts index 6881b5e277..34f136e073 100644 --- a/ui/desktop/src/hooks/useChatStream.ts +++ b/ui/desktop/src/hooks/useChatStream.ts @@ -84,6 +84,7 @@ export function useChatStream({ body: JSON.stringify({ session_id: sessionIdRef.current, messages: currentMessages, + background_mode: true, }), signal: abortControllerRef.current.signal, }); diff --git a/ui/desktop/src/hooks/useMessageStream.ts b/ui/desktop/src/hooks/useMessageStream.ts index 1b3781839e..99b81d16bd 100644 --- a/ui/desktop/src/hooks/useMessageStream.ts +++ b/ui/desktop/src/hooks/useMessageStream.ts @@ -408,6 +408,7 @@ export function useMessageStream({ }, body: JSON.stringify({ messages: requestMessages, + background_mode: true, ...extraMetadataRef.current.body, }), signal: abortController.signal,