pi skill dollar autocomplete extension

skill-dollar.ts

/**
 * pi-skill-dollar - `$` autocomplete shortcut for pi Agent Skills.
 */
import { type ExtensionAPI, type SlashCommandInfo } from "@earendil-works/pi-coding-agent";
// Use pi's bundled TUI so this patches the actual Editor class in use.
import { Editor, fuzzyFilter } from "/Users/maxktz/.nvm/versions/node/v23.7.0/lib/node_modules/@earendil-works/pi-coding-agent/node_modules/@earendil-works/pi-tui/dist/index.js";
import type { AutocompleteItem, AutocompleteProvider, AutocompleteSuggestions } from "/Users/maxktz/.nvm/versions/node/v23.7.0/lib/node_modules/@earendil-works/pi-coding-agent/node_modules/@earendil-works/pi-tui/dist/index.js";

const DIRECT_DOLLAR_SKILL_RE = /^\$([a-z0-9][a-z0-9-]{0,63})(?:\s+(.*))?\s*$/;
const PATCH_VERSION_SYMBOL = Symbol.for("pi-skill-dollar.editor-patch-version");
const PATCH_VERSION = 3;

export type DollarToken = {
  query: string;
  typedPrefix: string;
  completionPrefix: string;
};

/**
 * Extract a `$skill-fragment` token immediately before the cursor.
 *
 * Valid triggers are at prompt start or after horizontal whitespace. Escaped
 * dollars (`\$`) and embedded dollars (`foo$bar`) are intentionally ignored.
 */
export function extractDollarToken(beforeCursor: string): DollarToken | null {
  const match = beforeCursor.match(/(^|[ \t])\$([a-z0-9-]*)$/);
  if (!match) return null;

  const query = match[2] ?? "";
  const typedPrefix = `$${query}`;
  const tokenStart = beforeCursor.length - typedPrefix.length;
  const beforeToken = beforeCursor.slice(0, tokenStart);

  // At prompt start, use a slash-shaped prefix of the same length so the stock
  // editor treats selection as slash-command completion: Enter both selects and
  // submits. In the middle of text, keep a dollar-shaped prefix so selection
  // only inserts the command text.
  const completionPrefix = beforeToken.length === 0 ? `/${query}` : typedPrefix;

  return { query, typedPrefix, completionPrefix };
}

export function skillDisplayName(commandName: string): string {
  return commandName.replace(/^skill:/, "");
}

function skillQueryText(command: SlashCommandInfo): string {
  const name = skillDisplayName(command.name);
  return `${name} ${command.name} ${command.description ?? ""}`;
}

function skillToItem(command: SlashCommandInfo): AutocompleteItem {
  const name = skillDisplayName(command.name);
  return {
    value: command.name,
    label: `$${name}`,
    ...(command.description ? { description: command.description } : {}),
  };
}

export function buildSkillCompletion(
  lines: string[],
  cursorLine: number,
  cursorCol: number,
  item: AutocompleteItem,
  prefix: string,
): { lines: string[]; cursorLine: number; cursorCol: number } {
  const currentLine = lines[cursorLine] ?? "";
  const beforePrefix = currentLine.slice(0, cursorCol - prefix.length);
  const afterCursor = currentLine.slice(cursorCol);
  const command = item.value.startsWith("/") ? item.value : `/${item.value}`;
  const replacement = `${command} `;
  const adjustedAfterCursor = afterCursor.startsWith(" ") ? afterCursor.slice(1) : afterCursor;
  const newLines = [...lines];
  newLines[cursorLine] = `${beforePrefix}${replacement}${adjustedAfterCursor}`;
  return {
    lines: newLines,
    cursorLine,
    cursorCol: beforePrefix.length + replacement.length,
  };
}

function createDollarSkillProvider(pi: ExtensionAPI, current: AutocompleteProvider): AutocompleteProvider {
  return {
    async getSuggestions(lines, cursorLine, cursorCol, options): Promise<AutocompleteSuggestions | null> {
      const line = lines[cursorLine] ?? "";
      const beforeCursor = line.slice(0, cursorCol);
      const token = extractDollarToken(beforeCursor);
      if (!token) {
        return current.getSuggestions(lines, cursorLine, cursorCol, options);
      }

      const skills = pi.getCommands().filter((c): c is SlashCommandInfo => c.source === "skill");
      if (skills.length === 0) return null;

      const filtered = token.query ? fuzzyFilter(skills, token.query, skillQueryText) : skills;
      const items = filtered.map(skillToItem);
      if (items.length === 0) return null;

      return {
        items,
        prefix: token.completionPrefix,
      };
    },

    applyCompletion(lines, cursorLine, cursorCol, item, prefix) {
      if (prefix.startsWith("$") || prefix.startsWith("/")) {
        return buildSkillCompletion(lines, cursorLine, cursorCol, item, prefix);
      }

      return current.applyCompletion(lines, cursorLine, cursorCol, item, prefix);
    },

    shouldTriggerFileCompletion(lines, cursorLine, cursorCol) {
      const line = lines[cursorLine] ?? "";
      const beforeCursor = line.slice(0, cursorCol);
      if (extractDollarToken(beforeCursor)) return true;
      return current.shouldTriggerFileCompletion?.(lines, cursorLine, cursorCol) ?? true;
    },
  };
}

