Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
13 kB · 393 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394import { setKeybindings } from "@earendil-works/pi-tui";import { afterEach, beforeEach, describe, expect, it, type Mock, vi } from "vitest";import { KeybindingsManager } from "../src/core/keybindings.js";import { InteractiveMode } from "../src/modes/interactive/interactive-mode.js";
type FakeEditor = { text: string; getText: () => string; getExpandedText: () => string; setText: (text: string) => void;};
type FakeInteractiveMode = { ctrlCExitHintExpiresAt: number; ctrlCExitHintTimer: ReturnType<typeof setTimeout> | undefined; escapeRepeatAction: "tree" | "clear" | undefined; escapeRepeatExpiresAt: number; escapeRepeatTimer: ReturnType<typeof setTimeout> | undefined; traceUploadAllAbortController: AbortController | undefined; isShuttingDown: boolean; editor: FakeEditor; connectionState: { isStreaming: boolean; isCompacting: boolean; isBashRunning: boolean; retryAttempt: number; sessionActions: { queuedCount: number; steering: readonly string[]; followUps: readonly string[] }; }; agentConnection: { abort: Mock; clearQueue: Mock; abortAndClearQueue: Mock; abortRetry: Mock; abortCompaction: Mock; abortBranchSummary: Mock; abortBash: Mock; }; subagentSummaryLine: { invalidate: Mock }; ui: { requestRender: Mock; onDebug?: () => void }; updatePendingMessagesDisplay: Mock; showError: Mock; showTreeSelector: Mock; shutdown: Mock; updateEditorBorderColor: Mock; queueSelection: { isBrowsing: boolean; reset: () => string }; defaultEditor?: { onAction: Mock; getHeaderLine?: () => string | undefined; onEscape?: () => void; onCtrlD?: () => void; onPasteImage?: () => void; onMoveBelowPrompt?: () => boolean; onChange?: (text: string) => void; }; keybindings?: KeybindingsManager; handleDebugCommand?: Mock; showShortcutGuide?: Mock;};
function createEditor(text = ""): FakeEditor { const editor: FakeEditor = { text, getText() { return this.text; }, getExpandedText() { return this.text; }, setText(nextText: string) { this.text = nextText; }, }; return editor;}
function createInteractiveFake(options: { editorText?: string; streaming?: boolean; compacting?: boolean; bashRunning?: boolean; retryAttempt?: number;}): FakeInteractiveMode { const editor = createEditor(options.editorText ?? ""); const fake: FakeInteractiveMode = { ctrlCExitHintExpiresAt: 0, ctrlCExitHintTimer: undefined, escapeRepeatAction: undefined, escapeRepeatExpiresAt: 0, escapeRepeatTimer: undefined, traceUploadAllAbortController: undefined, isShuttingDown: false, editor, connectionState: { isStreaming: options.streaming ?? false, isCompacting: options.compacting ?? false, isBashRunning: options.bashRunning ?? false, retryAttempt: options.retryAttempt ?? 0, sessionActions: { queuedCount: 0, steering: [], followUps: [] }, }, agentConnection: { abort: vi.fn().mockResolvedValue(undefined), clearQueue: vi.fn().mockResolvedValue({ steering: [], followUp: [] }), abortAndClearQueue: vi.fn().mockResolvedValue({ steering: [], followUp: [] }), abortRetry: vi.fn(), abortCompaction: vi.fn(), abortBranchSummary: vi.fn(), abortBash: vi.fn(), }, subagentSummaryLine: { invalidate: vi.fn() }, ui: { requestRender: vi.fn() }, queueSelection: { isBrowsing: false, reset: () => "" }, updatePendingMessagesDisplay: vi.fn(), showError: vi.fn(), showTreeSelector: vi.fn(), shutdown: vi.fn().mockResolvedValue(undefined), updateEditorBorderColor: vi.fn(), showShortcutGuide: vi.fn(), }; Object.setPrototypeOf(fake, InteractiveMode.prototype); return fake;}
describe("InteractiveMode interrupt shortcuts", () => { beforeEach(() => { setKeybindings(new KeybindingsManager()); vi.useFakeTimers(); vi.setSystemTime(new Date("2026-06-08T12:00:00Z")); });
afterEach(() => { vi.useRealTimers(); });
it("interrupts streaming and shows the exit hint on first Ctrl+C", () => { const mode = createInteractiveFake({ streaming: true });
Reflect.get(InteractiveMode.prototype, "handleCtrlC").call(mode);
expect(mode.agentConnection.abort).toHaveBeenCalledTimes(1); expect(mode.shutdown).not.toHaveBeenCalled(); expect(Reflect.get(InteractiveMode.prototype, "getTrayOverrideLabel").call(mode)).toBe( "Press Ctrl+C again to exit", ); });
it("interrupts bash and streaming on the same Ctrl+C", () => { const mode = createInteractiveFake({ streaming: true, bashRunning: true });
Reflect.get(InteractiveMode.prototype, "handleCtrlC").call(mode);
expect(mode.agentConnection.abortBash).toHaveBeenCalledTimes(1); expect(mode.agentConnection.abort).toHaveBeenCalledTimes(1); expect(mode.shutdown).not.toHaveBeenCalled(); });
it.each([ ["Ctrl+C", "handleCtrlC"], ["Escape", "handleEscape"], ] as const)("cancels an upload-all operation on %s", (_label, handlerName) => { const mode = createInteractiveFake({}); const controller = new AbortController(); mode.traceUploadAllAbortController = controller;
Reflect.get(InteractiveMode.prototype, handlerName).call(mode);
expect(controller.signal.aborted).toBe(true); expect(controller.signal.reason).toEqual(new Error("Trace upload cancelled")); expect(mode.shutdown).not.toHaveBeenCalled(); });
it("preserves the queue and the draft when interrupting streaming", () => { const mode = createInteractiveFake({ editorText: "draft", streaming: true }); mode.connectionState.sessionActions = { queuedCount: 2, steering: ["steer"], followUps: ["follow"] };
Reflect.get(InteractiveMode.prototype, "handleCtrlC").call(mode);
expect(mode.agentConnection.abort).toHaveBeenCalledTimes(1); expect(mode.agentConnection.abortAndClearQueue).not.toHaveBeenCalled(); expect(mode.agentConnection.clearQueue).not.toHaveBeenCalled(); expect(mode.editor.getText()).toBe("draft"); expect(mode.connectionState.sessionActions).toEqual({ queuedCount: 2, steering: ["steer"], followUps: ["follow"], }); });
it("exits on the second Ctrl+C while the hint is visible", () => { const mode = createInteractiveFake({ streaming: true }); const handleCtrlC = Reflect.get(InteractiveMode.prototype, "handleCtrlC");
handleCtrlC.call(mode); handleCtrlC.call(mode);
expect(mode.agentConnection.abort).toHaveBeenCalledTimes(1); expect(mode.shutdown).toHaveBeenCalledTimes(1); });
it("clears the exit hint after two seconds", async () => { const mode = createInteractiveFake({ editorText: "draft" });
Reflect.get(InteractiveMode.prototype, "handleCtrlC").call(mode); expect(mode.editor.getText()).toBe("draft"); expect(Reflect.get(InteractiveMode.prototype, "getTrayOverrideLabel").call(mode)).toBe( "Press Ctrl+C again to exit", );
await vi.advanceTimersByTimeAsync(2000);
expect(Reflect.get(InteractiveMode.prototype, "getTrayOverrideLabel").call(mode)).toBeUndefined(); expect(mode.subagentSummaryLine.invalidate).toHaveBeenCalled(); expect(mode.ui.requestRender).toHaveBeenCalled(); });
it("preserves idle draft input on first Ctrl+C", () => { const mode = createInteractiveFake({ editorText: "draft" });
Reflect.get(InteractiveMode.prototype, "handleCtrlC").call(mode);
expect(mode.editor.getText()).toBe("draft"); expect(mode.agentConnection.abort).not.toHaveBeenCalled(); expect(mode.shutdown).not.toHaveBeenCalled(); });
it("cancels the tree repeat when typing after interrupting streaming", () => { const actionHandlers = new Map<string, () => void>(); const mode = createInteractiveFake({ editorText: "draft", streaming: true }); const defaultEditor: NonNullable<FakeInteractiveMode["defaultEditor"]> = { onAction: vi.fn((action: string, handler: () => void) => { actionHandlers.set(action, handler); }), }; Object.assign(mode, { defaultEditor, keybindings: new KeybindingsManager(), handleDebugCommand: vi.fn(), });
Reflect.get(InteractiveMode.prototype, "setupKeyHandlers").call(mode); expect(defaultEditor.onEscape).toBeDefined(); defaultEditor.onEscape?.(); expect(mode.agentConnection.abort).toHaveBeenCalledTimes(1); expect(mode.editor.getText()).toBe("draft"); mode.editor.setText("queued draft"); defaultEditor.onChange?.("queued draft");
defaultEditor.onEscape?.();
expect(mode.showTreeSelector).not.toHaveBeenCalled(); expect(mode.shutdown).not.toHaveBeenCalled(); });
it("preserves the tree repeat while browsing into a queued message", () => { const mode = createInteractiveFake({}); const defaultEditor: NonNullable<FakeInteractiveMode["defaultEditor"]> = { onAction: vi.fn(), }; Object.assign(mode, { defaultEditor, keybindings: new KeybindingsManager(), handleDebugCommand: vi.fn(), isApplyingQueueSelectionText: false, }); Reflect.get(InteractiveMode.prototype, "setupKeyHandlers").call(mode); const setText = mode.editor.setText.bind(mode.editor); mode.editor.setText = (text) => { setText(text); defaultEditor.onChange?.(text); }; mode.escapeRepeatAction = "tree"; mode.escapeRepeatExpiresAt = Date.now() + 500;
Reflect.get(InteractiveMode.prototype, "setEditorTextFromQueueSelection").call(mode, "queued item");
expect(mode.editor.getText()).toBe("queued item"); expect(mode.escapeRepeatAction).toBe("tree"); });
it("clears an idle draft on double Escape", () => { const actionHandlers = new Map<string, () => void>(); const mode = createInteractiveFake({ editorText: "draft" }); const defaultEditor: NonNullable<FakeInteractiveMode["defaultEditor"]> = { onAction: vi.fn((action: string, handler: () => void) => { actionHandlers.set(action, handler); }), }; Object.assign(mode, { defaultEditor, keybindings: new KeybindingsManager(), handleDebugCommand: vi.fn(), });
Reflect.get(InteractiveMode.prototype, "setupKeyHandlers").call(mode); defaultEditor.onEscape?.(); expect(mode.editor.getText()).toBe("draft");
defaultEditor.onEscape?.();
expect(mode.editor.getText()).toBe(""); expect(mode.agentConnection.abort).not.toHaveBeenCalled(); });
it("opens the tree on double Escape with an empty idle prompt", () => { const mode = createInteractiveFake({}); const handleEscape = Reflect.get(InteractiveMode.prototype, "handleEscape");
handleEscape.call(mode); handleEscape.call(mode);
expect(mode.showTreeSelector).toHaveBeenCalledTimes(1); expect(mode.editor.getText()).toBe(""); });
it("clears a whitespace draft on double Escape", () => { const mode = createInteractiveFake({ editorText: " " }); const handleEscape = Reflect.get(InteractiveMode.prototype, "handleEscape");
handleEscape.call(mode); handleEscape.call(mode);
expect(mode.showTreeSelector).not.toHaveBeenCalled(); expect(mode.editor.getText()).toBe(""); });
for (const [label, options] of [ ["a retry", { retryAttempt: 1 }], ["compaction", { compacting: true }], ["a bash command", { bashRunning: true }], ] as const) { it(`opens the tree after cancelling ${label} without clearing the draft`, () => { const mode = createInteractiveFake({ editorText: "draft", ...options }); const handleEscape = Reflect.get(InteractiveMode.prototype, "handleEscape");
handleEscape.call(mode); handleEscape.call(mode);
expect(mode.showTreeSelector).toHaveBeenCalledTimes(1); expect(mode.editor.getText()).toBe("draft"); }); }
it("clears the Escape repeat before a separate interrupt", () => { const mode = createInteractiveFake({}); mode.escapeRepeatAction = "tree"; mode.escapeRepeatExpiresAt = Date.now() + 500;
Reflect.get(InteractiveMode.prototype, "handleInterruptKey").call(mode);
expect(mode.escapeRepeatAction).toBeUndefined(); });
it("expires the Escape repeat window", async () => { const actionHandlers = new Map<string, () => void>(); const mode = createInteractiveFake({ editorText: "draft" }); const defaultEditor: NonNullable<FakeInteractiveMode["defaultEditor"]> = { onAction: vi.fn((action: string, handler: () => void) => { actionHandlers.set(action, handler); }), }; Object.assign(mode, { defaultEditor, keybindings: new KeybindingsManager(), handleDebugCommand: vi.fn(), });
Reflect.get(InteractiveMode.prototype, "setupKeyHandlers").call(mode); defaultEditor.onEscape?.(); await vi.advanceTimersByTimeAsync(500); defaultEditor.onEscape?.();
expect(mode.editor.getText()).toBe("draft"); });
it("opens keyboard shortcuts from the configured app action", () => { const actionHandlers = new Map<string, () => void>(); const mode = createInteractiveFake({}); const defaultEditor: NonNullable<FakeInteractiveMode["defaultEditor"]> = { onAction: vi.fn((action: string, handler: () => void) => { actionHandlers.set(action, handler); }), }; Object.assign(mode, { defaultEditor, keybindings: new KeybindingsManager(), handleDebugCommand: vi.fn(), });
Reflect.get(InteractiveMode.prototype, "setupKeyHandlers").call(mode); actionHandlers.get("app.shortcuts")?.();
expect(mode.showShortcutGuide).toHaveBeenCalledTimes(1); });});