Merge pull request #7272 from continuedev/nate/cli-fixes

fix: max abort listeners warning

Author: sestinj

extensions/cli/src/streamChatResponse.ts

@@ -300,7 +300,7 @@ export async function processStreamingResponse(
   } = options;
   const requestStartTime = Date.now();
 
-  const streamFactory = async () => {
+  const streamFactory = async (retryAbortSignal: AbortSignal) => {
     logger.debug("Creating chat completion stream", {
       model,
       messageCount: chatHistory.length,
@@ -315,7 +315,7 @@ export async function processStreamingResponse(
         tools,
         ...getDefaultCompletionOptions(model.defaultCompletionOptions),
       },
-      abortController.signal,
+      retryAbortSignal,
     );
   };
 

extensions/cli/src/ui/hooks/useChat.ts

@@ -111,6 +111,15 @@ export function useChat({
     return null;
   });
 
+  // Clean up abort controller on unmount
+  useEffect(() => {
+    return () => {
+      if (abortController && !abortController.signal.aborted) {
+        abortController.abort();
+      }
+    };
+  }, [abortController]);
+
   // Remote mode polling
   useEffect(() => {
     if (!isRemoteMode || !remoteUrl || process.env.NODE_ENV === "test") return;
@@ -158,6 +167,92 @@ export function useChat({
     }
   }, [initialPrompt, isChatHistoryInitialized]);
 
+  const executeStreamingResponse = async (
+    newHistory: ChatCompletionMessageParam[],
+    currentCompactionIndex: number | null,
+    message: string,
+  ) => {
+    // Clean up previous abort controller if it exists
+    if (abortController && !abortController.signal.aborted) {
+      abortController.abort();
+    }
+
+    // Start streaming response
+    const controller = new AbortController();
+    setAbortController(controller);
+    setIsWaitingForResponse(true);
+    setResponseStartTime(Date.now());
+    setInputMode(false);
+    logger.debug("Starting chat response stream", {
+      messageLength: message.length,
+      historyLength: newHistory.length,
+    });
+
+    try {
+      const currentStreamingMessageRef = {
+        current: null as DisplayMessage | null,
+      };
+      const streamCallbacks = createStreamCallbacks(
+        { setMessages, setActivePermissionRequest },
+        currentStreamingMessageRef,
+      );
+
+      // Execute streaming chat response
+      await executeStreaming({
+        newHistory,
+        model,
+        llmApi,
+        controller,
+        streamCallbacks,
+        currentCompactionIndex,
+      });
+
+      if (
+        currentStreamingMessageRef.current &&
+        currentStreamingMessageRef.current.content
+      ) {
+        const messageContent = currentStreamingMessageRef.current.content;
+        setMessages((prev) => [
+          ...prev,
+          {
+            role: "assistant",
+            content: messageContent,
+            isStreaming: false,
+          },
+        ]);
+      }
+
+      // Update the chat history with the complete conversation after streaming
+      setChatHistory(newHistory);
+      logger.debug("Chat history updated", {
+        finalHistoryLength: newHistory.length,
+      });
+
+      // Save the updated history to session
+      logger.debug("Saving session", { historyLength: newHistory.length });
+      saveSession(newHistory);
+      logger.debug("Session saved");
+    } catch (error: any) {
+      const errorMessage = `Error: ${formatError(error)}`;
+      setMessages((prev) => [
+        ...prev,
+        {
+          role: "system",
+          content: errorMessage,
+          messageType: "system" as const,
+        },
+      ]);
+    } finally {
+      // Stop active time tracking
+      telemetryService.stopActiveTime();
+
+      setAbortController(null);
+      setIsWaitingForResponse(false);
+      setResponseStartTime(null);
+      setInputMode(true);
+    }
+  };
+
   const handleUserMessage = async (message: string) => {
     // Handle special commands
     const handled = await handleSpecialCommands({
@@ -247,80 +342,8 @@ export function useChat({
     setChatHistory(newHistory);
     setMessages((prev) => [...prev, { role: "user", content: message }]);
 
-    // Start streaming response
-    const controller = new AbortController();
-    setAbortController(controller);
-    setIsWaitingForResponse(true);
-    setResponseStartTime(Date.now());
-    setInputMode(false);
-    logger.debug("Starting chat response stream", {
-      messageLength: message.length,
-      historyLength: newHistory.length,
-    });
-
-    try {
-      const currentStreamingMessageRef = {
-        current: null as DisplayMessage | null,
-      };
-      const streamCallbacks = createStreamCallbacks(
-        { setMessages, setActivePermissionRequest },
-        currentStreamingMessageRef,
-      );
-
-      // Execute streaming chat response
-      await executeStreaming({
-        newHistory,
-        model,
-        llmApi,
-        controller,
-        streamCallbacks,
-        currentCompactionIndex,
-      });
-
-      if (
-        currentStreamingMessageRef.current &&
-        currentStreamingMessageRef.current.content
-      ) {
-        const messageContent = currentStreamingMessageRef.current.content;
-        setMessages((prev) => [
-          ...prev,
-          {
-            role: "assistant",
-            content: messageContent,
-            isStreaming: false,
-          },
-        ]);
-      }
-
-      // Update the chat history with the complete conversation after streaming
-      setChatHistory(newHistory);
-      logger.debug("Chat history updated", {
-        finalHistoryLength: newHistory.length,
-      });
-
-      // Save the updated history to session
-      logger.debug("Saving session", { historyLength: newHistory.length });
-      saveSession(newHistory);
-      logger.debug("Session saved");
-    } catch (error: any) {
-      const errorMessage = `Error: ${formatError(error)}`;
-      setMessages((prev) => [
-        ...prev,
-        {
-          role: "system",
-          content: errorMessage,
-          messageType: "system" as const,
-        },
-      ]);
-    } finally {
-      // Stop active time tracking
-      telemetryService.stopActiveTime();
-
-      setAbortController(null);
-      setIsWaitingForResponse(false);
-      setResponseStartTime(null);
-      setInputMode(true);
-    }
+    // Execute the streaming response
+    await executeStreamingResponse(newHistory, currentCompactionIndex, message);
   };
 
   const handleInterrupt = () => {

extensions/cli/src/util/exponentialBackoff.ts

@@ -192,23 +192,38 @@ export async function chatCompletionStreamWithBackoff(
  * If the generator throws an error during iteration, it will retry the entire generator.
  */
 export async function* withExponentialBackoff<T>(
-  generatorFactory: () => Promise<AsyncGenerator<T>>,
+  generatorFactory: (
+    retryAbortSignal: AbortSignal,
+  ) => Promise<AsyncGenerator<T>>,
   abortSignal: AbortSignal,
   options: ExponentialBackoffOptions = {},
 ): AsyncGenerator<T> {
   const opts = { ...DEFAULT_OPTIONS, ...options };
   let lastError: any;
 
   for (let attempt = 0; attempt <= opts.maxRetries; attempt++) {
-    const yieldedValues: T[] = [];
+    // Create a new AbortController for this retry attempt
+    // This prevents accumulating listeners on the original signal
+    const retryAbortController = new AbortController();
+
+    // Forward abort from the original signal to the retry signal
+    const abortListener = () => {
+      retryAbortController.abort();
+    };
+
+    if (abortSignal.aborted) {
+      retryAbortController.abort();
+    } else if (typeof abortSignal.addEventListener === "function") {
+      abortSignal.addEventListener("abort", abortListener);
+    }
 
     try {
       // Check if we should abort before creating the generator
       if (abortSignal.aborted) {
         throw new Error("Request aborted");
       }
 
-      const generator = await generatorFactory();
+      const generator = await generatorFactory(retryAbortController.signal);
 
       // Iterate through the generator and yield each value
       for await (const chunk of generator) {
@@ -217,15 +232,24 @@ export async function* withExponentialBackoff<T>(
           throw new Error("Request aborted");
         }
 
-        yieldedValues.push(chunk);
         yield chunk;
       }
 
+      // Clean up the abort listener since we succeeded
+      if (typeof abortSignal.removeEventListener === "function") {
+        abortSignal.removeEventListener("abort", abortListener);
+      }
+
       // If we successfully completed the generator, we're done
       return;
     } catch (err: any) {
       lastError = err;
 
+      // Clean up the abort listener
+      if (typeof abortSignal.removeEventListener === "function") {
+        abortSignal.removeEventListener("abort", abortListener);
+      }
+
       // Don't retry if the request was aborted
       if (abortSignal.aborted) {
         throw err;

extensions/cli/src/util/withExponentialBackoff.test.ts

@@ -20,7 +20,7 @@ describe("withExponentialBackoff", () => {
 
   it("should successfully yield all values from generator", async () => {
     const mockData = ["chunk1", "chunk2", "chunk3"];
-    const generatorFactory = vi.fn(async () => {
+    const generatorFactory = vi.fn(async (retryAbortSignal: AbortSignal) => {
       return (async function* () {
         for (const chunk of mockData) {
           yield chunk;
@@ -44,7 +44,7 @@ describe("withExponentialBackoff", () => {
 
   it("should retry on retryable errors during generator creation", async () => {
     let callCount = 0;
-    const generatorFactory = vi.fn(async () => {
+    const generatorFactory = vi.fn(async (retryAbortSignal: AbortSignal) => {
       callCount++;
       if (callCount === 1) {
         const error = new Error("Connection reset");
@@ -75,7 +75,7 @@ describe("withExponentialBackoff", () => {
   });
 
   it("should not retry on non-retryable errors", async () => {
-    const generatorFactory = vi.fn(async () => {
+    const generatorFactory = vi.fn(async (retryAbortSignal: AbortSignal) => {
       const error = new Error("Bad request");
       (error as any).status = 400;
       throw error;
@@ -96,7 +96,7 @@ describe("withExponentialBackoff", () => {
   });
 
   it("should handle abort signal", async () => {
-    const generatorFactory = vi.fn(async () => {
+    const generatorFactory = vi.fn(async (retryAbortSignal: AbortSignal) => {
       return (async function* () {
         yield "should-not-yield";
       })();