@codewithdan/zingit 0.7.0 → 0.8.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@codewithdan/zingit",
-  "version": "0.7.0",
+  "version": "0.8.0",
   "description": "AI-powered UI annotation tool - point, annotate, and let AI fix it",
   "type": "module",
   "engines": {
@@ -25,6 +25,9 @@
     "dev": "concurrently \"npm run dev:server\" \"npm run dev:client\"",
     "dev:server": "cd server && npm run dev",
     "dev:client": "cd client && npm run dev",
+    "test": "cd client && npm run test",
+    "test:watch": "cd client && npm run test:watch",
+    "test:ui": "cd client && npm run test:ui",
     "release": "npm version minor && npm run build && npm publish --access public"
   },
   "keywords": [

package/server/dist/agents/base.d.ts

@@ -1,11 +1,10 @@
-import type { WebSocket } from 'ws';
-import type { Agent, AgentSession, BatchData, ImageContent } from '../types.js';
+import type { Agent, AgentSession, BatchData, ImageContent, WebSocketRef } from '../types.js';
 export declare abstract class BaseAgent implements Agent {
     abstract name: string;
     abstract model: string;
     abstract start(): Promise<void>;
     abstract stop(): Promise<void>;
-    abstract createSession(ws: WebSocket, projectDir: string): Promise<AgentSession>;
+    abstract createSession(wsRef: WebSocketRef, projectDir: string, resumeSessionId?: string): Promise<AgentSession>;
     /**
      * Format a prompt with image metadata for agents that don't support native multimodal.
      * This adds text descriptions of the images to the prompt.

package/server/dist/agents/claude.d.ts

@@ -1,4 +1,3 @@
-import type { WebSocket } from 'ws';
 import { BaseAgent } from './base.js';
 import type { AgentSession } from '../types.js';
 export declare class ClaudeCodeAgent extends BaseAgent {
@@ -14,5 +13,5 @@ export declare class ClaudeCodeAgent extends BaseAgent {
      * Create a generator that yields the initial user message with optional images
      */
     private createMessageGenerator;
-    createSession(ws: WebSocket, projectDir: string): Promise<AgentSession>;
+    createSession(wsRef: import('../types.js').WebSocketRef, projectDir: string, resumeSessionId?: string): Promise<AgentSession>;
 }

package/server/dist/agents/claude.js

@@ -52,14 +52,18 @@ export class ClaudeCodeAgent extends BaseAgent {
             session_id: '' // SDK will assign the actual session ID
         };
     }
