oh-my-pi/packages/coding-agent/src/modes/components/custom-editor.ts

import { Editor, type KeyId, matchesKey, parseKittySequence } from "@oh-my-pi/pi-tui";
import type { AppKeybinding } from "../../config/keybindings";
import { highlightMagicKeywords } from "../magic-keywords";

type ConfigurableEditorAction = Extract<
	AppKeybinding,
	| "app.interrupt"
	| "app.clear"
	| "app.exit"
	| "app.suspend"
	| "app.thinking.cycle"
	| "app.model.cycleForward"
	| "app.model.cycleBackward"
	| "app.model.select"
	| "app.model.selectTemporary"
	| "app.tools.expand"
	| "app.thinking.toggle"
	| "app.editor.external"
	| "app.history.search"
	| "app.message.dequeue"
	| "app.clipboard.pasteImage"
	| "app.clipboard.copyPrompt"
>;

const DEFAULT_ACTION_KEYS: Record<ConfigurableEditorAction, KeyId[]> = {
	"app.interrupt": ["escape"],
	"app.clear": ["ctrl+c"],
	"app.exit": ["ctrl+d"],
	"app.suspend": ["ctrl+z"],
	"app.thinking.cycle": ["shift+tab"],
	"app.model.cycleForward": ["ctrl+p"],
	"app.model.cycleBackward": ["shift+ctrl+p"],
	"app.model.select": ["ctrl+l"],
	"app.model.selectTemporary": ["alt+p"],
	"app.tools.expand": ["ctrl+o"],
	"app.thinking.toggle": ["ctrl+t"],
	"app.editor.external": ["ctrl+g"],
	"app.history.search": ["ctrl+r"],
	"app.message.dequeue": ["alt+up"],
	"app.clipboard.pasteImage": ["ctrl+v"],
	"app.clipboard.copyPrompt": ["alt+shift+c"],
};

/**
 * Custom editor that handles configurable app-level shortcuts for coding-agent.
 */
export class CustomEditor extends Editor {
	/** Gradient-highlight the "ultrathink" / "orchestrate" / "workflow" keywords as the user types
	 *  them, skipping any occurrence inside code spans, fenced blocks, or XML sections. */
	decorateText = (text: string): string => highlightMagicKeywords(text);
	onEscape?: () => void;
	onClear?: () => void;
	onExit?: () => void;
	onCycleThinkingLevel?: () => void;
	onCycleModelForward?: () => void;
	onCycleModelBackward?: () => void;
	onSelectModel?: () => void;
	onExpandTools?: () => void;
	onToggleThinking?: () => void;
	onExternalEditor?: () => void;
	onHistorySearch?: () => void;
	onSuspend?: () => void;
	onSelectModelTemporary?: () => void;
	/** Called when the configured copy-prompt shortcut is pressed. */
	onCopyPrompt?: () => void;
	/** Called when the configured image-paste shortcut is pressed. */
	onPasteImage?: () => Promise<boolean>;
	/** Called when the configured dequeue shortcut is pressed. */
	onDequeue?: () => void;
	/** Called when Caps Lock is pressed. */
	onCapsLock?: () => void;

