From c4b793cd0501221e7dd60a10e165d0556d8f0484 Mon Sep 17 00:00:00 2001
From: jwang19 <jwang19@atlassian.com>
Date: Thu, 8 Jan 2026 13:44:00 -0800
Subject: [PATCH] fix: check pending cancellation while waiting for response

---
 src/rovo-dev/rovoDevChatProvider.ts | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/src/rovo-dev/rovoDevChatProvider.ts b/src/rovo-dev/rovoDevChatProvider.ts
index 06ca0efda..b69efc03b 100644
--- a/src/rovo-dev/rovoDevChatProvider.ts
+++ b/src/rovo-dev/rovoDevChatProvider.ts
@@ -301,6 +301,14 @@ export class RovoDevChatProvider {
         const replayBuffer: RovoDevResponse[] = [];
 
         while (replayInProgress) {
+            // Check for cancellation during replay
+            if (this.pendingCancellation) {
+                await reader.cancel();
+                isDone = true;
+                this._pendingCancellation = false;
+                break;
+            }
+
             const { done, value } = await reader.read();
             if (done) {
                 isDone = true;
@@ -323,6 +331,14 @@ export class RovoDevChatProvider {
         }
 
         while (!isDone) {
+            // Check for cancellation before each read
+            if (this.pendingCancellation) {
+                await reader.cancel();
+                isDone = true;
+                this._pendingCancellation = false;
+                break;
+            }
+
             const { done, value } = await reader.read();
             if (done) {
                 isDone = true;
@@ -336,6 +352,14 @@ export class RovoDevChatProvider {
 
             const data = decoder.decode(value, { stream: true });
             for (const msg of parser.parse(data)) {
+                // Check for cancellation while processing messages
+                if (this.pendingCancellation) {
+                    await reader.cancel();
+                    isDone = true;
+                    this._pendingCancellation = false;
+                    break;
+                }
+
                 if (isFirstMessage) {
                     telemetryProvider?.perfLogger.promptFirstMessageReceived(this._currentPromptId);
                     isFirstMessage = false;
@@ -674,6 +698,10 @@ export class RovoDevChatProvider {
             } catch (error) {
                 // the error is retriable only when it happens during the streaming of a 'chat' response
                 await this.processError(error, { isRetriable: sourceApi === 'chat' });
+            } finally {
+                if (this.pendingCancellation) {
+                    this._pendingCancellation = false;
+                }
             }
         } else {
             await this.processError(new Error('RovoDev client not initialized'));