-    async createSession(ws, projectDir) {
+    async createSession(wsRef, projectDir, resumeSessionId) {
         const send = (data) => {
-            if (ws.readyState === ws.OPEN) {
+            const ws = wsRef.current;
+            if (ws && ws.readyState === ws.OPEN) {
                 ws.send(JSON.stringify(data));
             }
         };
         // Track session ID for conversation continuity (stable V1 resume feature)
-        let sessionId;
+        // Start with provided sessionId if resuming previous conversation
+        let sessionId = resumeSessionId;
+        // Track sent content to avoid duplicates (in case SDK sends same message multiple times)
+        const sentContent = new Set();
         return {
             send: async (msg) => {
                 try {
@@ -73,21 +77,31 @@ export class ClaudeCodeAgent extends BaseAgent {
                             permissionMode: 'acceptEdits', // Auto-approve file edits (no interactive terminal)
                             // Resume previous session if we have a session ID (enables follow-up conversations)
                             ...(sessionId && { resume: sessionId }),
-                            systemPrompt: `You are a UI debugging assistant. When given annotations about UI elements,
-you search for the corresponding code using the selectors and HTML context provided,
-then make the requested changes. Be thorough in finding the right files and making precise edits.
+                            systemPrompt: `You are a direct, efficient UI modification assistant.
 
-When screenshots are provided, use them to:
-- Better understand the visual context and styling of the elements
-- Identify the exact appearance that needs to be changed
-- Verify you're targeting the correct element based on its visual representation
+CRITICAL EFFICIENCY RULES:
+1. Use the provided selector and HTML context to quickly locate the target element
+2. Make the requested change immediately - don't explore or explain unless there's ambiguity
+3. For simple changes (text, styles, attributes), be concise - just do it and confirm
+4. Only search/explore if the selector doesn't match or you need to understand complex context
+5. Avoid explaining what annotations are or describing the codebase unnecessarily
 
-IMPORTANT: Format all responses using markdown:
-- Use **bold** for emphasis on important points
-- Use numbered lists for sequential steps (1. 2. 3.)
-- Use bullet points for non-sequential items
-- Use code blocks with \`\`\`language syntax for code examples
-- Use inline \`code\` for file paths, selectors, and technical terms`
+WHEN TO BE BRIEF (90% of cases):
+- Text changes: Find, change, confirm (1-2 sentences)
+- Style changes: Find, modify CSS, confirm
+- Simple DOM changes: Make the change, state what you did
+
+WHEN TO BE THOROUGH (10% of cases):
+- Ambiguous selectors (multiple matches)
+- Complex architectural changes
+- Need to understand component interactions
+
+Response format:
+- **Action taken:** Brief statement of what changed
+- **File:** Path to modified file
+- **Summary:** 1-2 sentence confirmation
+
+Use screenshots to verify you're targeting the right element, but don't over-explain their purpose.`
                         }
                     });
                     // Process streaming response
@@ -97,14 +111,23 @@ IMPORTANT: Format all responses using markdown:
                                 // Capture session ID from init message for follow-up conversations
                                 if ('subtype' in message && message.subtype === 'init') {
                                     sessionId = message.session_id;
+                                    console.log('[Claude Agent] Session initialized:', sessionId);
                                 }
                                 break;
                             case 'assistant':
-                                // Handle assistant message - extract text from BetaMessage content
-                                if (message.message?.content) {
-                                    for (const block of message.message.content) {
-                                        if (block.type === 'text') {
-                                            send({ type: 'delta', content: block.text });
+                                // Extract text content from assistant messages (content is nested in message.content)
+                                if ('message' in message && message.message && 'content' in message.message) {
+                                    const content = message.message.content;
+                                    if (Array.isArray(content)) {
+                                        for (const block of content) {
+                                            if (block.type === 'text' && block.text) {
+                                                // Deduplicate content (SDK may replay conversation history)
+                                                const contentHash = `${block.text.substring(0, 100)}_${block.text.length}`;
+                                                if (!sentContent.has(contentHash)) {
+                                                    sentContent.add(contentHash);
+                                                    send({ type: 'delta', content: block.text });
+                                                }
+                                            }
                                         }
                                     }
                                 }
@@ -114,20 +137,30 @@ IMPORTANT: Format all responses using markdown:
                                 if (message.event?.type === 'content_block_delta') {
                                     const delta = message.event.delta;
                                     if (delta && 'text' in delta) {
+                                        console.log('[Claude Agent] Content block delta, text length:', delta.text.length);
                                         send({ type: 'delta', content: delta.text });
                                     }
                                 }
+                                else if (message.event?.type === 'content_block_stop') {
+                                    console.log('[Claude Agent] Content block stopped');
+                                }
+                                else if (message.event?.type === 'message_stop') {
+                                    console.log('[Claude Agent] Message stopped');
+                                }
                                 break;
                             case 'tool_progress':
                                 // Tool is being executed
+                                console.log('[Claude Agent] Tool executing:', message.tool_name);
                                 send({ type: 'tool_start', tool: message.tool_name });
                                 break;
                             case 'result':
                                 // Query completed
+                                console.log('[Claude Agent] Query completed, sending idle');
                                 send({ type: 'idle' });
                                 break;
                         }
                     }
+                    console.log('[Claude Agent] Response stream ended');
                 }
                 catch (err) {
                     send({ type: 'error', message: err.message });
@@ -135,7 +168,8 @@ IMPORTANT: Format all responses using markdown:
             },
             destroy: async () => {
                 // SDK handles session cleanup automatically
-            }
+            },
+            getSessionId: () => sessionId || null
         };
     }
 }

package/server/dist/agents/codex.d.ts

@@ -1,4 +1,3 @@
-import type { WebSocket } from 'ws';
 import { BaseAgent } from './base.js';
 import type { AgentSession } from '../types.js';
 export declare class CodexAgent extends BaseAgent {
@@ -8,5 +7,5 @@ export declare class CodexAgent extends BaseAgent {
     constructor();
     start(): Promise<void>;
     stop(): Promise<void>;
-    createSession(ws: WebSocket, projectDir: string): Promise<AgentSession>;
+    createSession(wsRef: import('../types.js').WebSocketRef, projectDir: string, resumeSessionId?: string): Promise<AgentSession>;
 }

package/server/dist/agents/codex.js

@@ -24,19 +24,24 @@ export class CodexAgent extends BaseAgent {
         // Codex SDK doesn't require explicit cleanup
         this.codex = null;
     }
-    async createSession(ws, projectDir) {
+    async createSession(wsRef, projectDir, resumeSessionId) {
         if (!this.codex) {
             throw new Error('Codex client not initialized');
         }
         const send = (data) => {
-            if (ws.readyState === ws.OPEN) {
+            const ws = wsRef.current;
+            if (ws && ws.readyState === ws.OPEN) {
                 ws.send(JSON.stringify(data));
             }
         };
-        // Start a Codex thread with the project directory
-        const thread = this.codex.startThread({
+        // Thread options for both new and resumed threads
+        const threadOptions = {
             workingDirectory: projectDir,
-        });
+        };
+        // Resume existing thread if we have a thread ID, otherwise start new thread
+        const thread = resumeSessionId
+            ? this.codex.resumeThread(resumeSessionId, threadOptions)
+            : this.codex.startThread(threadOptions);
         let abortController = null;
         // Track temp files for cleanup on session destroy (prevents race condition)
         const sessionTempFiles = [];
@@ -132,6 +137,7 @@ IMPORTANT: Format all responses using markdown:
                                 break;
                             case 'turn.completed':
                                 // Turn finished
+                                console.log('[Codex Agent] Turn completed, sending idle');
                                 send({ type: 'idle' });
                                 break;
                             case 'turn.failed':
@@ -143,6 +149,7 @@ IMPORTANT: Format all responses using markdown:
                                 break;
                         }
                     }
+                    console.log('[Codex Agent] Event stream ended');
                 }
                 catch (err) {
                     send({ type: 'error', message: err.message });
@@ -171,7 +178,8 @@ IMPORTANT: Format all responses using markdown:
                     }
                     sessionTempFiles.length = 0; // Clear the array
                 }
-            }
+            },
+            getSessionId: () => thread.id
         };
     }
 }

package/server/dist/agents/copilot.d.ts

@@ -1,4 +1,3 @@
-import type { WebSocket } from 'ws';
 import { BaseAgent } from './base.js';
 import type { AgentSession } from '../types.js';
 export declare class CopilotAgent extends BaseAgent {
@@ -8,5 +7,5 @@ export declare class CopilotAgent extends BaseAgent {
     constructor();
     start(): Promise<void>;
     stop(): Promise<void>;
-    createSession(ws: WebSocket, projectDir: string): Promise<AgentSession>;
+    createSession(wsRef: import('../types.js').WebSocketRef, projectDir: string, resumeSessionId?: string): Promise<AgentSession>;
 }

package/server/dist/agents/copilot.js

@@ -29,18 +29,18 @@ export class CopilotAgent extends BaseAgent {
             this.client = null;
         }
     }
-    async createSession(ws, projectDir) {
+    async createSession(wsRef, projectDir, resumeSessionId) {
         if (!this.client) {
             throw new Error('Copilot client not initialized');
         }
         const send = (data) => {
-            if (ws.readyState === ws.OPEN) {
+            const ws = wsRef.current;
+            if (ws && ws.readyState === ws.OPEN) {
                 ws.send(JSON.stringify(data));
             }
         };
-        // Create a Copilot session with streaming enabled
-        // Note: Copilot SDK doesn't support cwd in session config, so we include it in the system message
-        const session = await this.client.createSession({
+        // System message and permission config for both new and resumed sessions
+        const sessionConfig = {
             model: this.model,
             streaming: true,
             systemMessage: {
@@ -75,7 +75,11 @@ IMPORTANT: Format all responses using markdown:
                 }
                 return { kind: 'approved' };
             },
-        });
+        };
+        // Resume existing session if we have a sessionId, otherwise create new session
+        const session = resumeSessionId
+            ? await this.client.resumeSession(resumeSessionId, sessionConfig)
+            : await this.client.createSession(sessionConfig);
         // Track temp files for cleanup on session destroy (prevents race condition)
         const sessionTempFiles = [];
         // Subscribe to streaming events and capture unsubscribe function
@@ -87,17 +91,22 @@ IMPORTANT: Format all responses using markdown:
                     break;
                 case 'assistant.message':
                     // Final message (we already sent deltas, so just log)
+                    console.log('[Copilot Agent] Assistant message complete');
                     break;
                 case 'tool.execution_start':
+                    console.log('[Copilot Agent] Tool executing:', event.data.toolName);
                     send({ type: 'tool_start', tool: event.data.toolName });
                     break;
                 case 'tool.execution_complete':
+                    console.log('[Copilot Agent] Tool complete:', event.data.toolCallId);
                     send({ type: 'tool_end', tool: event.data.toolCallId });
                     break;
                 case 'session.idle':
+                    console.log('[Copilot Agent] Session idle, sending idle message');
                     send({ type: 'idle' });
                     break;
                 case 'session.error':
+                    console.error('[Copilot Agent] Session error:', event.data.message);
                     send({ type: 'error', message: event.data.message });
                     break;
             }
@@ -105,6 +114,7 @@ IMPORTANT: Format all responses using markdown:
         return {
             send: async (msg) => {
                 try {
+                    console.log('[Copilot Agent] send() called, processing request...');
                     // If images are provided, save them as temp files and attach them
                     // Copilot SDK supports file attachments for images
                     const attachments = [];
@@ -134,12 +144,15 @@ IMPORTANT: Format all responses using markdown:
                             });
                         }
                     }
+                    console.log('[Copilot Agent] Calling session.sendAndWait...');
                     await session.sendAndWait({
                         prompt: msg.prompt,
                         attachments: attachments.length > 0 ? attachments : undefined
                     });
+                    console.log('[Copilot Agent] session.sendAndWait completed');
                 }
                 catch (err) {
+                    console.error('[Copilot Agent] Error in send():', err.message);
                     send({ type: 'error', message: err.message });
                 }
                 // Note: Temp files cleaned up on session destroy to avoid race condition
@@ -162,7 +175,8 @@ IMPORTANT: Format all responses using markdown:
                     }
                     sessionTempFiles.length = 0; // Clear the array
                 }
-            }
+            },
+            getSessionId: () => session.sessionId
         };
     }
 }

package/server/dist/handlers/messageHandlers.d.ts

@@ -8,6 +8,8 @@ export interface ConnectionState {
     agent: Agent | null;
     gitManager: GitManager | null;
     currentCheckpointId: string | null;
+    sessionId: string | null;
+    wsRef: import('../types.js').WebSocketRef | null;
 }
 export interface MessageHandlerDeps {
     projectDir: string;
@@ -30,7 +32,7 @@ export declare function handleBatch(ws: WebSocket, state: ConnectionState, msg:
 /**
  * Handle message (follow-up message to agent)
  */
-export declare function handleMessage(ws: WebSocket, state: ConnectionState, msg: WSIncomingMessage): Promise<void>;
+export declare function handleMessage(ws: WebSocket, state: ConnectionState, msg: WSIncomingMessage, deps: MessageHandlerDeps): Promise<void>;
 /**
  * Handle reset message
  */

package/server/dist/handlers/messageHandlers.js

@@ -110,9 +110,11 @@ export async function handleBatch(ws, state, msg, deps) {
     }
     // Use client-specified projectDir, or fall back to server default
     const projectDir = batchData.projectDir || deps.projectDir;
+    console.log('[Batch] ===== Request started =====');
     // Create a checkpoint before AI modifications (if git manager available)
     if (state.gitManager) {
         try {
+            console.log('[Batch] Creating checkpoint...');
             const checkpoint = await state.gitManager.createCheckpoint({
                 annotations: batchData.annotations,
                 pageUrl: batchData.pageUrl,
@@ -120,6 +122,7 @@ export async function handleBatch(ws, state, msg, deps) {
                 agentName: state.agentName,
             });
             state.currentCheckpointId = checkpoint.id;
+            console.log('[Batch] Checkpoint created, sending to client');
             sendMessage(ws, {
                 type: 'checkpoint_created',
                 checkpoint: {
@@ -140,13 +143,60 @@ export async function handleBatch(ws, state, msg, deps) {
             console.warn('Failed to create checkpoint:', err.message);
         }
     }
+    console.log('[Batch] Creating session if needed...');
     if (!state.session) {
-        state.session = await state.agent.createSession(ws, projectDir);
+        // Create or update WebSocket reference for reconnection support
+        if (!state.wsRef) {
+            const { WebSocketRef } = await import('../types.js');
+            state.wsRef = new WebSocketRef(ws);
+        }
+        else {
+            // Update existing reference with new WebSocket (reconnection case)
+            state.wsRef.current = ws;
+        }
+        // Pass the preserved sessionId to resume conversation history
+        state.session = await state.agent.createSession(state.wsRef, projectDir, state.sessionId || undefined);
+        console.log('[Batch] New session created', state.sessionId ? '(resuming previous conversation)' : '(fresh start)');
+        // Store the sessionId if the agent provides it
+        if (state.session.getSessionId) {
+            const newSessionId = state.session.getSessionId();
+            if (newSessionId && newSessionId !== state.sessionId) {
+                state.sessionId = newSessionId;
+                console.log('[Batch] Session ID captured for future resumption');
+            }
+        }
+    }
+    else {
+        console.log('[Batch] Reusing existing session');
+    }
+    // Log user's annotations before formatting
+    console.log('[Batch] Annotation count:', batchData.annotations?.length || 0);
+    if (batchData.annotations && batchData.annotations.length > 0) {
+        batchData.annotations.forEach((ann, idx) => {
+            const notePreview = ann.notes?.substring(0, 200) || '(no notes)';
+            console.log(`[Batch] Annotation ${idx + 1}: ${notePreview}`);
+        });
     }
+    if (batchData.pageUrl) {
+        console.log('[Batch] Page URL:', batchData.pageUrl);
+    }
+    console.log('[Batch] Formatting prompt and extracting images...');
     const prompt = state.agent.formatPrompt(batchData, projectDir);
     const images = state.agent.extractImages(batchData);
+    console.log('[Batch] Image count:', images.length);
+    console.log('[Batch] Sending processing message to client');
     sendMessage(ws, { type: 'processing' });
+    console.log('[Batch] Starting agent session.send()...');
     await state.session.send({ prompt, images: images.length > 0 ? images : undefined });
+    console.log('[Batch] Agent session.send() completed');
+    // Update sessionId after send (it may be assigned during the first message)
+    if (state.session.getSessionId) {
+        const currentSessionId = state.session.getSessionId();
+        if (currentSessionId && currentSessionId !== state.sessionId) {
+            state.sessionId = currentSessionId;
+            console.log('[Batch] Session ID updated after message');
+        }
+    }
     // Finalize checkpoint after processing
     if (state.gitManager && state.currentCheckpointId) {
         try {
@@ -166,30 +216,75 @@ export async function handleBatch(ws, state, msg, deps) {
         }
         state.currentCheckpointId = null;
     }
+    console.log('[Batch] ===== Request completed =====');
 }
 /**
  * Handle message (follow-up message to agent)
  */
-export async function handleMessage(ws, state, msg) {
-    if (state.session && msg.content) {
-        try {
-            console.log(`[ZingIt] Sending follow-up message: "${msg.content.substring(0, 50)}..."`);
-            sendMessage(ws, { type: 'processing' });
-            // Add timeout to detect if SDK hangs
-            const timeoutMs = 120000; // 2 minutes
-            const sendPromise = state.session.send({ prompt: msg.content });
-            const timeoutPromise = new Promise((_, reject) => setTimeout(() => reject(new Error('Agent response timeout')), timeoutMs));
-            await Promise.race([sendPromise, timeoutPromise]);
-            console.log('[ZingIt] Follow-up message sent to agent');
+export async function handleMessage(ws, state, msg, deps) {
+    if (!msg.content) {
+        return;
+    }
+    console.log('[Message] ===== Request started =====');
+    // Create session if it doesn't exist (allows direct messaging without annotations)
+    if (!state.session) {
+        if (!state.agent) {
+            console.warn('[ZingIt] No agent selected for message');
+            sendMessage(ws, { type: 'error', message: 'No agent selected. Please select an agent first.' });
+            return;
         }
-        catch (err) {
-            console.error('[ZingIt] Error sending follow-up message:', err.message);
-            sendMessage(ws, { type: 'error', message: `Failed to send message: ${err.message}` });
+        console.log('[ZingIt] Creating session for direct message', state.sessionId ? '(resuming conversation)' : '(fresh start)');
+        // Create or update WebSocket reference for reconnection support
+        if (!state.wsRef) {
+            const { WebSocketRef } = await import('../types.js');
+            state.wsRef = new WebSocketRef(ws);
+        }
+        else {
+            // Update existing reference with new WebSocket (reconnection case)
+            state.wsRef.current = ws;
+        }
+        // Pass the preserved sessionId to resume conversation history
+        state.session = await state.agent.createSession(state.wsRef, deps.projectDir, state.sessionId || undefined);
+        // Store the sessionId if the agent provides it
+        if (state.session.getSessionId) {
+            const newSessionId = state.session.getSessionId();
+            if (newSessionId && newSessionId !== state.sessionId) {
+                state.sessionId = newSessionId;
+                console.log('[Message] Session ID captured for future resumption');
+            }
         }
     }
-    else if (!state.session) {
-        console.warn('[ZingIt] No active session for follow-up message');
-        sendMessage(ws, { type: 'error', message: 'No active session. Please create annotations first.' });
+    try {
+        const messagePreview = msg.content.length > 500 ? msg.content.substring(0, 500) + '...' : msg.content;
+        console.log('[Message] User message:', messagePreview);
+        if (msg.pageUrl) {
+            console.log('[Message] Page URL:', msg.pageUrl);
+        }
+        sendMessage(ws, { type: 'processing' });
+        // Prepend page URL context if provided
+        let prompt = msg.content;
+        if (msg.pageUrl) {
+            prompt = `[Context: User is currently on page: ${msg.pageUrl}]\n\n${msg.content}`;
+        }
+        // Add timeout to detect if SDK hangs
+        const timeoutMs = 120000; // 2 minutes
+        const sendPromise = state.session.send({ prompt });
+        const timeoutPromise = new Promise((_, reject) => setTimeout(() => reject(new Error('Agent response timeout')), timeoutMs));
+        await Promise.race([sendPromise, timeoutPromise]);
+        console.log('[Message] Agent processing completed');
+        // Update sessionId after send (it may be assigned during the first message)
+        if (state.session && state.session.getSessionId) {
+            const currentSessionId = state.session.getSessionId();
+            if (currentSessionId && currentSessionId !== state.sessionId) {
+                state.sessionId = currentSessionId;
+                console.log('[Message] Session ID updated after message');
+            }
+        }
+        console.log('[Message] ===== Request completed =====');
+    }
+    catch (err) {
+        console.error('[Message] Error sending message:', err.message);
+        sendMessage(ws, { type: 'error', message: `Failed to send message: ${err.message}` });
     }
 }
 /**

package/server/dist/index.js

@@ -94,6 +94,8 @@ async function main() {
         agent: DEFAULT_AGENT ? initializedAgents.get(DEFAULT_AGENT) || null : null,
         gitManager,
         currentCheckpointId: null,
+        sessionId: null, // Preserved across reconnections for conversation continuity
+        wsRef: null, // Will be created when first session is established
     };
     // Track cleanup timer to prevent race conditions
     let sessionCleanupTimer = null;
@@ -107,6 +109,11 @@ async function main() {
             sessionCleanupTimer = null;
             console.log('Cancelled session cleanup - client reconnected');
         }
+        // Update WebSocket reference if session exists (reconnection case)
+        if (state.wsRef) {
+            state.wsRef.current = ws;
+            console.log('Updated WebSocket reference for existing session');
+        }
         // Heartbeat mechanism to detect dead connections
         let isAlive = true;
         ws.on('pong', () => {
@@ -149,7 +156,7 @@ async function main() {
                         await handleBatch(ws, state, msg, deps);
                         break;
                     case 'message':
-                        await handleMessage(ws, state, msg);
+                        await handleMessage(ws, state, msg, deps);
                         break;
                     case 'reset':
                         await handleReset(ws, state);
@@ -180,29 +187,32 @@ async function main() {
             connections.delete(ws);
             // Clean up heartbeat interval
             clearInterval(heartbeatInterval);
-            // Don't destroy session immediately - keep it alive for reconnection (page reload)
-            // Clear any existing cleanup timer to prevent race conditions
-            if (sessionCleanupTimer) {
-                clearTimeout(sessionCleanupTimer);
-                sessionCleanupTimer = null;
-            }
-            // Set new cleanup timer only if no connections remain
-            if (connections.size === 0) {
+            // Don't immediately destroy the session - it may still be processing
+            // Instead, schedule cleanup after a delay to allow for reconnection
+            // If client reconnects, the wsRef will be updated with the new WebSocket
+            if (state.session) {
+                console.log('Scheduling session cleanup (allowing time for reconnection)');
+                // Clear any existing cleanup timer
+                if (sessionCleanupTimer) {
+                    clearTimeout(sessionCleanupTimer);
+                }
+                // Destroy session after 5 seconds if client doesn't reconnect
                 sessionCleanupTimer = setTimeout(async () => {
-                    if (state.session && connections.size === 0) {
-                        console.log('Cleaning up stale session after 5 minutes of inactivity');
+                    if (state.session) {
                         try {
                             await state.session.destroy();
+                            console.log('Session destroyed after reconnection timeout');
                         }
                         catch (err) {
                             console.error('Error destroying session during cleanup:', err.message);
                         }
                         finally {
                             state.session = null;
-                            sessionCleanupTimer = null;
+                            state.wsRef = null;
                         }
                     }
-                }, 300000); // 5 minutes
+                    sessionCleanupTimer = null;
+                }, 5000);
             }
         });
         ws.on('error', (err) => {