Websocket followup (#49)

* websocket followup added

* websocket followup added

Author: shivammittal274

packages/agent/src/agent/BaseAgent.ts

@@ -132,6 +132,19 @@ export abstract class BaseAgent {
    */
   abstract destroy(): Promise<void>;
 
+  /**
+   * Abort current execution
+   * Triggers the abort signal to stop the current task
+   * Must be implemented by concrete agent classes
+   */
+  abstract abort(): void;
+
+  /**
+   * Check if agent is currently executing
+   * Must be implemented by concrete agent classes
+   */
+  abstract isExecuting(): boolean;
+
   /**
    * Get current agent metadata
    */

packages/agent/src/agent/ClaudeSDKAgent.ts

@@ -191,9 +191,9 @@ export class ClaudeSDKAgent extends BaseAgent {
           logger.info(
             '⚠️  Agent execution aborted (caught during iterator wait)',
           );
-          // Cleanup iterator
+          // Cleanup iterator (fire-and-forget to avoid blocking)
           if (iterator.return) {
-            await iterator.return(undefined).catch(() => {});
+            iterator.return(undefined).catch(() => {});
           }
           return;
         }
@@ -375,6 +375,26 @@ export class ClaudeSDKAgent extends BaseAgent {
     }
   }
 
+  /**
+   * Abort current execution
+   * Triggers abort signal to stop the current task gracefully
+   */
+  abort(): void {
+    if (this.abortController) {
+      logger.info('🛑 Aborting ClaudeSDKAgent execution');
+      this.abortController.abort();
+    } else {
+      logger.warn('⚠️  Cancel not fully supported - no active execution');
+    }
+  }
+
+  /**
+   * Check if agent is currently executing
+   */
+  isExecuting(): boolean {
+    return this.metadata.state === 'executing' && this.abortController !== null;
+  }
+
   /**
    * Cleanup agent resources
    *

packages/agent/src/agent/CodexSDKAgent.ts

@@ -6,7 +6,7 @@
 import {accessSync, constants as fsConstants} from 'node:fs';
 import {dirname, join} from 'node:path';
 
-import {Codex, type McpServerConfig} from '@browseros/codex-sdk-ts';
+import {Codex, Thread, type McpServerConfig} from '@browseros/codex-sdk-ts';
 import {logger} from '@browseros/common';
 import type {ControllerBridge} from '@browseros/controller-server';
 import {allControllerTools} from '@browseros/tools/controller-based';
@@ -66,6 +66,7 @@ export class CodexSDKAgent extends BaseAgent {
   private codex: Codex | null = null;
   private codexExecutablePath: string | null = null;
   private codexConfigPath: string | null = null;
+  private currentThread: Thread | null = null;
 
   constructor(config: AgentConfig, _controllerBridge: ControllerBridge) {
     const mcpServerConfig = buildMcpServerConfig(config);
@@ -338,7 +339,17 @@ export class CodexSDKAgent extends BaseAgent {
         });
       }
 
-      const thread = this.codex.startThread(threadOptions);
+      // Reuse existing thread for follow-up messages, or create new one
+      // CRITICAL: Check both existence AND thread ID (ID is null if cancelled before thread.started event)
+      if (!this.currentThread || !this.currentThread.id) {
+        this.currentThread = this.codex.startThread(threadOptions);
+        logger.info('🆕 Created new thread for session');
+      } else {
+        logger.info('♻️  Reusing existing thread for follow-up message', {
+          threadId: this.currentThread.id,
+        });
+      }
+      const thread = this.currentThread;
 
       // Get streaming events from thread
       const messages: Array<{type: 'text'; text: string}> = [];
@@ -368,6 +379,9 @@ export class CodexSDKAgent extends BaseAgent {
             logger.info(
               '⚠️  Agent execution aborted by client (breaking loop)',
             );
+            // Clear thread - next message will create fresh thread
+            this.currentThread = null;
+            logger.debug('🔄 Cleared thread reference due to abort');
             break;
           }
 
@@ -454,9 +468,14 @@ export class CodexSDKAgent extends BaseAgent {
         }
       } finally {
         // CRITICAL: Close iterator to trigger SIGKILL in forked SDK's finally block
+        // Fire-and-forget to avoid blocking markIdle() - subprocess cleanup can happen async
         if (iterator.return) {
           logger.debug('🔒 Closing iterator to terminate Codex subprocess');
-          await iterator.return(undefined);
+          iterator.return(undefined).catch((error) => {
+            logger.warn('⚠️  Iterator cleanup error (non-fatal)', {
+              error: error instanceof Error ? error.message : String(error),
+            });
+          });
         }
       }
 
@@ -469,6 +488,10 @@ export class CodexSDKAgent extends BaseAgent {
         duration: Date.now() - this.executionStartTime,
       });
     } catch (error) {
+      // Clear thread on error - next call will create fresh thread
+      this.currentThread = null;
+      logger.debug('🔄 Cleared thread reference due to error');
+
       // Mark execution error
       this.errorExecution(
         error instanceof Error ? error : new Error(String(error)),
@@ -486,6 +509,24 @@ export class CodexSDKAgent extends BaseAgent {
     }
   }
 
+  /**
+   * Abort current execution
+   * Triggers abort signal to stop the current task gracefully
+   */
+  abort(): void {
+    if (this.abortController) {
+      logger.info('🛑 Aborting CodexSDKAgent execution');
+      this.abortController.abort();
+    }
+  }
+
+  /**
+   * Check if agent is currently executing
+   */
+  isExecuting(): boolean {
+    return this.metadata.state === 'executing' && this.abortController !== null;
+  }
+
   /**
    * Cleanup agent resources
    *
@@ -500,6 +541,9 @@ export class CodexSDKAgent extends BaseAgent {
 
     this.markDestroyed();
 
+    // Clear thread reference
+    this.currentThread = null;
+
     // Trigger abort controller for cleanup
     if (this.abortController) {
       this.abortController.abort();

packages/agent/src/session/SessionManager.ts

@@ -260,6 +260,45 @@ export class SessionManager {
     logger.debug('Session marked as idle', {sessionId});
   }
 
+  /**
+   * Cancel current execution for a session
+   * Triggers abort on the agent if it's executing
+   * CRITICAL: Does NOT mark session as idle - let processMessage() handle that
+   *
+   * @param sessionId - Session ID
+   * @returns true if cancel was triggered, false if not executing or agent not found
+   */
+  cancelExecution(sessionId: string): boolean {
+    const agent = this.agents.get(sessionId);
+    if (!agent) {
+      logger.warn('⚠️  Cancel requested but no agent found', {sessionId});
+      return false;
+    }
+
+    // Defensive: check abort support
+    if (typeof agent.abort !== 'function') {
+      logger.warn('⚠️  Agent does not support cancel', {
+        sessionId,
+        agentType: agent.getMetadata().type,
+      });
+      return false;
+    }
+
+    if (!agent.isExecuting()) {
+      logger.debug('⚠️  Cancel requested but agent not executing', {sessionId});
+      return false;
+    }
+
+    logger.info('🛑 Cancelling execution', {sessionId});
+    agent.abort();
+
+    // CRITICAL: Do NOT mark idle here!
+    // Let the original processMessage() call mark idle when it completes
+    // Otherwise we get race condition: new messages can start while execute() is still in finally block
+
+    return true;
+  }
+
   /**
    * Delete a session and its agent
    *

packages/agent/src/websocket/protocol.ts

@@ -17,13 +17,29 @@ import {z} from 'zod';
 // ============================================================================
 
 /**
- * Message sent from client to server
+ * Regular message from client
  */
-export const ClientMessageSchema = z.object({
+export const ClientRegularMessageSchema = z.object({
   type: z.literal('message'),
   content: z.string().min(1, 'Message content cannot be empty'),
 });
 
+/**
+ * Cancel message from client
+ */
+export const ClientCancelMessageSchema = z.object({
+  type: z.literal('cancel'),
+  sessionId: z.string().optional(),
+});
+
+/**
+ * Discriminated union of all client message types
+ */
+export const ClientMessageSchema = z.discriminatedUnion('type', [
+  ClientRegularMessageSchema,
+  ClientCancelMessageSchema,
+]);
+
 export type ClientMessage = z.infer<typeof ClientMessageSchema>;
 
 // ============================================================================
@@ -63,6 +79,17 @@ export const AgentEventSchema = z.object({
 
 export type AgentEvent = z.infer<typeof AgentEventSchema>;
 
+/**
+ * Cancelled event (acknowledgment of cancel request)
+ */
+export const CancelledEventSchema = z.object({
+  type: z.literal('cancelled'),
+  sessionId: z.string(),
+  message: z.string().optional(),
+});
+
+export type CancelledEvent = z.infer<typeof CancelledEventSchema>;
+
 /**
  * Error event
  */
@@ -80,6 +107,7 @@ export type ErrorEvent = z.infer<typeof ErrorEventSchema>;
 export const ServerEventSchema = z.union([
   ConnectionEventSchema,
   AgentEventSchema,
+  CancelledEventSchema,
   ErrorEventSchema,
 ]);
 

packages/agent/src/websocket/server.ts

@@ -94,23 +94,23 @@ export function createServer(
   // Track WebSocket connections (needed to close idle sessions)
   const wsConnections = new Map<string, ServerWebSocket<WebSocketData>>();
 
-  // Cleanup idle sessions callback (now async)
-  const cleanupIdle = async () => {
-    const idleSessionIds = sessionManager.findIdleSessions();
-
-    for (const sessionId of idleSessionIds) {
-      const ws = wsConnections.get(sessionId);
-      if (ws) {
-        logger.info('🧹 Closing idle session', {sessionId});
-        ws.close(1001, 'Idle timeout');
-        wsConnections.delete(sessionId);
-      }
-      await sessionManager.deleteSession(sessionId);
-    }
-  };
-
-  // Run cleanup check with the timer
-  setInterval(cleanupIdle, 60000);
+  // Cleanup idle sessions callback (now async) -> commenting out for now as we should let BrowserOS agent handle this
+  // const cleanupIdle = async () => {
+  //   const idleSessionIds = sessionManager.findIdleSessions();
+
+  //   for (const sessionId of idleSessionIds) {
+  //     const ws = wsConnections.get(sessionId);
+  //     if (ws) {
+  //       logger.info('🧹 Closing idle session', {sessionId});
+  //       ws.close(1001, 'Idle timeout');
+  //       wsConnections.delete(sessionId);
+  //     }
+  //     await sessionManager.deleteSession(sessionId);
+  //   }
+  // };
+
+  // // Run cleanup check with the timer
+  // setInterval(cleanupIdle, 60000);
 
   const server = Bun.serve<WebSocketData>({
     port: config.port,
@@ -244,15 +244,6 @@ export function createServer(
             return;
           }
 
-          // Try to mark session as processing (reject if already processing)
-          if (!sessionManager.markProcessing(sessionId)) {
-            sendError(
-              ws,
-              'Session is already processing a message. Please wait.',
-            );
-            return;
-          }
-
           // Parse message
           const messageStr =
             typeof message === 'string'
@@ -266,11 +257,40 @@ export function createServer(
           const clientMessage = tryParseClientMessage(parsedData);
 
           if (!clientMessage) {
-            sessionManager.markIdle(sessionId); // Mark idle before returning
             sendError(ws, 'Invalid message format');
             return;
           }
 
+          // Handle cancel message
+          if (clientMessage.type === 'cancel') {
+            logger.info('🛑 Cancel request received', {sessionId});
+
+            const success = sessionManager.cancelExecution(sessionId);
+
+            // Send cancelled acknowledgment
+            const cancelledEvent = {
+              type: 'cancelled',
+              sessionId,
+              message: success
+                ? 'Execution cancelled'
+                : 'No active execution to cancel',
+            };
+            ws.send(JSON.stringify(cancelledEvent));
+
+            logger.info(success ? '✅ Cancel successful' : '⚠️  Nothing to cancel', {sessionId});
+            return;
+          }
+
+          // Handle regular message
+          // Try to mark session as processing (reject if already processing)
+          if (!sessionManager.markProcessing(sessionId)) {
+            sendError(
+              ws,
+              'Session is already processing a message. Please wait.',
+            );
+            return;
+          }
+
           // Update stats
           stats.messagesProcessed++;
 
@@ -408,8 +428,10 @@ async function processMessage(
       try {
         result = await Promise.race([iterator.next(), timeoutPromise]);
       } catch (timeoutError) {
-        // Cleanup iterator
-        if (iterator.return) await iterator.return(undefined);
+        // Cleanup iterator (fire-and-forget - session will be deleted anyway)
+        if (iterator.return) {
+          iterator.return(undefined).catch(() => {});
+        }
         throw new Error(
           `Event gap timeout: No activity for ${config.eventGapTimeoutMs / 1000}s`,
         );