fix(vscode): prevent race conditions in prompt cancellation and streaming

packages/cli/src/acp-integration/session/Session.ts

@@ -90,6 +90,14 @@ const debugLogger = createDebugLogger('SESSION');
  */
 export class Session implements SessionContext {
   private pendingPrompt: AbortController | null = null;
+  /**
+   * Tracks the completion of the current prompt so that the next prompt
+   * can await it.  This prevents a new prompt from reading chat history
+   * before the previous prompt's tool results have been added —
+   * a race condition that causes malformed history on Windows where
+   * process termination is slow.
+   */
+  private pendingPromptCompletion: Promise<void> | null = null;
   private turn: number = 0;
 
   // Modular components
@@ -143,10 +151,43 @@ export class Session implements SessionContext {
   }
 
   async prompt(params: PromptRequest): Promise<PromptResponse> {
+    // Install this prompt's AbortController before awaiting the previous
+    // prompt, so that a session/cancel during the wait targets us.
     this.pendingPrompt?.abort();
     const pendingSend = new AbortController();
     this.pendingPrompt = pendingSend;
 
+    // Wait for the previous prompt to finish so chat history is consistent.
+    if (this.pendingPromptCompletion) {
+      try {
+        await this.pendingPromptCompletion;
+      } catch {
+        // Expected: previous prompt was cancelled or errored
+      }
+    }
+
+    // Cancelled while waiting for the previous prompt to finish.
+    if (pendingSend.signal.aborted) {
+      return { stopReason: 'cancelled' };
+    }
+
+    // Track this prompt's completion for the next prompt to await
+    let resolveCompletion!: () => void;
+    this.pendingPromptCompletion = new Promise<void>((resolve) => {
+      resolveCompletion = resolve;
+    });
+
+    try {
+      return await this.#executePrompt(params, pendingSend);
+    } finally {
+      resolveCompletion();
+    }
+  }
+
+  async #executePrompt(
+    params: PromptRequest,
+    pendingSend: AbortController,
+  ): Promise<PromptResponse> {
     // Increment turn counter for each user prompt
     this.turn += 1;
 

packages/vscode-ide-companion/src/webview/App.tsx

@@ -307,22 +307,24 @@ export const App: React.FC = () => {
   // Emit a cancel to the extension and immediately reflect interruption locally.
   const handleCancel = useCallback(() => {
     if (messageHandling.isStreaming || messageHandling.isWaitingForResponse) {
-      // Proactively end local states and add an 'Interrupted' line
-      try {
-        messageHandling.endStreaming?.();
-      } catch {
-        /* no-op */
-      }
-      try {
-        messageHandling.clearWaitingForResponse?.();
-      } catch {
-        /* no-op */
+      // End streaming state and add an 'Interrupted' line.
+      // IMPORTANT: Do NOT clear isWaitingForResponse here — let the
+      // extension's streamEnd message clear it after the cancel is
+      // properly processed on the backend.  This keeps the submit
+      // guard active and prevents any cached input from being
+      // auto-submitted during the cancel → confirmed window.
+      if (messageHandling.isStreaming) {
+        try {
+          messageHandling.endStreaming?.();
+        } catch {
+          /* no-op */
+        }
+        messageHandling.addMessage({
+          role: 'assistant',
+          content: 'Interrupted',
+          timestamp: Date.now(),
+        });
       }
-      messageHandling.addMessage({
-        role: 'assistant',
-        content: 'Interrupted',
-        timestamp: Date.now(),
-      });
     }
     // Notify extension/agent to cancel server-side work
     vscode.postMessage({

packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.ts

@@ -159,17 +159,50 @@ export class SessionMessageHandler extends BaseMessageHandler {
     this.currentStreamContent = '';
   }
 
+  /**
+   * Monotonically increasing request counter used to tag streamStart/streamEnd
+   * so the WebView can detect and discard stale events from previous requests.
+   */
+  private requestCounter = 0;
+  private currentRequestId: string | null = null;
+  private streamEndSent = false;
+
   /**
    * Notify the webview that streaming has finished.
+   * Includes the `requestId` so the webview can ignore stale events.
+   * Guarded by `streamEndSent` to prevent duplicate streamEnd for the
+   * same request (e.g. cancel handler + error handler both sending one).
+   *
+   * @param reason  Optional reason string (e.g. 'user_cancelled').
+   * @param forRequestId  When provided, the call is scoped to a specific
+   *   request invocation.  If a newer request has since overwritten
+   *   `this.currentRequestId`, the call is silently dropped — this
+   *   prevents a stale `handleSendMessage` invocation (resumed after
+   *   cancellation) from emitting a streamEnd tagged as the newer request.
    */
-  private sendStreamEnd(reason?: string): void {
-    const data: { timestamp: number; reason?: string } = {
+  private sendStreamEnd(reason?: string, forRequestId?: string): void {
+    if (this.streamEndSent) {
+      return;
+    }
+    // If the caller captured a request ID, only proceed when it still
+    // matches the active request.  A mismatch means a newer request has
+    // taken over the shared state; emitting now would incorrectly tag
+    // the event with the newer request's ID.
+    if (forRequestId && this.currentRequestId !== forRequestId) {
+      return;
+    }
+    this.streamEndSent = true;
+
+    const data: { timestamp: number; reason?: string; requestId?: string } = {
       timestamp: Date.now(),
     };
 
     if (reason) {
       data.reason = reason;
     }
+    if (this.currentRequestId) {
+      data.requestId = this.currentRequestId;
+    }
 
     this.sendToWebView({
       type: 'streamEnd',
@@ -388,12 +421,28 @@ export class SessionMessageHandler extends BaseMessageHandler {
     }
 
     // Send to agent
+    //
+    // Generate a unique requestId so the webview can correlate
+    // streamStart/streamEnd and discard stale events.
+    this.requestCounter += 1;
+    this.currentRequestId = `req-${this.requestCounter}-${Date.now()}`;
+    this.streamEndSent = false;
+
+    // Capture locally so that if a newer handleSendMessage() overwrites
+    // the shared fields while we are awaiting, our sendStreamEnd calls
+    // will detect the mismatch and silently no-op instead of emitting
+    // a streamEnd tagged with the newer request's ID.
+    const myRequestId = this.currentRequestId;
+
     try {
       this.resetStreamContent();
 
       this.sendToWebView({
         type: 'streamStart',
-        data: { timestamp: Date.now() },
+        data: {
+          timestamp: Date.now(),
+          requestId: myRequestId,
+        },
       });
 
       await this.agentManager.sendMessage(formattedText);
@@ -411,7 +460,7 @@ export class SessionMessageHandler extends BaseMessageHandler {
         );
       }
 
-      this.sendStreamEnd();
+      this.sendStreamEnd(undefined, myRequestId);
     } catch (error) {
       console.error('[SessionMessageHandler] Error sending message:', error);
 
@@ -433,7 +482,7 @@ export class SessionMessageHandler extends BaseMessageHandler {
       if (isAbortLike) {
         // Do not show VS Code error popup for intentional cancellations.
         // Ensure the webview knows the stream ended due to user action.
-        this.sendStreamEnd('user_cancelled');
+        this.sendStreamEnd('user_cancelled', myRequestId);
         return;
       }
       // Check for session not found error and handle it appropriately
@@ -451,7 +500,7 @@ export class SessionMessageHandler extends BaseMessageHandler {
           type: 'sessionExpired',
           data: { message: 'Session expired. Please login again.' },
         });
-        this.sendStreamEnd('session_expired');
+        this.sendStreamEnd('session_expired', myRequestId);
       } else {
         const isTimeoutError =
           lower.includes('timeout') || lower.includes('timed out');
@@ -474,15 +523,15 @@ export class SessionMessageHandler extends BaseMessageHandler {
             type: 'message',
             data: timeoutMessage,
           });
-          this.sendStreamEnd('timeout');
+          this.sendStreamEnd('timeout', myRequestId);
         } else {
           // Handling of Non-Timeout Errors
           vscode.window.showErrorMessage(`Error sending message: ${errorMsg}`);
           this.sendToWebView({
             type: 'error',
             data: { message: errorMsg },
           });
-          this.sendStreamEnd('error');
+          this.sendStreamEnd('error', myRequestId);
         }
       }
     }
@@ -790,21 +839,15 @@ export class SessionMessageHandler extends BaseMessageHandler {
       // Cancel the current streaming operation in the agent manager
       await this.agentManager.cancelCurrentPrompt();
 
-      // Send streamEnd message to WebView to update UI
-      this.sendToWebView({
-        type: 'streamEnd',
-        data: { timestamp: Date.now(), reason: 'user_cancelled' },
-      });
+      // Use sendStreamEnd to include requestId for proper correlation
+      this.sendStreamEnd('user_cancelled');
 
       console.log('[SessionMessageHandler] Streaming cancelled successfully');
     } catch (_error) {
       console.log('[SessionMessageHandler] Streaming cancelled (interrupted)');
 
-      // Always send streamEnd to update UI, regardless of errors
-      this.sendToWebView({
-        type: 'streamEnd',
-        data: { timestamp: Date.now(), reason: 'user_cancelled' },
-      });
+      // Use sendStreamEnd (with duplicate guard) to include requestId
+      this.sendStreamEnd('user_cancelled');
     }
   }
 

packages/vscode-ide-companion/src/webview/hooks/useWebViewMessages.ts

@@ -168,6 +168,9 @@ export const useWebViewMessages = ({
   // keep the bottom "waiting" message visible until all of them complete.
   const activeExecToolCallsRef = useRef<Set<string>>(new Set());
   const modelInfoRef = useRef<ModelInfo | null>(null);
+  // Track the active requestId from the latest streamStart so we can
+  // discard stale streamEnd events from cancelled/previous requests.
+  const activeRequestIdRef = useRef<string | null>(null);
   // Use ref to store callbacks to avoid useEffect dependency issues
   const handlersRef = useRef({
     sessionManagement,
@@ -461,11 +464,15 @@ export const useWebViewMessages = ({
           break;
         }
 
-        case 'streamStart':
-          handlers.messageHandling.startStreaming(
-            (message.data as { timestamp?: number } | undefined)?.timestamp,
-          );
+        case 'streamStart': {
+          const startData = message.data as
+            | { timestamp?: number; requestId?: string }
+            | undefined;
+          // Store the requestId so we can validate streamEnd events
+          activeRequestIdRef.current = startData?.requestId ?? null;
+          handlers.messageHandling.startStreaming(startData?.timestamp);
           break;
+        }
 
         case 'streamChunk': {
           handlers.messageHandling.appendStreamChunk(message.data.chunk);
@@ -479,6 +486,24 @@ export const useWebViewMessages = ({
         }
 
         case 'streamEnd': {
+          const endData = message.data as
+            | { reason?: string; requestId?: string }
+            | undefined;
+          const endRequestId = endData?.requestId ?? null;
+
+          // Drop stale or untagged streamEnd when a tagged stream is active.
+          if (activeRequestIdRef.current) {
+            if (endRequestId !== activeRequestIdRef.current) {
+              console.log(
+                '[useWebViewMessages] Ignoring stale/untagged streamEnd:',
+                endRequestId,
+                'active:',
+                activeRequestIdRef.current,
+              );
+              break;
+            }
+          }
+
           // Always end local streaming state and clear thinking state
           handlers.messageHandling.endStreaming();
           handlers.messageHandling.clearThinking();
@@ -488,9 +513,7 @@ export const useWebViewMessages = ({
           // This avoids UI getting stuck with Stop button visible after
           // rejecting a permission request.
           try {
-            const reason = (
-              (message.data as { reason?: string } | undefined)?.reason || ''
-            ).toLowerCase();
+            const reason = (endData?.reason || '').toLowerCase();
 
             /**
              * Handle different types of stream end reasons that require a full reset:

scripts/prepare-package.js

@@ -13,7 +13,6 @@
 import fs from 'node:fs';
 import path from 'node:path';
 import { fileURLToPath } from 'node:url';
-import { execSync } from 'node:child_process';
 
 const __filename = fileURLToPath(import.meta.url);
 const __dirname = path.dirname(__filename);
@@ -179,4 +178,17 @@ fs.writeFileSync(
 
 console.log('\n✅ Package prepared for publishing at dist/');
 console.log('\nPackage structure:');
-execSync('ls -lh dist/', { stdio: 'inherit', cwd: rootDir });
+// Use Node.js to list directory contents (cross-platform)
+const distFiles = fs.readdirSync(distDir);
+for (const file of distFiles) {
+  const filePath = path.join(distDir, file);
+  const stats = fs.statSync(filePath);
+  const size = stats.isDirectory() ? '<DIR>' : formatBytes(stats.size);
+  console.log(`  ${size.padEnd(12)} ${file}`);
+}
+
+function formatBytes(bytes) {
+  if (bytes < 1024) return `${bytes}B`;
+  if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)}KB`;
+  return `${(bytes / (1024 * 1024)).toFixed(1)}MB`;
+}