	/** Custom key handlers from extensions and non-built-in app actions. */
	#customKeyHandlers = new Map<KeyId, () => void>();
	#actionKeys = new Map<ConfigurableEditorAction, KeyId[]>(
		Object.entries(DEFAULT_ACTION_KEYS).map(([action, keys]) => [action as ConfigurableEditorAction, [...keys]]),
	);

	setActionKeys(action: ConfigurableEditorAction, keys: KeyId[]): void {
		this.#actionKeys.set(action, [...keys]);
	}

	#matchesAction(data: string, action: ConfigurableEditorAction): boolean {
		const keys = this.#actionKeys.get(action);
		if (!keys) return false;
		for (const key of keys) {
			if (matchesKey(data, key)) return true;
		}
		return false;
	}

	/**
	 * Register a custom key handler. Extensions use this for shortcuts.
	 */
	setCustomKeyHandler(key: KeyId, handler: () => void): void {
		this.#customKeyHandlers.set(key, handler);
	}

	/**
	 * Remove a custom key handler.
	 */
	removeCustomKeyHandler(key: KeyId): void {
		this.#customKeyHandlers.delete(key);
	}

	/**
	 * Clear all custom key handlers.
	 */
	clearCustomKeyHandlers(): void {
		this.#customKeyHandlers.clear();
	}

	handleInput(data: string): void {
		const parsed = parseKittySequence(data);
		if (parsed && (parsed.modifier & 64) !== 0 && this.onCapsLock) {
			// Caps Lock is modifier bit 64
			this.onCapsLock();
			return;
		}

		// Intercept configured image paste (async - fires and handles result)
		if (this.#matchesAction(data, "app.clipboard.pasteImage") && this.onPasteImage) {
			void this.onPasteImage();
			return;
		}

		// Intercept configured external editor shortcut
		if (this.#matchesAction(data, "app.editor.external") && this.onExternalEditor) {
			this.onExternalEditor();
			return;
		}

		// Intercept configured temporary model selector shortcut
		if (this.#matchesAction(data, "app.model.selectTemporary") && this.onSelectModelTemporary) {
			this.onSelectModelTemporary();
			return;
		}

		// Intercept configured suspend shortcut
		if (this.#matchesAction(data, "app.suspend") && this.onSuspend) {
			this.onSuspend();
			return;
		}

		// Intercept configured thinking block visibility toggle
		if (this.#matchesAction(data, "app.thinking.toggle") && this.onToggleThinking) {
			this.onToggleThinking();
			return;
		}

		// Intercept configured model selector shortcut
		if (this.#matchesAction(data, "app.model.select") && this.onSelectModel) {
			this.onSelectModel();
			return;
		}

		// Intercept configured history search shortcut
		if (this.#matchesAction(data, "app.history.search") && this.onHistorySearch) {
			this.onHistorySearch();
			return;
		}

		// Intercept configured tool output expansion shortcut
		if (this.#matchesAction(data, "app.tools.expand") && this.onExpandTools) {
			this.onExpandTools();
			return;
		}

		// Intercept configured backward model cycling (check before forward cycling)
		if (this.#matchesAction(data, "app.model.cycleBackward") && this.onCycleModelBackward) {
			this.onCycleModelBackward();
			return;
		}

		// Intercept configured forward model cycling
		if (this.#matchesAction(data, "app.model.cycleForward") && this.onCycleModelForward) {
			this.onCycleModelForward();
			return;
		}

		// Intercept configured thinking level cycling
		if (this.#matchesAction(data, "app.thinking.cycle") && this.onCycleThinkingLevel) {
			this.onCycleThinkingLevel();
			return;
		}

		// Intercept configured interrupt shortcut.
		// When the autocomplete popup is visible, ESC's first job is to dismiss
		// the popup — let super.handleInput() route it to #cancelAutocomplete().
		// The user can press ESC again afterward to fire the global interrupt
		// handler. This matches the standard TUI/IDE pattern and prevents a
		// single ESC from both closing an @ completion and aborting an active
		// agent run (#1655).
		if (this.#matchesAction(data, "app.interrupt") && this.onEscape && !this.isShowingAutocomplete()) {
			this.onEscape();
			return;
		}

		// Intercept configured clear shortcut
		if (this.#matchesAction(data, "app.clear") && this.onClear) {
			this.onClear();
			return;
		}

		// Intercept configured exit shortcut. Always consume the shortcut so it
		// never reaches the parent handler; firing onExit is the controller's
		// chance to snapshot the current text as a draft before shutting down.
		if (this.#matchesAction(data, "app.exit")) {
			this.onExit?.();
			return;
		}

		// Intercept configured dequeue shortcut (restore queued message to editor)
		if (this.#matchesAction(data, "app.message.dequeue") && this.onDequeue) {
			this.onDequeue();
			return;
		}

		// Intercept configured copy-prompt shortcut
		if (this.#matchesAction(data, "app.clipboard.copyPrompt") && this.onCopyPrompt) {
			this.onCopyPrompt();
			return;
		}

		// Check custom key handlers (extensions)
		for (const [keyId, handler] of this.#customKeyHandlers) {
			if (matchesKey(data, keyId)) {
				handler();
				return;
			}
		}

		// Pass to parent for normal handling
		super.handleInput(data);
	}
}