fix(cli): handle print abort and completion loading

im10furry · im10furry · commit 895ebdb32349 · 2026-07-03T13:15:03.000-05:00
diff --git a/apps/cli/src/app.tsx b/apps/cli/src/app.tsx
@@ -23,6 +23,7 @@ import {
   writeToStderr,
   writeToStdout,
 } from '#cli-utils/stdio'
+import { isPrintModeSignalAbortHandlingActive } from '#host-cli/entrypoints/cli/print/signalState'
 
 import { cursorShow } from 'ansi-escapes'
 import { openSync } from 'fs'
@@ -204,10 +205,15 @@ async function gracefulExit(code = 0) {
   process.exit(code)
 }
 
-process.on('SIGINT', () => void gracefulExit(0))
-process.on('SIGTERM', () => void gracefulExit(0))
+function handleProcessSignalExit(code = 0): void {
+  if (isPrintModeSignalAbortHandlingActive()) return
+  void gracefulExit(code)
+}
+
+process.on('SIGINT', () => handleProcessSignalExit(0))
+process.on('SIGTERM', () => handleProcessSignalExit(0))
 // Windows CTRL+BREAK
-process.on('SIGBREAK', () => void gracefulExit(0))
+process.on('SIGBREAK', () => handleProcessSignalExit(0))
 process.on('unhandledRejection', err => {
   logError(err)
   void gracefulExit(1)
diff --git a/apps/cli/src/entrypoints/cli/print/runSingleTurn.ts b/apps/cli/src/entrypoints/cli/print/runSingleTurn.ts
@@ -4,6 +4,7 @@ import type { QueryToolUseContext } from '#core/engine'
 import { MaxBudgetUsdExceededError } from '#core/errors/maxBudgetUsd'
 import { MaxTurnsExceededError } from '#core/errors/maxTurns'
 import { randomUUID } from 'crypto'
+import { beginPrintModeSignalAbortHandling } from './signalState'
 
 const PRINT_MODE_ABORT_SIGNALS: NodeJS.Signals[] = [
   'SIGINT',
@@ -14,6 +15,7 @@ const PRINT_MODE_ABORT_SIGNALS: NodeJS.Signals[] = [
 function installPrintModeSignalAbort(
   abortController: AbortController,
 ): () => void {
+  const endPrintModeSignalAbortHandling = beginPrintModeSignalAbortHandling()
   const abort = () => {
     abortController.abort()
   }
@@ -26,6 +28,7 @@ function installPrintModeSignalAbort(
     for (const signal of PRINT_MODE_ABORT_SIGNALS) {
       process.removeListener(signal, abort)
     }
+    endPrintModeSignalAbortHandling()
   }
 }
 
diff --git a/apps/cli/src/entrypoints/cli/print/signalState.ts b/apps/cli/src/entrypoints/cli/print/signalState.ts
@@ -0,0 +1,19 @@
+let activePrintSignalAbortHandlers = 0
+
+export function beginPrintModeSignalAbortHandling(): () => void {
+  activePrintSignalAbortHandlers += 1
+  let disposed = false
+
+  return () => {
+    if (disposed) return
+    disposed = true
+    activePrintSignalAbortHandlers = Math.max(
+      0,
+      activePrintSignalAbortHandlers - 1,
+    )
+  }
+}
+
+export function isPrintModeSignalAbortHandlingActive(): boolean {
+  return activePrintSignalAbortHandlers > 0
+}
diff --git a/apps/cli/src/ui/hooks/useUnifiedCompletion/actions.ts b/apps/cli/src/ui/hooks/useUnifiedCompletion/actions.ts
@@ -1,8 +1,9 @@
 import { useCallback } from 'react'
 
-import type {
-  CompletionContext,
-  UnifiedSuggestion,
+import {
+  isLoadingSuggestion,
+  type CompletionContext,
+  type UnifiedSuggestion,
 } from '#cli-utils/completion/types'
 
 export function useCompletionActions(args: {
@@ -18,6 +19,8 @@ export function useCompletionActions(args: {
 } {
   const completeWith = useCallback(
     (suggestion: UnifiedSuggestion, context: CompletionContext) => {
+      if (isLoadingSuggestion(suggestion)) return
+
       let completion: string
 
       if (context.type === 'command') {
diff --git a/apps/cli/src/ui/hooks/useUnifiedCompletion/hook.test.ts b/apps/cli/src/ui/hooks/useUnifiedCompletion/hook.test.ts
@@ -0,0 +1,88 @@
+import { describe, expect, test } from 'bun:test'
+
+import { __computeCompletionRefreshForTests } from './hook'
+import type { CompletionState } from './types'
+
+function makeState(overrides: Partial<CompletionState>): CompletionState {
+  return {
+    suggestions: [],
+    selectedIndex: 0,
+    isActive: false,
+    context: null,
+    preview: null,
+    emptyDirMessage: '',
+    suppressUntil: 0,
+    ...overrides,
+  }
+}
+
+describe('__computeCompletionRefreshForTests', () => {
+  test('refreshes active loading suggestions when new command matches arrive', () => {
+    const state = makeState({
+      isActive: true,
+      context: {
+        type: 'file',
+        prefix: 'kub',
+        startPos: 0,
+        endPos: 3,
+      },
+      suggestions: [
+        {
+          value: 'loading...',
+          displayValue: 'Loading system commands...',
+          type: 'file',
+          score: 0,
+          metadata: { isLoading: true },
+        },
+      ],
+    })
+
+    const result = __computeCompletionRefreshForTests({
+      isEnabled: true,
+      state,
+      suggestions: [
+        {
+          value: 'kubectl',
+          displayValue: '$ kubectl',
+          type: 'command',
+          score: 100,
+        },
+      ],
+    })
+
+    expect(result.action).toBe('refresh')
+    if (result.action !== 'refresh') return
+    expect(result.suggestions[0]?.value).toBe('kubectl')
+    expect(result.selectedIndex).toBe(0)
+  })
+
+  test('does not refresh while a completion preview is active', () => {
+    const result = __computeCompletionRefreshForTests({
+      isEnabled: true,
+      state: makeState({
+        isActive: true,
+        context: {
+          type: 'file',
+          prefix: 'gi',
+          startPos: 0,
+          endPos: 2,
+        },
+        preview: {
+          isActive: true,
+          originalInput: 'gi',
+          wordRange: [0, 3],
+        },
+      }),
+      suggestions: [
+        {
+          value: 'git',
+          displayValue: '$ git',
+          type: 'command',
+          score: 100,
+        },
+      ],
+    })
+
+    expect(result.action).toBe('none')
+  })
+})
diff --git a/apps/cli/src/ui/hooks/useUnifiedCompletion/hook.ts b/apps/cli/src/ui/hooks/useUnifiedCompletion/hook.ts
@@ -25,6 +25,34 @@ export function __getCompletionContextForTests(args: {
   return getCompletionContext(args)
 }
 
+export function __computeCompletionRefreshForTests(args: {
+  isEnabled: boolean
+  state: CompletionState
+  suggestions: UnifiedSuggestion[]
+}):
+  | { action: 'none' }
+  | { action: 'reset' }
+  | {
+      action: 'refresh'
+      suggestions: UnifiedSuggestion[]
+      selectedIndex: number
+    } {
+  if (!args.isEnabled || !args.state.isActive || !args.state.context) {
+    return { action: 'none' }
+  }
+  if (args.state.preview?.isActive) return { action: 'none' }
+  if (args.suggestions.length === 0) return { action: 'reset' }
+
+  return {
+    action: 'refresh',
+    suggestions: args.suggestions,
+    selectedIndex: Math.min(
+      args.state.selectedIndex,
+      args.suggestions.length - 1,
+    ),
+  }
+}
+
 export function useUnifiedCompletion({
   input,
   cursorOffset,
@@ -110,6 +138,38 @@ export function useUnifiedCompletion({
     }
   }, [isEnabled, resetCompletion, state.isActive])
 
+  useEffect(() => {
+    if (!state.context) return
+
+    const nextSuggestions = generateSuggestions(state.context)
+    const result = __computeCompletionRefreshForTests({
+      isEnabled,
+      state,
+      suggestions: nextSuggestions,
+    })
+
+    if (result.action === 'reset') {
+      resetCompletion()
+      return
+    }
+
+    if (result.action === 'refresh') {
+      setState(prev => ({
+        ...prev,
+        suggestions: result.suggestions,
+        selectedIndex: result.selectedIndex,
+      }))
+    }
+  }, [
+    generateSuggestions,
+    isEnabled,
+    resetCompletion,
+    state.context,
+    state.isActive,
+    state.selectedIndex,
+    state.preview?.isActive,
+  ])
+
   useUnifiedCompletionTabKey({
     input,
     state,
diff --git a/apps/cli/src/ui/hooks/useUnifiedCompletion/useNavigationKeys.ts b/apps/cli/src/ui/hooks/useUnifiedCompletion/useNavigationKeys.ts
@@ -1,9 +1,10 @@
 import { useKeypress } from '#ui-ink/hooks/useKeypress'
 import { KEYPRESS_PRIORITY } from '#ui-ink/constants/keypressPriority'
 
-import type {
-  CompletionContext,
-  UnifiedSuggestion,
+import {
+  isLoadingSuggestion,
+  type CompletionContext,
+  type UnifiedSuggestion,
 } from '#cli-utils/completion/types'
 import type { CompletionState } from './types'
 
@@ -63,6 +64,10 @@ export function useUnifiedCompletionNavigationKeys(args: {
         args.state.context
       ) {
         const context = args.state.context
+        const selectedSuggestion =
+          args.state.suggestions[args.state.selectedIndex]
+        if (isLoadingSuggestion(selectedSuggestion)) return true
+
         const isFileCompletion = context.type === 'file'
 
         if (isFileCompletion) {
@@ -76,8 +81,6 @@ export function useUnifiedCompletionNavigationKeys(args: {
           return false
         } else {
           // Command/agent/other completion: fill the selected suggestion
-          const selectedSuggestion =
-            args.state.suggestions[args.state.selectedIndex]
           args.completeWith(selectedSuggestion, context)
           args.resetCompletion()
           return true
@@ -94,6 +97,10 @@ export function useUnifiedCompletionNavigationKeys(args: {
         }
 
         const suggestion = args.state.suggestions[newIndex]
+        if (isLoadingSuggestion(suggestion)) {
+          args.updateState({ selectedIndex: newIndex })
+          return
+        }
         const previewValue = getPreviewText(suggestion, args.state.context)
 
         if (args.state.preview?.isActive && args.state.context) {
@@ -151,6 +158,8 @@ export function useUnifiedCompletionNavigationKeys(args: {
       if (key.rightArrow) {
         const selectedSuggestion =
           args.state.suggestions[args.state.selectedIndex]
+        if (isLoadingSuggestion(selectedSuggestion)) return true
+
         const isDirectory = selectedSuggestion.value.endsWith('/')
 
         if (!args.state.context) return false
diff --git a/apps/cli/src/ui/hooks/useUnifiedCompletion/useTabKey.ts b/apps/cli/src/ui/hooks/useUnifiedCompletion/useTabKey.ts
@@ -1,9 +1,10 @@
 import { useKeypress, type Key } from '#ui-ink/hooks/useKeypress'
 import { KEYPRESS_PRIORITY } from '#ui-ink/constants/keypressPriority'
 
-import type {
-  CompletionContext,
-  UnifiedSuggestion,
+import {
+  isLoadingSuggestion,
+  type CompletionContext,
+  type UnifiedSuggestion,
 } from '#cli-utils/completion/types'
 import type { CompletionState } from './types'
 
@@ -43,6 +44,7 @@ export function useUnifiedCompletionTabKey(args: {
         const nextIndex =
           (args.state.selectedIndex + 1) % args.state.suggestions.length
         const nextSuggestion = args.state.suggestions[nextIndex]
+        if (isLoadingSuggestion(nextSuggestion)) return true
 
         if (args.state.context) {
           const currentWord = args.input.slice(args.state.context.startPos)
@@ -95,13 +97,18 @@ export function useUnifiedCompletionTabKey(args: {
         return false
       }
       if (currentSuggestions.length === 1) {
+        if (isLoadingSuggestion(currentSuggestions[0])) {
+          args.activateCompletion(currentSuggestions, context)
+          return true
+        }
         args.completeWith(currentSuggestions[0], context)
         return true
       }
 
       args.activateCompletion(currentSuggestions, context)
 
       const firstSuggestion = currentSuggestions[0]
+      if (isLoadingSuggestion(firstSuggestion)) return true
       const currentWord = args.input.slice(context.startPos)
       const wordEnd = currentWord.search(/\s/)
       const actualEndPos =
diff --git a/apps/cli/src/utils/completion/types.ts b/apps/cli/src/utils/completion/types.ts
@@ -17,3 +17,9 @@ export interface CompletionContext {
   endPos: number
   trigger?: '@' | '/' | null
 }
+
+export function isLoadingSuggestion(
+  suggestion: UnifiedSuggestion | undefined,
+): boolean {
+  return Boolean(suggestion?.metadata?.isLoading)
+}
diff --git a/packages/core/src/test/unit/print-mode-signal-abort.test.ts b/packages/core/src/test/unit/print-mode-signal-abort.test.ts
@@ -1,5 +1,6 @@
 import { describe, expect, test } from 'bun:test'
 import { __installPrintModeSignalAbortForTests } from '#host-cli/entrypoints/cli/print/runSingleTurn'
+import { isPrintModeSignalAbortHandlingActive } from '#host-cli/entrypoints/cli/print/signalState'
 
 describe('print mode signal cancellation', () => {
   test('SIGINT aborts the active print turn controller', () => {
@@ -27,4 +28,18 @@ describe('print mode signal cancellation', () => {
 
     expect(controller.signal.aborted).toBe(false)
   })
+
+  test('tracks active print signal handling for global handler deferral', () => {
+    const controller = new AbortController()
+    expect(isPrintModeSignalAbortHandlingActive()).toBe(false)
+
+    const cleanup = __installPrintModeSignalAbortForTests(controller)
+    try {
+      expect(isPrintModeSignalAbortHandlingActive()).toBe(true)
+    } finally {
+      cleanup()
+    }
+
+    expect(isPrintModeSignalAbortHandlingActive()).toBe(false)
+  })
 })

Original file line number	Diff line number	Diff line change
`@@ -4,6 +4,7 @@ import type { QueryToolUseContext } from '#core/engine'`
`4`	`4`	`import { MaxBudgetUsdExceededError } from '#core/errors/maxBudgetUsd'`
`5`	`5`	`import { MaxTurnsExceededError } from '#core/errors/maxTurns'`
`6`	`6`	`import { randomUUID } from 'crypto'`
	`7`	`+import { beginPrintModeSignalAbortHandling } from './signalState'`
`7`	`8`
`8`	`9`	`const PRINT_MODE_ABORT_SIGNALS: NodeJS.Signals[] = [`
`9`	`10`	`'SIGINT',`
`@@ -14,6 +15,7 @@ const PRINT_MODE_ABORT_SIGNALS: NodeJS.Signals[] = [`
`14`	`15`	`function installPrintModeSignalAbort(`
`15`	`16`	`abortController: AbortController,`
`16`	`17`	`): () => void {`
	`18`	`+ const endPrintModeSignalAbortHandling = beginPrintModeSignalAbortHandling()`
`17`	`19`	`const abort = () => {`
`18`	`20`	`abortController.abort()`
`19`	`21`	`}`
`@@ -26,6 +28,7 @@ function installPrintModeSignalAbort(`
`26`	`28`	`for (const signal of PRINT_MODE_ABORT_SIGNALS) {`
`27`	`29`	`process.removeListener(signal, abort)`
`28`	`30`	`}`
	`31`	`+ endPrintModeSignalAbortHandling()`
`29`	`32`	`}`
`30`	`33`	`}`
`31`	`34`