type DollarEditorInternals = {
  getLines?: () => string[];
  getCursor?: () => { line: number; col: number };
  getText?: () => string;
  isShowingAutocomplete?: () => boolean;
  tryTriggerAutocomplete?: () => void;
  autocompleteState?: unknown;
  autocompleteList?: {
    getSelectedItem?: () => AutocompleteItem | null;
    filteredItems?: AutocompleteItem[];
  };
  autocompleteProvider?: AutocompleteProvider;
  state?: { lines: string[]; cursorLine: number };
  pushUndoSnapshot?: () => void;
  setCursorCol?: (col: number) => void;
  cancelAutocomplete?: () => void;
  onChange?: (text: string) => void;
  lastAction?: unknown;
};

function getEditorLinesAndCursor(editor: DollarEditorInternals): { lines: string[]; cursor: { line: number; col: number } } {
  const lines = editor.getLines?.() ?? editor.getText?.().split("\n") ?? [""];
  const cursor = editor.getCursor?.() ?? {
    line: lines.length - 1,
    col: lines[lines.length - 1]?.length ?? 0,
  };
  return { lines, cursor };
}

function maybeTriggerDollarAutocomplete(editor: unknown): void {
  const candidate = editor as DollarEditorInternals;
  if (!candidate.tryTriggerAutocomplete || candidate.isShowingAutocomplete?.()) return;

  const { lines, cursor } = getEditorLinesAndCursor(candidate);
  const line = lines[cursor.line] ?? "";
  const beforeCursor = line.slice(0, cursor.col);

  if (extractDollarToken(beforeCursor)) {
    candidate.tryTriggerAutocomplete();
  }
}

function findExactSkillItem(editor: DollarEditorInternals, query: string): AutocompleteItem | null {
  const selected = editor.autocompleteList?.getSelectedItem?.();
  if (selected && skillDisplayName(selected.value) === query) return selected;

  const exact = editor.autocompleteList?.filteredItems?.find((item) => skillDisplayName(item.value) === query);
  return exact ?? null;
}

function handleDollarAutocompleteSpace(editor: unknown): boolean {
  const candidate = editor as DollarEditorInternals;
  if (!candidate.autocompleteState || !candidate.autocompleteList || !candidate.autocompleteProvider) return false;

  const { lines, cursor } = getEditorLinesAndCursor(candidate);
  const line = lines[cursor.line] ?? "";
  const beforeCursor = line.slice(0, cursor.col);
  const token = extractDollarToken(beforeCursor);
  if (!token) return false;

  const item = token.query ? findExactSkillItem(candidate, token.query) : null;
  if (!item) {
    // In `$...` autocomplete, Space without an exact skill should simply leave
    // autocomplete mode. Cancel first so the built-in editor does not re-query
    // and fall through to file suggestions after inserting the space.
    candidate.cancelAutocomplete?.();
    return false;
  }

  candidate.pushUndoSnapshot?.();
  candidate.lastAction = null;
  const result = candidate.autocompleteProvider.applyCompletion(lines, cursor.line, cursor.col, item, token.completionPrefix);

  if (candidate.state) {
    candidate.state.lines = result.lines;
    candidate.state.cursorLine = result.cursorLine;
  }
  candidate.setCursorCol?.(result.cursorCol);
  candidate.cancelAutocomplete?.();
  candidate.onChange?.(candidate.getText?.() ?? result.lines.join("\n"));
  return true;
}

function patchEditorDollarTrigger(): void {
  const proto = Editor.prototype as unknown as {
    [key: symbol]: unknown;
    handleInput?: (data: string) => void;
    insertCharacter?: (char: string, skipUndoCoalescing?: boolean) => void;
  };

  if (proto[PATCH_VERSION_SYMBOL] === PATCH_VERSION) return;

  if (typeof proto.handleInput === "function") {
    const originalHandleInput = proto.handleInput;
    proto.handleInput = function patchedHandleInput(this: unknown, data: string): void {
      if (data === " " && handleDollarAutocompleteSpace(this)) {
        return;
      }
      originalHandleInput.call(this, data);
    };
  }

  if (typeof proto.insertCharacter === "function") {
    const originalInsertCharacter = proto.insertCharacter;
    proto.insertCharacter = function patchedInsertCharacter(this: unknown, char: string, skipUndoCoalescing?: boolean): void {
      originalInsertCharacter.call(this, char, skipUndoCoalescing);

      if (char === "$" || /^[a-zA-Z0-9_-]$/.test(char)) {
        maybeTriggerDollarAutocomplete(this);
      }
    };
  }

  proto[PATCH_VERSION_SYMBOL] = PATCH_VERSION;
}

export default function skillDollarExtension(pi: ExtensionAPI): void {
  patchEditorDollarTrigger();

  pi.on("session_start", async (_event, ctx) => {
    ctx.ui.addAutocompleteProvider((current) => createDollarSkillProvider(pi, current));
  });

  // Submit-time fallback. Autocomplete is the primary behavior, but this keeps
  // `$shisa-kb args` useful if a user submits without selecting.
  pi.on("input", async (event) => {
    if (event.source === "extension" || !event.text.trim()) {
      return { action: "continue" };
    }

    const text = event.text;

    if (text.startsWith("\\$")) {
      return { action: "transform", text: text.slice(1) };
    }

    const match = text.match(DIRECT_DOLLAR_SKILL_RE);
    if (match) {
      const skillName = match[1];
      const args = match[2] ? ` ${match[2].trim()}` : "";
      return { action: "transform", text: `/skill:${skillName}${args}` };
    }

    return { action: "continue" };
  });
}