From 7d3400ecee118b1ee083d9bb7e5870fe58cdef6f Mon Sep 17 00:00:00 2001
From: Jeff Agapitos <233853744+jeffa-block@users.noreply.github.com>
Date: Sat, 4 Apr 2026 13:53:25 +1100
Subject: [PATCH] feat: add retry and new session buttons to disconnect error
 screen
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Replace the single 'Go home' button on the Failed to Load Session
screen with:
- 'Retry connection' — clears the error and re-triggers session load
- 'New session' — navigates to hub to start fresh
- Session ID display (selectable) for context recovery

The retry works by adding a retryCount state to useChatStream that
is included in the session load effect's dependency array. Incrementing
it re-runs the effect, which calls resumeAgent again.
---
 ui/desktop/src/components/BaseChat.tsx | 32 ++++++++++++++++++--------
 ui/desktop/src/hooks/useChatStream.ts  | 13 +++++++++--
 2 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/ui/desktop/src/components/BaseChat.tsx b/ui/desktop/src/components/BaseChat.tsx
index 6ab1a9c4b065..05e7b5c2c892 100644
--- a/ui/desktop/src/components/BaseChat.tsx
+++ b/ui/desktop/src/components/BaseChat.tsx
@@ -92,6 +92,7 @@ export default function BaseChat({
     submitElicitationResponse,
     stopStreaming,
     sessionLoadError,
+    retrySessionLoad,
     setRecipeUserParams,
     tokenState,
     notifications: toolCallNotifications,
@@ -350,19 +351,30 @@ export default function BaseChat({
           {renderHeader && renderHeader()}
           <div className="flex flex-col flex-1 mb-0.5 min-h-0 relative">
             <div className="flex-1 bg-background-primary rounded-b-2xl flex items-center justify-center">
-              <div className="flex flex-col items-center justify-center p-8">
-                <div className="text-red-700 dark:text-red-300 bg-red-400/50 p-4 rounded-lg mb-4 max-w-md">
+              <div className="flex flex-col items-center justify-center p-8 max-w-md w-full">
+                <div className="text-red-700 dark:text-red-300 bg-red-400/50 p-4 rounded-lg mb-6 w-full">
                   <h3 className="font-semibold mb-2">Failed to Load Session</h3>
                   <p className="text-sm">{sessionLoadError}</p>
                 </div>
-                <button
-                  onClick={() => {
-                    setView('chat');
-                  }}
-                  className="px-4 py-2 text-center cursor-pointer text-text-primary border border-border-primary hover:bg-background-secondary rounded-lg transition-all duration-150"
-                >
-                  Go home
-                </button>
+                <div className="flex gap-3 mb-4">
+                  <button
+                    onClick={retrySessionLoad}
+                    className="px-4 py-2 text-center cursor-pointer text-text-on-primary bg-blue-600 hover:bg-blue-700 rounded-lg transition-all duration-150 font-medium"
+                  >
+                    Retry connection
+                  </button>
+                  <button
+                    onClick={() => {
+                      setView('chat');
+                    }}
+                    className="px-4 py-2 text-center cursor-pointer text-text-primary border border-border-primary hover:bg-background-secondary rounded-lg transition-all duration-150"
+                  >
+                    New session
+                  </button>
+                </div>
+                <p className="text-xs text-text-muted text-center">
+                  Session: <code className="select-all bg-background-secondary px-1 py-0.5 rounded">{sessionId}</code>
+                </p>
               </div>
             </div>
           </div>
diff --git a/ui/desktop/src/hooks/useChatStream.ts b/ui/desktop/src/hooks/useChatStream.ts
index 0306b3ea9041..225fbc177fd8 100644
--- a/ui/desktop/src/hooks/useChatStream.ts
+++ b/ui/desktop/src/hooks/useChatStream.ts
@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useMemo, useReducer, useRef } from 'react';
+import { useCallback, useEffect, useMemo, useReducer, useRef, useState } from 'react';
 import { v7 as uuidv7 } from 'uuid';
 import { AppEvents } from '../constants/events';
 import { ChatState } from '../types/chatState';
@@ -50,6 +50,7 @@ interface UseChatStreamReturn {
   setRecipeUserParams: (values: Record<string, string>) => Promise<void>;
   stopStreaming: () => void;
   sessionLoadError?: string;
+  retrySessionLoad: () => void;
   tokenState: TokenState;
   notifications: Map<string, NotificationEvent[]>;
   onMessageUpdate: (
@@ -349,6 +350,7 @@ export function useChatStream({
   onSessionLoaded,
 }: UseChatStreamProps): UseChatStreamReturn {
   const [state, dispatch] = useReducer(streamReducer, initialState);
+  const [retryCount, setRetryCount] = useState(0);
 
   // Long-lived SSE connection for this session
   const { addListener, setActiveRequestsHandler } = useSessionEvents(sessionId);
@@ -776,7 +778,7 @@ export function useChatStream({
     return () => {
       cancelled = true;
     };
-  }, [sessionId, onSessionLoaded]);
+  }, [sessionId, onSessionLoaded, retryCount]);
 
   const handleSubmit = useCallback(
     async (input: UserInput) => {
@@ -1058,8 +1060,15 @@ export function useChatStream({
     }, new Map<string, NotificationEvent[]>());
   }, [state.notifications]);
 
+  const retrySessionLoad = useCallback(() => {
+    dispatch({ type: 'SET_SESSION_LOAD_ERROR', payload: undefined });
+    dispatch({ type: 'SET_CHAT_STATE', payload: ChatState.LoadingConversation });
+    setRetryCount((c) => c + 1);
+  }, []);
+
   return {
     sessionLoadError: state.sessionLoadError,
+    retrySessionLoad,
     messages: maybe_cached_messages,
     session: maybe_cached_session,
     chatState: state.chatState,