fix(openai): abort print stream on signal

Author: im10furry

apps/cli/src/entrypoints/cli/print/runSingleTurn.ts

@@ -5,6 +5,30 @@ import { MaxBudgetUsdExceededError } from '#core/errors/maxBudgetUsd'
 import { MaxTurnsExceededError } from '#core/errors/maxTurns'
 import { randomUUID } from 'crypto'
 
+const PRINT_MODE_ABORT_SIGNALS: NodeJS.Signals[] = [
+  'SIGINT',
+  'SIGTERM',
+  'SIGBREAK',
+]
+
+function installPrintModeSignalAbort(
+  abortController: AbortController,
+): () => void {
+  const abort = () => {
+    abortController.abort()
+  }
+
+  for (const signal of PRINT_MODE_ABORT_SIGNALS) {
+    process.prependListener(signal, abort)
+  }
+
+  return () => {
+    for (const signal of PRINT_MODE_ABORT_SIGNALS) {
+      process.removeListener(signal, abort)
+    }
+  }
+}
+
 type RunTurnFn = (args: {
   messages: Message[]
   canUseTool: CanUseToolFn
@@ -54,6 +78,9 @@ export async function runSingleTurnPrint(args: {
 }): Promise<void> {
   let lastAssistant: Message | null = null
   let queryError: unknown = null
+  const cleanupSignalAbort = installPrintModeSignalAbort(
+    args.toolUseContext.abortController,
+  )
 
   try {
     for await (const m of args.runTurn({
@@ -75,6 +102,8 @@ export async function runSingleTurnPrint(args: {
       args.toolUseContext.abortController.abort()
     } catch {}
     queryError = e
+  } finally {
+    cleanupSignalAbort()
   }
 
   const totalCostUsd = args.getTotalCostUsd()
@@ -219,3 +248,5 @@ export async function runSingleTurnPrint(args: {
   }
   process.exit((resultMsg as any)?.is_error ? 1 : 0)
 }
+
+export const __installPrintModeSignalAbortForTests = installPrintModeSignalAbort

packages/core/src/ai/llm/openai/stream.ts

@@ -52,6 +52,12 @@ function messageReducer(
   return reduce(previous, choice.delta) as OpenAI.ChatCompletionMessage
 }
 
+function throwIfAborted(signal?: AbortSignal): void {
+  if (signal?.aborted) {
+    throw new Error('Request was cancelled')
+  }
+}
+
 export async function handleMessageStream(
   stream: ChatCompletionStream,
   signal?: AbortSignal,
@@ -71,13 +77,16 @@ export async function handleMessageStream(
 
   let id, model, created, object, usage
   try {
+    throwIfAborted(signal)
     for await (const chunk of stream) {
-      if (signal?.aborted) {
+      try {
+        throwIfAborted(signal)
+      } catch (error) {
         debugLogger.flow('OPENAI_STREAM_ABORTED', {
           chunkCount,
           timestamp: Date.now(),
         })
-        throw new Error('Request was cancelled')
+        throw error
       }
 
       chunkCount++
@@ -148,6 +157,8 @@ export async function handleMessageStream(
       }
     }
 
+    throwIfAborted(signal)
+
     debugLogger.api('OPENAI_STREAM_COMPLETE', {
       totalChunks: String(chunkCount),
       errorCount: String(errorCount),

packages/core/src/test/unit/openai-stream-abort.test.ts

@@ -0,0 +1,40 @@
+import { describe, expect, test } from 'bun:test'
+import { handleMessageStream } from '#core/ai/llm/openai/stream'
+
+function chunk(delta: Record<string, unknown>) {
+  return {
+    id: 'chatcmpl_test',
+    model: 'gpt-4',
+    created: 1,
+    object: 'chat.completion.chunk',
+    choices: [{ index: 0, delta, finish_reason: null }],
+  }
+}
+
+describe('OpenAI stream cancellation', () => {
+  test('rejects when signal is aborted before reading stream chunks', async () => {
+    const controller = new AbortController()
+    controller.abort()
+
+    async function* stream() {
+      yield chunk({ content: 'late' })
+    }
+
+    await expect(
+      handleMessageStream(stream() as any, controller.signal),
+    ).rejects.toThrow('Request was cancelled')
+  })
+
+  test('does not return a partial response when signal aborts after a chunk', async () => {
+    const controller = new AbortController()
+
+    async function* stream() {
+      yield chunk({ content: 'partial' })
+      controller.abort()
+    }
+
+    await expect(
+      handleMessageStream(stream() as any, controller.signal),
+    ).rejects.toThrow('Request was cancelled')
+  })
+})

packages/core/src/test/unit/print-mode-signal-abort.test.ts

@@ -0,0 +1,30 @@
+import { describe, expect, test } from 'bun:test'
+import { __installPrintModeSignalAbortForTests } from '#host-cli/entrypoints/cli/print/runSingleTurn'
+
+describe('print mode signal cancellation', () => {
+  test('SIGINT aborts the active print turn controller', () => {
+    const controller = new AbortController()
+    const listenerCountBefore = process.listenerCount('SIGINT')
+    const cleanup = __installPrintModeSignalAbortForTests(controller)
+
+    try {
+      expect(process.listenerCount('SIGINT')).toBe(listenerCountBefore + 1)
+      process.emit('SIGINT')
+      expect(controller.signal.aborted).toBe(true)
+    } finally {
+      cleanup()
+    }
+
+    expect(process.listenerCount('SIGINT')).toBe(listenerCountBefore)
+  })
+
+  test('cleanup removes print turn signal handlers', () => {
+    const controller = new AbortController()
+    const cleanup = __installPrintModeSignalAbortForTests(controller)
+    cleanup()
+
+    process.emit('SIGINT')
+
+    expect(controller.signal.aborted).toBe(false)
+  })
+})