diff --git a/extensions/cli/src/commands/chat.ts b/extensions/cli/src/commands/chat.ts index a32df348818..b953eb31f64 100644 --- a/extensions/cli/src/commands/chat.ts +++ b/extensions/cli/src/commands/chat.ts @@ -33,6 +33,7 @@ import { calculateContextUsagePercentage, countChatHistoryTokens, } from "../util/tokenizer.js"; +import { withSigintAbort } from "../util/withSigintAbort.js"; import { ExtendedCommandOptions } from "./BaseCommandOptions.js"; @@ -169,28 +170,14 @@ async function getStreamingResponse( model: ModelConfig, llmApi: BaseLlmApi, ): Promise { - const abortController = new AbortController(); + return withSigintAbort(async (abortController) => { + const history = + compactionIndex !== null && compactionIndex !== undefined + ? services.chatHistory.getHistoryForLLM(compactionIndex) + : services.chatHistory.getHistory(); - if (compactionIndex !== null && compactionIndex !== undefined) { - // Use service to compute history for LLM - const historyForLLM = - services.chatHistory.getHistoryForLLM(compactionIndex); - - return await streamChatResponse( - historyForLLM, - model, - llmApi, - abortController, - ); - } else { - // No compaction - get full history from service - return await streamChatResponse( - services.chatHistory.getHistory(), - model, - llmApi, - abortController, - ); - } + return await streamChatResponse(history, model, llmApi, abortController); + }); } // Helper function to process and output response in headless mode diff --git a/extensions/cli/src/util/withSigintAbort.test.ts b/extensions/cli/src/util/withSigintAbort.test.ts new file mode 100644 index 00000000000..0d22b2f3bf5 --- /dev/null +++ b/extensions/cli/src/util/withSigintAbort.test.ts @@ -0,0 +1,51 @@ +import { vi } from "vitest"; + +import { withSigintAbort } from "./withSigintAbort.js"; + +describe("withSigintAbort", () => { + it("aborts the active operation on SIGINT and prevents its late effect", async () => { + const listenersBefore = process.listeners("SIGINT"); + const lateEffect = vi.fn(); + let finishOperation: (() => void) | undefined; + let operationSignal: AbortSignal | undefined; + + const operation = withSigintAbort(async (abortController) => { + operationSignal = abortController.signal; + await new Promise((resolve) => { + finishOperation = resolve; + }); + + if (!abortController.signal.aborted) { + lateEffect(); + } + + return "finished"; + }); + + const addedListeners = process + .listeners("SIGINT") + .filter((listener) => !listenersBefore.includes(listener)); + + expect(addedListeners).toHaveLength(1); + addedListeners[0]("SIGINT"); + finishOperation?.(); + + await expect(operation).resolves.toBe("finished"); + expect(operationSignal?.aborted).toBe(true); + expect(lateEffect).not.toHaveBeenCalled(); + expect(process.listeners("SIGINT")).toEqual(listenersBefore); + }); + + it("removes its SIGINT listener when the operation rejects", async () => { + const listenersBefore = process.listeners("SIGINT"); + const expectedError = new Error("stream failed"); + + await expect( + withSigintAbort(async () => { + throw expectedError; + }), + ).rejects.toBe(expectedError); + + expect(process.listeners("SIGINT")).toEqual(listenersBefore); + }); +}); diff --git a/extensions/cli/src/util/withSigintAbort.ts b/extensions/cli/src/util/withSigintAbort.ts new file mode 100644 index 00000000000..61178ccdd93 --- /dev/null +++ b/extensions/cli/src/util/withSigintAbort.ts @@ -0,0 +1,14 @@ +export async function withSigintAbort( + operation: (abortController: AbortController) => Promise, +): Promise { + const abortController = new AbortController(); + const abortOnSigint = () => abortController.abort(); + + process.once("SIGINT", abortOnSigint); + + try { + return await operation(abortController); + } finally { + process.off("SIGINT", abortOnSigint); + } +}