diff --git a/src/components/Chat/Chat.test.tsx b/src/components/Chat/Chat.test.tsx index 926a94a2..123e4ecf 100644 --- a/src/components/Chat/Chat.test.tsx +++ b/src/components/Chat/Chat.test.tsx @@ -78,15 +78,27 @@ vi.mock('../../utils', async () => ({ }, })); +const interruptState = vi.hoisted(() => ({ + handler: undefined as (() => void) | undefined, + clear() { + this.handler = undefined; + }, +})); + vi.mock('./Input', () => ({ Input: (props: { onSubmit?: (value: string) => void; + onInterrupt?: () => void; isDisabled?: boolean; }) => { if (props.onSubmit) { mockState.handler = props.onSubmit; } + if (props.onInterrupt) { + interruptState.handler = props.onInterrupt; + } + if (props.isDisabled) { return null; } @@ -134,12 +146,17 @@ async function waitForStream() { await time.tick(10); } +function fireInterrupt() { + interruptState.handler?.(); +} + function resetChatMocks() { vi.restoreAllMocks(); vi.clearAllMocks(); mockState.clear(); planApprovalState.clear(); toolApprovalState.clear(); + interruptState.clear(); tools.TOOLS.splice(0, tools.TOOLS.length); vi.mocked(ollama.streamChat).mockImplementation(async function* () { await Promise.resolve(); @@ -312,6 +329,7 @@ describe('Chat', () => { expect.any(Array), 'llama3', expect.any(Array), + expect.any(AbortSignal), ); }); @@ -1021,8 +1039,10 @@ describe('Chat with tool calls', () => { await waitForStream(); rerender(chat); - // Should show rejection message - expect(lastFrame()).toContain('declined'); + expect(lastFrame()).not.toContain('Tool requires approval'); + expect(lastFrame()).toContain('❗ Tool call rejected.'); + expect(lastFrame()).toContain('>'); + expect(vi.mocked(ollama.streamChat)).toHaveBeenCalledOnce(); }); it('handles tool approval acceptance', async () => { @@ -1235,3 +1255,75 @@ describe('Chat with error', () => { expect(lastFrame()).toContain('Error: Plan generation crashed'); }); }); + +describe('Chat interrupt', () => { + beforeEach(() => { + resetChatMocks(); + }); + + it('shows interrupt notice and turn_aborted message when interrupted during streaming', async () => { + vi.mocked(ollama.streamChat).mockImplementation(async function* () { + yield { type: 'content', content: 'Partial' }; + await new Promise(() => undefined); + }); + + const chat = ( + + ); + const { lastFrame, rerender } = render(chat); + submitInput('hello'); + rerender(chat); + await time.tick(); + + fireInterrupt(); + rerender(chat); + await time.tick(); + + const frame = lastFrame() ?? ''; + expect(frame).toContain('❗ Execution interrupted'); + expect(frame).not.toContain('turn_aborted'); + expect(frame).toContain('>'); + }); + + it('clears interrupt notice on next submit', async () => { + vi.mocked(ollama.streamChat).mockImplementation(async function* () { + yield { type: 'content', content: 'Partial' }; + await new Promise(() => undefined); + }); + + const chat = ( + + ); + const { lastFrame, rerender } = render(chat); + submitInput('hello'); + rerender(chat); + await time.tick(); + + fireInterrupt(); + rerender(chat); + await time.tick(); + expect(lastFrame()).toContain('❗ Execution interrupted'); + + vi.mocked(ollama.streamChat).mockImplementation(async function* () { + await Promise.resolve(); + yield { type: 'content', content: 'New response' }; + }); + submitInput('continue'); + rerender(chat); + await waitForStream(); + + expect(lastFrame()).not.toContain('❗ Execution interrupted'); + }); +}); diff --git a/src/components/Chat/Chat.tsx b/src/components/Chat/Chat.tsx index 38c279a0..b99e91e6 100644 --- a/src/components/Chat/Chat.tsx +++ b/src/components/Chat/Chat.tsx @@ -1,13 +1,15 @@ -import { Box } from 'ink'; -import { useCallback, useEffect, useState } from 'react'; +import { Box, Text } from 'ink'; +import { useCallback, useEffect, useRef, useState } from 'react'; import { DECISION, MODE, PROMPT, ROLE } from '../../constants'; import { agents, ollama, tools } from '../../utils'; import { Messages } from '../Messages'; +import { TURN_ABORTED_MESSAGE } from '../Messages/constants'; import { PlanApproval } from '../PlanApproval'; import { ToolApproval } from '../ToolApproval'; import { ACTION_NOT_PERFORMED, + INTERRUPT_REASON, PLAN_CHECKLIST_REMINDER, PLAN_EXECUTION_REMINDER, } from './constants'; @@ -32,7 +34,9 @@ export function Chat({ const [messages, setMessages] = useState([]); const [streamingMessage, setStreamingMessage] = useState(null); + const [isLoading, setIsLoading] = useState(false); + const [pendingToolCall, setPendingToolCall] = useState(null); const [pendingPlan, setPendingPlan] = useState<{ @@ -40,12 +44,17 @@ export function Chat({ messages: ollama.Message[]; } | null>(null); + const [interruptReason, setInterruptReason] = + useState(null); + const abortControllerRef = useRef(null); + useEffect(() => { setMessages([]); setStreamingMessage(null); setIsLoading(false); setPendingToolCall(null); setPendingPlan(null); + setInterruptReason(null); }, [sessionId]); const buildToolResultMessage = useCallback( @@ -84,11 +93,25 @@ export function Chat({ [], ); + const handleInterrupt = useCallback(() => { + abortControllerRef.current?.abort(); + abortControllerRef.current = null; + setIsLoading(false); + setStreamingMessage(null); + setInterruptReason(INTERRUPT_REASON.INTERRUPTED); + setMessages((prev) => [ + ...prev, + { role: ROLE.USER, content: TURN_ABORTED_MESSAGE }, + ]); + }, []); + const processStream = useCallback( async ( currentMessages: ollama.Message[], executionMode: MODE.Name = mode, ) => { + const controller = new AbortController(); + abortControllerRef.current = controller; const assistantMessage: ollama.Message = { role: ROLE.ASSISTANT, content: '', @@ -129,7 +152,12 @@ export function Chat({ agents.withSystemMessage(currentMessages), model, tools.TOOLS, + controller.signal, )) { + // v8 ignore next 3 + if (controller.signal.aborted) { + return; + } if (chunk.type === 'content') { assistantMessage.content += chunk.content; setStreamingMessage({ ...assistantMessage }); @@ -182,9 +210,14 @@ export function Chat({ commitAssistantMessage(); } catch (error) { // v8 ignore next - assistantMessage.content = `Error: ${error instanceof Error ? error.message : String(error)}`; - commitAssistantMessage(); + if (!controller.signal.aborted) { + assistantMessage.content = `Error: ${error instanceof Error ? error.message : String(error)}`; + commitAssistantMessage(); + } } finally { + if (abortControllerRef.current === controller) { + abortControllerRef.current = null; + } setIsLoading(false); } }, @@ -194,10 +227,14 @@ export function Chat({ // Process stream with only read-only tools (for plan mode research phase) const processStreamReadOnly = useCallback( async (currentMessages: ollama.Message[]) => { + const controller = new AbortController(); + abortControllerRef.current = controller; + const assistantMessage: ollama.Message = { role: ROLE.ASSISTANT, content: '', }; + let committedMessages = currentMessages; let assistantCommitted = false; @@ -239,7 +276,12 @@ export function Chat({ agents.withSystemMessage(currentMessages), model, readOnlyTools, + controller.signal, )) { + // v8 ignore next 3 + if (controller.signal.aborted) { + return; + } if (chunk.type === 'content') { assistantMessage.content += chunk.content; setStreamingMessage({ ...assistantMessage }); @@ -310,7 +352,12 @@ export function Chat({ agents.withSystemMessage(planMessages), model, [], // No tools during plan generation output + controller.signal, )) { + // v8 ignore next 3 + if (controller.signal.aborted) { + return; + } if (chunk.type === 'content') { planAssistantMessage.content += chunk.content; setStreamingMessage({ ...planAssistantMessage }); @@ -346,9 +393,14 @@ export function Chat({ setIsLoading(false); } catch (error) { // v8 ignore next - assistantMessage.content = `Error: ${error instanceof Error ? error.message : String(error)}`; - commitAssistantMessage(); + if (!controller.signal.aborted) { + assistantMessage.content = `Error: ${error instanceof Error ? error.message : String(error)}`; + commitAssistantMessage(); + } } finally { + if (abortControllerRef.current === controller) { + abortControllerRef.current = null; + } setIsLoading(false); } }, @@ -430,18 +482,12 @@ export function Chat({ } case DECISION.REJECT: { - const rejectionMessage: ollama.Message = { - role: ROLE.SYSTEM, - content: `User declined to execute tool ${toolCall.function.name}`, - }; - - const newMessages = [...messages, rejectionMessage]; setMessages((previousMessages) => [ ...previousMessages, - rejectionMessage, + { role: ROLE.USER, content: TURN_ABORTED_MESSAGE }, ]); - - await processStream(newMessages); + setIsLoading(false); + setInterruptReason(INTERRUPT_REASON.REJECTED); break; } } @@ -451,7 +497,9 @@ export function Chat({ const handleSubmit = useCallback( async (value: string) => { + setInterruptReason(null); const userContent = value.trim(); + if (!userContent) { return; } @@ -504,9 +552,23 @@ export function Chat({ /> )} + {interruptReason && !isLoading && ( + + + {interruptReason === INTERRUPT_REASON.REJECTED + ? '❗ Tool call rejected.' + : '❗ Execution interrupted.'} + + + )} + {!pendingPlan && !pendingToolCall && ( - // eslint-disable-next-line @typescript-eslint/no-misused-promises - + )} ); diff --git a/src/components/Chat/Input.test.tsx b/src/components/Chat/Input.test.tsx index 22a9e4de..4cf2ebcf 100644 --- a/src/components/Chat/Input.test.tsx +++ b/src/components/Chat/Input.test.tsx @@ -480,6 +480,36 @@ describe('Input', () => { expect(onSubmit).not.toHaveBeenCalled(); }); + it('calls onInterrupt on Ctrl+C when disabled', async () => { + const onInterrupt = vi.fn(); + const { stdin } = render( + , + ); + stdin.write(KEY.CTRL_C); + await time.tick(); + expect(onInterrupt).toHaveBeenCalledOnce(); + }); + + it('calls onInterrupt on Esc when disabled', async () => { + const onInterrupt = vi.fn(); + const { stdin } = render( + , + ); + stdin.write(KEY.ESCAPE); + await time.tick(20); + expect(onInterrupt).toHaveBeenCalledOnce(); + }); + + it('does not call onInterrupt when not disabled', async () => { + const onInterrupt = vi.fn(); + const { stdin } = render( + , + ); + stdin.write(KEY.ESCAPE); + await time.tick(); + expect(onInterrupt).not.toHaveBeenCalled(); + }); + it('ignores file suggestion interactions when disabled', async () => { const onSubmit = vi.fn(); const { lastFrame, stdin } = render( diff --git a/src/components/Chat/Input.tsx b/src/components/Chat/Input.tsx index b456b7bf..fd10c9ad 100644 --- a/src/components/Chat/Input.tsx +++ b/src/components/Chat/Input.tsx @@ -9,6 +9,7 @@ import { FileSuggestions } from './FileSuggestions'; interface Props { isDisabled?: boolean; + onInterrupt?: () => void; onSubmit: (value: string) => void; } @@ -17,7 +18,7 @@ function hasFileSuggestionQuery(input: string): boolean { return /(^|\s)@\S+$/.test(input); } -export function Input({ isDisabled = false, onSubmit }: Props) { +export function Input({ isDisabled = false, onInterrupt, onSubmit }: Props) { const { exit } = useApp(); const [input, setInput] = useState(''); const [inputKey, setInputKey] = useState(0); @@ -88,13 +89,23 @@ export function Input({ isDisabled = false, onSubmit }: Props) { ); useInput((_input, key) => { - if (key.ctrl && _input === 'c') { + const isCtrlC = key.ctrl && _input === 'c'; + + if (isDisabled) { + if (key.escape || isCtrlC) { + onInterrupt?.(); + } + return; + } + + if (isCtrlC) { if (input) { setInput(''); remountTextInput(); - } else { - exit(); + return; } + + exit(); } }); diff --git a/src/components/Chat/constants.ts b/src/components/Chat/constants.ts index b0acc115..98604fd1 100644 --- a/src/components/Chat/constants.ts +++ b/src/components/Chat/constants.ts @@ -5,3 +5,8 @@ export const PLAN_CHECKLIST_REMINDER = export const PLAN_EXECUTION_REMINDER = 'Do not claim success and do not call write_file or run_shell until the user approves execution'; + +export enum INTERRUPT_REASON { + INTERRUPTED = 'interrupted', + REJECTED = 'rejected', +} diff --git a/src/components/Messages.test.tsx b/src/components/Messages/Messages.test.tsx similarity index 84% rename from src/components/Messages.test.tsx rename to src/components/Messages/Messages.test.tsx index 6c849b79..870374a4 100644 --- a/src/components/Messages.test.tsx +++ b/src/components/Messages/Messages.test.tsx @@ -1,7 +1,8 @@ import { Text } from 'ink'; import { render } from 'ink-testing-library'; -import { ROLE, UI } from '../constants'; +import { ROLE, UI } from '../../constants'; +import { TURN_ABORTED_MESSAGE } from './constants'; import { Messages } from './Messages'; vi.mock('@inkjs/ui', () => ({ @@ -70,13 +71,22 @@ describe('Messages', () => { const unknownMessage = { role: 'unknown', content: 'test', - } as unknown as import('../utils/ollama').Message; + } as unknown as import('../../utils/ollama').Message; const { lastFrame } = render( , ); expect(lastFrame()).toContain('test'); }); + it('does not render turn_aborted messages', () => { + const abortedMessage = { role: ROLE.USER, content: TURN_ABORTED_MESSAGE }; + const { lastFrame } = render( + , + ); + expect(lastFrame()).toContain('hello'); + expect(lastFrame()).not.toContain('turn_aborted'); + }); + it('renders a streaming message after committed messages', () => { const { lastFrame } = render( - {message.role === ROLE.USER ? UI.PROMPT_PREFIX : ''} + {message.role === ROLE.USER && UI.PROMPT_PREFIX} {message.content} @@ -45,14 +46,13 @@ const MessageRow = memo(function MessageRow({ message }: MessageRowProps) { export function Messages({ messages, isLoading, streamingMessage }: Props) { return ( - {messages.map((message, index) => ( - - ))} + {messages + .filter(({ content }) => content !== TURN_ABORTED_MESSAGE) + .map((message, index) => ( + + ))} - {streamingMessage && } + {streamingMessage && } {isLoading && !streamingMessage?.content && ( diff --git a/src/components/Messages/constants.ts b/src/components/Messages/constants.ts new file mode 100644 index 00000000..99e27bd7 --- /dev/null +++ b/src/components/Messages/constants.ts @@ -0,0 +1,5 @@ +export const TURN_ABORTED_MESSAGE = [ + '', + 'The user interrupted the previous turn on purpose. Any running commands may still be running in the background. If any tools were aborted, they may have partially executed.', + '', +].join('\n'); diff --git a/src/components/Messages/index.ts b/src/components/Messages/index.ts new file mode 100644 index 00000000..67ba5822 --- /dev/null +++ b/src/components/Messages/index.ts @@ -0,0 +1 @@ +export { Messages } from './Messages'; diff --git a/src/constants/ui.ts b/src/constants/ui.ts index 51e3df6c..6a7bff4e 100644 --- a/src/constants/ui.ts +++ b/src/constants/ui.ts @@ -1,2 +1,2 @@ -export const HEADER_PREFIX = '🦙'; +export const HEADER_PREFIX = '🦙 '; export const PROMPT_PREFIX = '> '; diff --git a/src/utils/ollama.ts b/src/utils/ollama.ts index 71a3f9df..fc40e6fd 100644 --- a/src/utils/ollama.ts +++ b/src/utils/ollama.ts @@ -31,21 +31,43 @@ export async function* streamChat( messages: Message[], model: string = DEFAULT_MODEL, tools?: Tool[], + signal?: AbortSignal, ): AsyncGenerator { const response = await client.chat({ model, messages, stream: true, tools, + // v8 ignore next + ...(signal ? { signal } : {}), }); - for await (const chunk of response) { - if (chunk.message.content) { - yield { type: 'content', content: chunk.message.content }; + try { + for await (const chunk of response) { + // v8 ignore next 3 + if (signal?.aborted) { + return; + } + + if (chunk.message.content) { + yield { type: 'content', content: chunk.message.content }; + } + + if (chunk.message.tool_calls) { + yield { type: 'tool_calls', tool_calls: chunk.message.tool_calls }; + } } - if (chunk.message.tool_calls) { - yield { type: 'tool_calls', tool_calls: chunk.message.tool_calls }; + } catch (error) { + // v8 ignore start + if ( + error instanceof Error && + (error.name === 'AbortError' || signal?.aborted) + ) { + return; } + + throw error; + // v8 ignore stop } } diff --git a/vite.config.mts b/vite.config.mts index 47505985..db27cc74 100644 --- a/vite.config.mts +++ b/vite.config.mts @@ -10,15 +10,6 @@ export default defineConfig({ output: { entryFileNames: 'cli.js', }, - external: [ - 'cac', - 'node:child_process', - 'node:fs', - 'node:os', - 'node:path', - 'node:util', - 'ollama', - ], }, },