ai-space/goose
3
0
Fork 0
mirror of https://github.com/aaif-goose/goose.git synced 2026-07-03 14:10:03 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix resume session on refresh with agent in background

Browse Source
This commit is contained in:
Zane Staggs
2025-10-13 19:10:56 -07:00
parent 3368f68bc6
commit fdd5b348ba
8 changed files with 204 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/routes/reply.rs
+31 -27
View File
@@ -144,11 +144,7 @@ enum MessageEvent {
Ping,
}
async fn stream_event(
event: MessageEvent,
tx: &mpsc::Sender<String>,
cancel_token: &CancellationToken,
) {
async fn stream_event(event: MessageEvent, tx: &mpsc::Sender<String>) -> 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;
}));
crates/goose-server/src/routes/session.rs
+27 -4
View File
@@ -341,9 +341,18 @@ async fn stream_session(
State(_state): State<Arc<AppState>>,
Path(session_id): Path<String>,
) -> Result<SessionSseResponse, StatusCode> {
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))
crates/goose/src/agents/agent.rs
+2 -4
View File
@@ -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<SessionConfig>,
cancel_token: Option<CancellationToken>,
) -> Result<BoxStream<'_, Result<AgentEvent>>> {
// 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?;
crates/goose/src/session/session_manager.rs
+36 -6
View File
@@ -23,6 +23,8 @@ use utoipa::ToSchema;
const CURRENT_SCHEMA_VERSION: i32 = 3;
pub const STALE_LOCK_MINUTES: u64 = 10;
static SESSION_STORAGE: OnceCell<Arc<SessionStorage>> = 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<usize> = Some(initial_session.message_count);
let mut last_updated_at: Option<String> = Some(initial_session.updated_at.to_rfc3339());
let mut last_message_count: Option<usize> = Some(session_to_yield.message_count);
let mut last_updated_at: Option<String> = 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);
ui/desktop/src/components/BaseChat2.tsx
+48 -33
View File
@@ -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 (
<div className="h-full flex flex-col min-h-0">
@@ -405,24 +424,20 @@ function BaseChatContent({
</ScrollArea>
{/* Fixed loading indicator at bottom left of chat container */}
{((sessionLoading && !sessionLoadComplete) ||
(messages.length === 0 && !resumeSessionId) ||
isCompacting) &&
!sessionStreamError && (
<div className="absolute bottom-1 left-4 z-20 pointer-events-none">
<LoadingGoose
message={
(sessionLoading && !sessionLoadComplete) ||
(messages.length === 0 && !resumeSessionId)
? 'loading conversation...'
: isCompacting
? 'goose is compacting the conversation...'
: undefined
}
chatState={chatState}
/>
</div>
)}
{(isLoadingSession || isCompacting) && !sessionStreamError && (
<div className="absolute bottom-1 left-4 z-20 pointer-events-none">
<LoadingGoose
message={
isLoadingSession
? 'loading conversation...'
: isCompacting
? 'goose is compacting the conversation...'
: undefined
}
chatState={chatState}
/>
</div>
)}
{/* Fixed session in progress notice at bottom left */}
{isSessionInUse && !isStreamingRef.current && (
ui/desktop/src/components/Pair2.tsx
+15
View File
@@ -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 (
<BaseChat2
chat={chat}
ui/desktop/src/contexts/SessionStreamContext.tsx
+23
View File
@@ -104,10 +104,30 @@ export function SessionStreamProvider({ children }: { children: ReactNode }) {
const connect = useCallback(
async (sessionId: string) => {
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
ui/desktop/src/hooks/useSessionStream.ts
+22
View File
@@ -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]);