diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs index c8db3b47ce..314d381fb2 100644 --- a/crates/goose-server/src/routes/reply.rs +++ b/crates/goose-server/src/routes/reply.rs @@ -144,11 +144,7 @@ enum MessageEvent { Ping, } -async fn stream_event( - event: MessageEvent, - tx: &mpsc::Sender, - cancel_token: &CancellationToken, -) { +async fn stream_event(event: MessageEvent, tx: &mpsc::Sender) -> bool { let json = serde_json::to_string(&event).unwrap_or_else(|e| { format!( r#"{{"type":"Error","error":"Failed to serialize event: {}"}}"#, @@ -156,9 +152,13 @@ async fn stream_event( ) }); if tx.send(format!("data: {}\n\n", json)).await.is_err() { - tracing::info!("client hung up"); - cancel_token.cancel(); + 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 + return false; } + true } #[allow(clippy::too_many_lines)] @@ -224,7 +224,6 @@ pub async fn reply( error: format!("Failed to get session agent: {}", e), }, &task_tx, - &task_cancel, ) .await; return; @@ -240,7 +239,6 @@ pub async fn reply( error: format!("Failed to read session: {}", e), }, &task_tx, - &cancel_token, ) .await; return; @@ -272,7 +270,6 @@ pub async fn reply( error: e.to_string(), }, &task_tx, - &cancel_token, ) .await; return; @@ -282,6 +279,7 @@ pub async fn reply( let mut all_messages = messages.clone(); let mut heartbeat_interval = tokio::time::interval(Duration::from_millis(500)); + let mut client_connected = true; loop { tokio::select! { _ = task_cancel.cancelled() => { @@ -289,7 +287,9 @@ pub async fn reply( break; } _ = heartbeat_interval.tick() => { - stream_event(MessageEvent::Ping, &tx, &cancel_token).await; + if client_connected { + client_connected = stream_event(MessageEvent::Ping, &tx).await; + } } response = timeout(Duration::from_millis(500), stream.next()) => { match response { @@ -300,9 +300,9 @@ pub async fn reply( all_messages.push(message.clone()); - // Only send message to client if it's user_visible - if message.is_user_visible() { - stream_event(MessageEvent::Message { message }, &tx, &cancel_token).await; + // 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; } } Ok(Some(Ok(AgentEvent::HistoryReplaced(new_messages)))) => { @@ -312,24 +312,29 @@ pub async fn reply( // The client will see the compaction notification message that was sent before this event } Ok(Some(Ok(AgentEvent::ModelChange { model, mode }))) => { - stream_event(MessageEvent::ModelChange { model, mode }, &tx, &cancel_token).await; + if client_connected { + client_connected = stream_event(MessageEvent::ModelChange { model, mode }, &tx).await; + } } Ok(Some(Ok(AgentEvent::McpNotification((request_id, n))))) => { - stream_event(MessageEvent::Notification{ - request_id: request_id.clone(), - message: n, - }, &tx, &cancel_token).await; + if client_connected { + client_connected = stream_event(MessageEvent::Notification{ + request_id: request_id.clone(), + message: n, + }, &tx).await; + } } Ok(Some(Err(e))) => { tracing::error!("Error processing message: {}", e); - stream_event( - MessageEvent::Error { - error: e.to_string(), - }, - &tx, - &cancel_token, - ).await; + if client_connected { + stream_event( + MessageEvent::Error { + error: e.to_string(), + }, + &tx, + ).await; + } break; } Ok(None) => { @@ -401,7 +406,6 @@ pub async fn reply( reason: "stop".to_string(), }, &task_tx, - &cancel_token, ) .await; })); diff --git a/crates/goose-server/src/routes/session.rs b/crates/goose-server/src/routes/session.rs index ec63363b1b..ea03e92311 100644 --- a/crates/goose-server/src/routes/session.rs +++ b/crates/goose-server/src/routes/session.rs @@ -341,9 +341,18 @@ async fn stream_session( State(_state): State>, Path(session_id): Path, ) -> Result { + tracing::info!("[session.rs] Starting stream for session {}", session_id); + let session_stream = SessionManager::stream_updates(session_id.clone()) .await - .map_err(|_| StatusCode::NOT_FOUND)?; + .map_err(|e| { + tracing::error!( + "[session.rs] Failed to create stream for {}: {}", + session_id, + e + ); + StatusCode::NOT_FOUND + })?; let (tx, rx) = mpsc::channel(100); let stream = ReceiverStream::new(rx); @@ -357,6 +366,13 @@ async fn stream_session( while let Some(result) = session_stream.next().await { match result { Ok(session) => { + tracing::info!( + "[session.rs] Streaming session {} update: in_use={}, message_count={}", + session_id, + session.in_use, + session.message_count + ); + if !stream_session_event( SessionEvent::Session { session: Box::new(session), @@ -365,12 +381,19 @@ async fn stream_session( ) .await { - tracing::info!("Session stream client disconnected"); + tracing::info!( + "[session.rs] Session stream client disconnected for {}", + session_id + ); break; } } Err(e) => { - tracing::error!("Error in session stream: {}", e); + tracing::error!( + "[session.rs] Error in session stream for {}: {}", + session_id, + e + ); let _ = stream_session_event( SessionEvent::Error { error: format!("Failed to fetch session: {}", e), @@ -382,7 +405,7 @@ async fn stream_session( } } } - tracing::info!("Session stream completed for {}", session_id); + tracing::info!("[session.rs] Session stream completed for {}", session_id); }); Ok(SessionSseResponse::new(stream)) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 4ce8055ac6..c46814dbcc 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -64,7 +64,7 @@ use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DEC use crate::agents::subagent_task_config::TaskConfig; use crate::conversation::message::{Message, ToolRequest}; use crate::session::extension_data::{EnabledExtensionsState, ExtensionState}; -use crate::session::SessionManager; +use crate::session::session_manager::{SessionManager, STALE_LOCK_MINUTES}; const DEFAULT_MAX_TURNS: u32 = 1000; @@ -995,9 +995,7 @@ impl Agent { session: Option, cancel_token: Option, ) -> Result>> { - // Check if session is already in use (with 10 minute stale threshold) - const STALE_LOCK_MINUTES: u64 = 10; - + // Check if session is already in use (with stale threshold from STALE_LOCK_MINUTES) if let Some(ref session_config) = session { let is_in_use = SessionManager::is_session_in_use(&session_config.id, STALE_LOCK_MINUTES).await?; diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index 1ea9165a09..d02bcc992e 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -23,6 +23,8 @@ use utoipa::ToSchema; const CURRENT_SCHEMA_VERSION: i32 = 3; +pub const STALE_LOCK_MINUTES: u64 = 10; + static SESSION_STORAGE: OnceCell> = OnceCell::const_new(); #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] @@ -268,23 +270,51 @@ impl SessionManager { let initial_session = Self::get_session(&session_id, true).await?; let stream = async_stream::stream! { - yield Ok(initial_session.clone()); + // Check if session is stale using STALE_LOCK_MINUTES threshold + let is_actually_in_use = Self::is_session_in_use(&session_id, STALE_LOCK_MINUTES).await.unwrap_or(false); - // If session is not in use, close the stream immediately - if !initial_session.in_use { + // If session is marked in_use but is stale, clear the flag + let mut session_to_yield = initial_session.clone(); + if initial_session.in_use && !is_actually_in_use { + // Session is stale, clear the in_use flag + warn!("[stream_updates] Session {} is stale, clearing in_use flag", session_id); + if let Err(e) = Self::mark_in_use(&session_id, false).await { + warn!("Failed to clear stale in_use flag: {}", e); + } + session_to_yield.in_use = false; + } + + yield Ok(session_to_yield.clone()); + + // If session is not in use (or was stale), close the stream immediately + if !session_to_yield.in_use { + info!("[stream_updates] Session {} not in use, closing stream immediately", session_id); return; } + info!("[stream_updates] Session {} is in use, starting polling loop", session_id); + // Session is in use, continue streaming until it's no longer in use let mut interval = tokio::time::interval(Duration::from_secs(1)); - let mut last_message_count: Option = Some(initial_session.message_count); - let mut last_updated_at: Option = Some(initial_session.updated_at.to_rfc3339()); + let mut last_message_count: Option = Some(session_to_yield.message_count); + let mut last_updated_at: Option = Some(session_to_yield.updated_at.to_rfc3339()); loop { interval.tick().await; match Self::get_session(&session_id, true).await { - Ok(session) => { + Ok(mut session) => { + // Check if session is stale + let is_actually_in_use = Self::is_session_in_use(&session_id, STALE_LOCK_MINUTES).await.unwrap_or(false); + + // If session is marked in_use but is stale, clear the flag + if session.in_use && !is_actually_in_use { + if let Err(e) = Self::mark_in_use(&session_id, false).await { + warn!("Failed to clear stale in_use flag: {}", e); + } + session.in_use = false; + } + let updated_at_str = session.updated_at.to_rfc3339(); let has_updates = last_message_count != Some(session.message_count) || last_updated_at.as_ref() != Some(&updated_at_str); diff --git a/ui/desktop/src/components/BaseChat2.tsx b/ui/desktop/src/components/BaseChat2.tsx index 3f4df6d352..7418317ecb 100644 --- a/ui/desktop/src/components/BaseChat2.tsx +++ b/ui/desktop/src/components/BaseChat2.tsx @@ -77,17 +77,29 @@ function BaseChatContent({ // }); // Keep session streaming active for multi-window sync, but track initial load - const [sessionLoadComplete, setSessionLoadComplete] = useState(!resumeSessionId); - const sessionIdToStream = resumeSessionId || chat.sessionId; + const [sessionLoadComplete, setSessionLoadComplete] = useState(false); + const [messages, setMessages] = useState(chat.messages || []); + const sessionId = resumeSessionId || chat.sessionId; + + // Track the last resumeSessionId to detect when it changes + const prevResumeSessionIdRef = useRef(resumeSessionId); + + // Reset state when resumeSessionId changes + useEffect(() => { + if (resumeSessionId && resumeSessionId !== prevResumeSessionIdRef.current) { + // Reset all state for the new session + setSessionLoadComplete(false); + setMessages([]); + lastStreamedMessageCountRef.current = 0; + prevResumeSessionIdRef.current = resumeSessionId; + } + }, [resumeSessionId]); const { session: streamedSession, - isLoading: sessionLoading, error: sessionStreamError, // isConnected: sessionStreamConnected, // maybe we show an indicator somewhere? - } = useSessionStream(sessionIdToStream || undefined); - - const [messages, setMessages] = useState(chat.messages || []); + } = useSessionStream(sessionId || undefined); const isStreamingRef = useRef(false); const lastStreamedMessageCountRef = useRef(0); const [isSessionInUse, setIsSessionInUse] = useState(false); @@ -97,22 +109,21 @@ function BaseChatContent({ if (streamedSession) { const conversation = streamedSession.conversation || []; - // Mark initial load as complete if we're resuming - if (resumeSessionId && !sessionLoadComplete) { - setSessionLoadComplete(true); - } + // Mark initial load as complete if we're resuming and have data + const isInitialLoad = resumeSessionId && !sessionLoadComplete; // Always update in_use state for multi-window coordination setIsSessionInUse(streamedSession.in_use || false); // Only update messages if: - // 1. We're resuming a session (initial load), OR - // 2. We're not actively streaming locally AND message count changed + // 1. We're on initial load (resuming a session), OR + // 2. Message count changed from the stream + // Note: We continue updating even if locally streaming or session is in_use elsewhere + // This ensures all windows see real-time updates const shouldUpdateMessages = - resumeSessionId || - (!isStreamingRef.current && conversation.length !== lastStreamedMessageCountRef.current); + isInitialLoad || conversation.length !== lastStreamedMessageCountRef.current; - if (shouldUpdateMessages && !isStreamingRef.current) { + if (shouldUpdateMessages) { lastStreamedMessageCountRef.current = conversation.length; const loadedChat: ChatType = { @@ -126,6 +137,11 @@ function BaseChatContent({ setChat(loadedChat); setMessages(conversation); + // Mark load complete AFTER we've updated the messages + if (isInitialLoad) { + setSessionLoadComplete(true); + } + // Log for debugging window.electron.logInfo( `Session updated from stream: ${streamedSession.id}, messages: ${conversation.length}, in_use: ${streamedSession.in_use}` @@ -274,6 +290,9 @@ function BaseChatContent({ // TODO(Douwe): get this from the backend const isCompacting = false; + // Determine if we're truly loading: we have a resumeSessionId but haven't loaded messages yet + const isLoadingSession = resumeSessionId && messages.length === 0 && !sessionLoadComplete; + const initialPrompt = messages.length == 0 && recipe?.prompt ? recipe.prompt : ''; return (
@@ -405,24 +424,20 @@ function BaseChatContent({ {/* Fixed loading indicator at bottom left of chat container */} - {((sessionLoading && !sessionLoadComplete) || - (messages.length === 0 && !resumeSessionId) || - isCompacting) && - !sessionStreamError && ( -
- -
- )} + {(isLoadingSession || isCompacting) && !sessionStreamError && ( +
+ +
+ )} {/* Fixed session in progress notice at bottom left */} {isSessionInUse && !isStreamingRef.current && ( diff --git a/ui/desktop/src/components/Pair2.tsx b/ui/desktop/src/components/Pair2.tsx index fbec745714..48b03b259e 100644 --- a/ui/desktop/src/components/Pair2.tsx +++ b/ui/desktop/src/components/Pair2.tsx @@ -1,3 +1,5 @@ +import { useEffect } from 'react'; +import { useSearchParams } from 'react-router-dom'; import { View, ViewOptions } from '../utils/navigationUtils'; import 'react-toastify/dist/ReactToastify.css'; @@ -23,6 +25,19 @@ export default function Pair({ setIsGoosehintsModalOpen, resumeSessionId, }: PairProps & PairRouteState) { + const [_searchParams, setSearchParams] = useSearchParams(); + + // Update URL with sessionId to persist across refreshes + // Only update if resumeSessionId is not already set (to avoid overwriting it on mount) + useEffect(() => { + if (chat.sessionId && !resumeSessionId) { + setSearchParams((prev) => { + prev.set('resumeSessionId', chat.sessionId); + return prev; + }); + } + }, [chat.sessionId, resumeSessionId, setSearchParams]); + return ( { try { + // DEBUG LOGGING + window.electron.logInfo( + JSON.stringify({ + context: 'SessionStreamContext', + event: 'connect_start', + sessionId, + }) + ); + // Clean up existing connection const existingController = abortControllersRef.current.get(sessionId); if (existingController) { + console.log( + `[SessionStreamContext] Aborting existing connection for session ${sessionId}` + ); + window.electron.logInfo( + JSON.stringify({ + context: 'SessionStreamContext', + event: 'aborting_existing_connection', + sessionId, + }) + ); existingController.abort(); + abortControllersRef.current.delete(sessionId); } const secretKey = await window.electron.getSecretKey(); @@ -123,9 +143,12 @@ export function SessionStreamProvider({ children }: { children: ReactNode }) { if (response.ok) { const initialSession: Session = await response.json(); updateSessionState(sessionId, { session: initialSession, isLoading: false }); + } else { + updateSessionState(sessionId, { isLoading: false }); } } catch (err) { console.warn('Failed to fetch initial session data:', err); + updateSessionState(sessionId, { isLoading: false }); } // Create new abort controller for this stream diff --git a/ui/desktop/src/hooks/useSessionStream.ts b/ui/desktop/src/hooks/useSessionStream.ts index 6d86e237a7..28c44f85f4 100644 --- a/ui/desktop/src/hooks/useSessionStream.ts +++ b/ui/desktop/src/hooks/useSessionStream.ts @@ -53,9 +53,31 @@ export function useSessionStream( return; } + // DEBUG LOGGING + if (typeof window !== 'undefined' && window.electron?.logInfo) { + window.electron.logInfo( + JSON.stringify({ + hook: 'useSessionStream', + event: 'registerStream', + sessionId, + enabled, + }) + ); + } + registerStream(sessionId); return () => { + // DEBUG LOGGING + if (typeof window !== 'undefined' && window.electron?.logInfo) { + window.electron.logInfo( + JSON.stringify({ + hook: 'useSessionStream', + event: 'unregisterStream', + sessionId, + }) + ); + } unregisterStream(sessionId); }; }, [sessionId, enabled, registerStream, unregisterStream]);