diff --git a/src/lib/agents/search/index.ts b/src/lib/agents/search/index.ts
index 859183293..79cf8a4fd 100644
--- a/src/lib/agents/search/index.ts
+++ b/src/lib/agents/search/index.ts
@@ -11,6 +11,7 @@ import { TextBlock } from '@/lib/types';
 
 class SearchAgent {
   async searchAsync(session: SessionManager, input: SearchAgentInput) {
+    try {
     const exists = await db.query.messages.findFirst({
       where: and(
         eq(messages.chatId, input.chatId),
@@ -180,6 +181,24 @@ class SearchAgent {
         ),
       )
       .execute();
+    } catch (err) {
+      console.error('Search agent error:', err);
+
+      session.emit('error', {
+        data: err instanceof Error ? err.message : 'An error occurred during search',
+      });
+
+      await db
+        .update(messages)
+        .set({ status: 'error' })
+        .where(
+          and(
+            eq(messages.chatId, input.chatId),
+            eq(messages.messageId, input.messageId),
+          ),
+        )
+        .execute();
+    }
   }
 }
 
diff --git a/src/lib/hooks/useChat.tsx b/src/lib/hooks/useChat.tsx
index 5ee6d9fbe..9c9da70ba 100644
--- a/src/lib/hooks/useChat.tsx
+++ b/src/lib/hooks/useChat.tsx
@@ -425,6 +425,20 @@ export const ChatProvider = ({ children }: { children: React.ReactNode }) => {
           method: 'POST',
         });
 
+        if (!res.ok) {
+          // Session no longer exists (e.g. server restarted) - mark message as error
+          setLoading(false);
+          isReconnectingRef.current = false;
+          setMessages((prev) =>
+            prev.map((msg) =>
+              msg.messageId === lastMsg.messageId
+                ? { ...msg, status: 'error' as const }
+                : msg,
+            ),
+          );
+          return;
+        }
+
         if (!res.body) throw new Error('No response body');
 
         const reader = res.body?.getReader();