pi-dsh-mimic 0.1.0

package/src/editor.ts

@@ -0,0 +1,133 @@
+import {
+  existsSync,
+  mkdirSync,
+  readFileSync,
+  readdirSync,
+  statSync,
+  writeFileSync,
+} from "node:fs";
+import { dirname, isAbsolute, join, relative } from "node:path";
+
+import { EDITOR_TRUNCATED_MESSAGE } from "./constants.js";
+
+export type EditorCommand = "view" | "create" | "str_replace" | "insert";
+
+export interface EditorArgs {
+  path: string;
+  command: EditorCommand;
+  view_range?: number[];
+  old_str?: string;
+  new_str?: string;
+  insert_line?: number;
+  file_text?: string;
+}
+
+const OUTPUT_LIMIT = 16_000;
+
+function requireAbsolutePath(candidate: string): string {
+  if (candidate.trim() === "" || !isAbsolute(candidate)) {
+    throw new Error(`Path is not an absolute path: ${candidate}`);
+  }
+  return candidate;
+}
+
+function requireArgument(value: string | undefined, argument: string, command: EditorCommand): string {
+  if (value === undefined) throw new Error(`${command} requires ${argument}`);
+  return value;
+}
+
+function clip(value: string, limit: number): string {
+  if (value.length <= limit) return value;
+  return `${value.slice(0, limit)}\n${EDITOR_TRUNCATED_MESSAGE}`;
+}
+
+function directoryView(root: string): string {
+  const rows: string[] = [root];
+  const walk = (directory: string, depth: number): void => {
+    if (depth > 2) return;
+    const children = readdirSync(directory, { withFileTypes: true })
+      .filter((item) => !item.name.startsWith(".")
+        && item.name !== "node_modules"
+        && item.name !== "__pycache__")
+      .sort((a, b) => a.name.localeCompare(b.name));
+    for (const child of children) {
+      const absolute = join(directory, child.name);
+      rows.push(`${relative(root, absolute)}${child.isDirectory() ? "/" : ""}`);
+      if (child.isDirectory()) walk(absolute, depth + 1);
+    }
+  };
+  walk(root, 1);
+  return rows.join("\n");
+}
+
+function fileView(path: string, range?: number[]): string {
+  const lines = readFileSync(path, "utf8").split("\n");
+  const start = range?.[0] ?? 1;
+  const requestedEnd = range?.[1] ?? -1;
+  const end = requestedEnd === -1 ? lines.length : requestedEnd;
+  if ((range !== undefined && range.length !== 2)
+    || !Number.isInteger(start)
+    || !Number.isInteger(end)
+    || start < 1
+    || end < start
+    || end > lines.length) {
+    throw new Error(`Invalid view_range for ${path}`);
+  }
+  const body = lines.slice(start - 1, end)
+    .map((line, offset) => `${String(start + offset).padStart(6, " ")}  ${line}`)
+    .join("\n");
+  return `Contents of ${path}:\n${body}`;
+}
+
+function replaceUnique(source: string, needle: string, replacement: string): string {
+  const match = source.indexOf(needle);
+  if (match === -1) throw new Error("old_str was not found");
+  if (source.indexOf(needle, match + needle.length) !== -1) {
+    throw new Error("Multiple matches found for old_str");
+  }
+  return source.slice(0, match) + replacement + source.slice(match + needle.length);
+}
+
+export function executeEditor(args: EditorArgs, outputLimit = OUTPUT_LIMIT): string {
+  const path = requireAbsolutePath(args.path);
+
+  if (args.command === "create") {
+    if (existsSync(path)) throw new Error(`Cannot create an existing path: ${path}`);
+    mkdirSync(dirname(path), { recursive: true });
+    writeFileSync(path, requireArgument(args.file_text, "file_text", args.command), {
+      encoding: "utf8",
+      flag: "wx",
+    });
+    return `Created successfully: ${path}`;
+  }
+
+  if (!existsSync(path)) throw new Error(`Path does not exist: ${path}`);
+  const info = statSync(path);
+
+  if (args.command === "view") {
+    const output = info.isDirectory() ? directoryView(path) : fileView(path, args.view_range);
+    return clip(output, outputLimit);
+  }
+  if (!info.isFile()) throw new Error(`Editing requires a regular file: ${path}`);
+
+  const source = readFileSync(path, "utf8");
+  if (args.command === "str_replace") {
+    const next = replaceUnique(
+      source,
+      requireArgument(args.old_str, "old_str", args.command),
+      args.new_str ?? "",
+    );
+    writeFileSync(path, next, "utf8");
+    return `Replaced text in ${path}`;
+  }
+
+  if (!Number.isInteger(args.insert_line)) throw new Error("insert requires insert_line");
+  const insertAfter = args.insert_line as number;
+  const lines = source.split("\n");
+  if (insertAfter < 0 || insertAfter > lines.length) {
+    throw new Error(`insert_line is outside the file: ${insertAfter}`);
+  }
+  lines.splice(insertAfter, 0, requireArgument(args.new_str, "new_str", args.command));
+  writeFileSync(path, lines.join("\n"), "utf8");
+  return `Inserted text in ${path}`;
+}

package/src/index.ts

@@ -0,0 +1,188 @@
+import type {
+  ExtensionAPI,
+  ExtensionContext,
+} from "@earendil-works/pi-coding-agent";
+import { StringEnum } from "@earendil-works/pi-ai";
+import { Type } from "typebox";
+
+import {
+  BOOTSTRAP_TOOL_NAMES,
+  EDITOR_DESCRIPTION,
+  MINIMAL_BASH_DESCRIPTION,
+  MINIMAL_PERSONA,
+  STATE_ENTRY_TYPE,
+  TARGET_MODEL_ID,
+  TARGET_PROVIDERS,
+} from "./constants.js";
+import { executeEditor, type EditorArgs } from "./editor.js";
+import {
+  isJsonObject,
+  keepMinimalPersona,
+  makeBootstrapRequest,
+} from "./protocol.js";
+import {
+  conversationHasStarted,
+  type LiveStage,
+  restoreStage,
+  supportedRoute,
+} from "./session-stage.js";
+
+const STATUS_ID = "pi-dsh-mimic";
+
+export default function piDshAnchor(api: ExtensionAPI): void {
+  const sessions = new Map<string, LiveStage>();
+  let editorAvailable = false;
+
+  const sessionKey = (context: ExtensionContext): string => context.sessionManager.getSessionId();
+
+  const exposeEditor = (): void => {
+    if (editorAvailable) return;
+    api.registerTool({
+      name: "str_replace_editor",
+      description: EDITOR_DESCRIPTION,
+      parameters: Type.Object({
+        command: StringEnum(["view", "create", "str_replace", "insert"] as const),
+        path: Type.String(),
+        file_text: Type.Optional(Type.String()),
+        insert_line: Type.Optional(Type.Integer()),
+        new_str: Type.Optional(Type.String()),
+        old_str: Type.Optional(Type.String()),
+        view_range: Type.Optional(Type.Array(Type.Integer())),
+      }),
+      label: "str_replace_editor",
+      async execute(_callId, input) {
+        const args = input as EditorArgs;
+        return {
+          content: [{ type: "text", text: executeEditor(args) }],
+          details: { command: args.command, path: args.path },
+        };
+      },
+    });
+    editorAvailable = true;
+  };
+
+  const showStage = (context: ExtensionContext, stage?: LiveStage): void => {
+    if (!context.hasUI) return;
+    context.ui.setStatus(
+      STATUS_ID,
+      stage?.stage === "bootstrap" ? "V4 one-shot: bash/editor" : undefined,
+    );
+  };
+
+  const saveStage = (stage: LiveStage): void => {
+    api.appendEntry(STATE_ENTRY_TYPE, {
+      schema: stage.schema,
+      route: stage.route,
+      stage: stage.stage,
+    });
+  };
+
+  const locateStage = (context: ExtensionContext, create: boolean): LiveStage | undefined => {
+    const route = supportedRoute(context.model);
+    if (route === undefined) return undefined;
+
+    const key = sessionKey(context);
+    const cached = sessions.get(key);
+    if (cached?.route === route) return cached;
+
+    const branch = context.sessionManager.getBranch();
+    const restored = restoreStage(branch, route);
+    if (restored !== undefined) {
+      sessions.set(key, restored);
+      return restored;
+    }
+    if (!create || conversationHasStarted(branch)) return undefined;
+
+    const stage: LiveStage = {
+      schema: 1,
+      route,
+      stage: "bootstrap",
+      requestIssued: false,
+    };
+    sessions.set(key, stage);
+    saveStage(stage);
+    return stage;
+  };
+
+  const enterExecution = (context: ExtensionContext, stage: LiveStage): void => {
+    if (stage.stage === "execute") return;
+    stage.stage = "execute";
+    stage.requestIssued = false;
+    saveStage(stage);
+    showStage(context, stage);
+  };
+
+  const reloadSession = (context: ExtensionContext): void => {
+    sessions.delete(sessionKey(context));
+    const stage = locateStage(context, false);
+    if (stage !== undefined) exposeEditor();
+    showStage(context, stage);
+  };
+
+  api.on("session_start", (_event, context) => reloadSession(context));
+  api.on("session_tree", (_event, context) => reloadSession(context));
+  api.on("session_shutdown", (_event, context) => {
+    sessions.delete(sessionKey(context));
+    showStage(context);
+  });
+
+  api.on("model_select", (_event, context) => {
+    const stage = locateStage(context, false);
+    if (stage !== undefined) exposeEditor();
+    showStage(context, stage);
+  });
+
+  api.on("input", (event, context) => {
+    if (event.source === "extension" || supportedRoute(context.model) === undefined) return;
+    const stage = locateStage(context, true);
+    if (stage !== undefined) exposeEditor();
+    showStage(context, stage);
+  });
+
+  api.on("before_agent_start", (event, context) => {
+    const stage = locateStage(context, false);
+    if (stage === undefined) {
+      showStage(context);
+      return;
+    }
+    exposeEditor();
+    showStage(context, stage);
+    if (event.systemPrompt !== MINIMAL_PERSONA) return { systemPrompt: MINIMAL_PERSONA };
+  });
+
+  api.on("before_provider_request", (event, context) => {
+    const stage = locateStage(context, false);
+    if (stage === undefined || !isJsonObject(event.payload)) return;
+    exposeEditor();
+
+    if (stage.stage === "bootstrap") {
+      stage.requestIssued = true;
+      return makeBootstrapRequest(event.payload);
+    }
+    return keepMinimalPersona(event.payload);
+  });
+
+  api.on("message_end", (event, context) => {
+    const stage = locateStage(context, false);
+    if (stage?.stage !== "bootstrap" || !stage.requestIssued) return;
+    if (event.message.role !== "assistant"
+      || event.message.stopReason === "error"
+      || event.message.stopReason === "aborted") return;
+    enterExecution(context, stage);
+  });
+
+  api.on("tool_call", (_event, context) => {
+    const stage = locateStage(context, false);
+    if (stage?.stage === "bootstrap" && stage.requestIssued) enterExecution(context, stage);
+  });
+}
+
+export {
+  BOOTSTRAP_TOOL_NAMES,
+  EDITOR_DESCRIPTION,
+  MINIMAL_BASH_DESCRIPTION,
+  MINIMAL_PERSONA,
+  STATE_ENTRY_TYPE,
+  TARGET_MODEL_ID,
+  TARGET_PROVIDERS,
+};

package/src/protocol.ts

@@ -0,0 +1,134 @@
+import {
+  BOOTSTRAP_TOOL_NAMES,
+  EDITOR_DESCRIPTION,
+  MINIMAL_BASH_DESCRIPTION,
+  MINIMAL_PERSONA,
+} from "./constants.js";
+
+export type JsonObject = Record<string, unknown>;
+
+export function isJsonObject(value: unknown): value is JsonObject {
+  return value !== null && typeof value === "object" && !Array.isArray(value);
+}
+
+function encodeMinimalMessage(value: unknown): unknown {
+  if (!isJsonObject(value)) return value;
+  if (value.role === "system" || value.role === "developer") {
+    return { ...value, content: MINIMAL_PERSONA };
+  }
+  if (!Array.isArray(value.content)) return value;
+
+  const textOnly = value.content.every((part) => isJsonObject(part)
+    && part.type === "text"
+    && typeof part.text === "string");
+  if (!textOnly) return value;
+  return {
+    ...value,
+    content: value.content.map((part) => String((part as JsonObject).text)).join(""),
+  };
+}
+
+function minimalToolSchemas(): unknown[] {
+  const bash = {
+    type: "function",
+    function: {
+      name: BOOTSTRAP_TOOL_NAMES[0],
+      description: MINIMAL_BASH_DESCRIPTION,
+      parameters: {
+        type: "object",
+        properties: {
+          command: {
+            type: "string",
+            description: "The bash command to run. Relative path is preferred in the command.",
+          },
+        },
+        required: ["command"],
+      },
+      strict: false,
+    },
+  };
+  const editor = {
+    type: "function",
+    function: {
+      name: BOOTSTRAP_TOOL_NAMES[1],
+      description: EDITOR_DESCRIPTION,
+      parameters: {
+        type: "object",
+        properties: {
+          command: {
+            type: "string",
+            description: "The commands to run. Allowed options are: `view`, `create`, `str_replace`, `insert`.",
+            enum: ["view", "create", "str_replace", "insert"],
+          },
+          path: {
+            type: "string",
+            description: "Absolute path to file or directory, e.g. `/repo/file.py` or `/repo`.",
+          },
+          file_text: {
+            type: "string",
+            description: "Required parameter of `create` command, with the content of the file to be created.",
+          },
+          insert_line: {
+            type: "integer",
+            description: "Required parameter of `insert` command. The `new_str` will be inserted AFTER the line `insert_line` of `path`.",
+          },
+          new_str: {
+            type: "string",
+            description: "Optional parameter of `str_replace` command containing the new string (if not given, no string will be added). Required parameter of `insert` command containing the string to insert.",
+          },
+          old_str: {
+            type: "string",
+            description: "Required parameter of `str_replace` command containing the string in `path` to replace.",
+          },
+          view_range: {
+            type: "array",
+            description: "Optional parameter of `view` command when `path` points to a file. If none is given, the full file is shown. If provided, the file will be shown in the indicated line number range, e.g. [11, 12] will show lines 11 and 12. Indexing at 1 to start. Setting `[start_line, -1]` shows all lines from `start_line` to the end of the file.",
+            items: { type: "integer" },
+          },
+        },
+        required: ["command", "path"],
+      },
+      strict: false,
+    },
+  };
+  return [bash, editor];
+}
+
+export function makeBootstrapRequest(source: JsonObject): JsonObject {
+  const messages = Array.isArray(source.messages)
+    ? source.messages.map(encodeMinimalMessage)
+    : [];
+  const request: JsonObject = {};
+  const values: JsonObject = {
+    ...source,
+    messages,
+    tools: minimalToolSchemas(),
+  };
+
+  for (const field of [
+    "model",
+    "messages",
+    "stream",
+    "stream_options",
+    "thinking",
+    "reasoning_effort",
+    "tools",
+    "max_tokens",
+  ]) {
+    if (Object.prototype.hasOwnProperty.call(values, field)) request[field] = values[field];
+  }
+  return request;
+}
+
+export function keepMinimalPersona(source: JsonObject): JsonObject | undefined {
+  if (!Array.isArray(source.messages)) return undefined;
+  let changed = false;
+  const messages = source.messages.map((message) => {
+    if (!isJsonObject(message)
+      || (message.role !== "system" && message.role !== "developer")
+      || message.content === MINIMAL_PERSONA) return message;
+    changed = true;
+    return { ...message, content: MINIMAL_PERSONA };
+  });
+  return changed ? { ...source, messages } : undefined;
+}

package/src/session-stage.ts

@@ -0,0 +1,67 @@
+import {
+  STATE_ENTRY_TYPE,
+  TARGET_MODEL_ID,
+  TARGET_PROVIDERS,
+} from "./constants.js";
+import { isJsonObject } from "./protocol.js";
+
+export type Stage = "bootstrap" | "execute";
+
+export interface StageRecord {
+  schema: 1;
+  route: string;
+  stage: Stage;
+}
+
+export interface LiveStage extends StageRecord {
+  requestIssued: boolean;
+}
+
+export function supportedRoute(model: { provider?: unknown; id?: unknown } | undefined): string | undefined {
+  if (typeof model?.provider !== "string" || typeof model.id !== "string") return undefined;
+  const provider = model.provider.toLowerCase();
+  const modelId = model.id.toLowerCase();
+  if (modelId !== TARGET_MODEL_ID || !TARGET_PROVIDERS.some((candidate) => candidate === provider)) {
+    return undefined;
+  }
+  return `${provider}/${modelId}`;
+}
+
+export function conversationHasStarted(entries: readonly unknown[]): boolean {
+  return entries.some((entry) => isJsonObject(entry)
+    && entry.type === "message"
+    && isJsonObject(entry.message)
+    && ["user", "assistant", "toolResult"].includes(String(entry.message.role)));
+}
+
+function completedMessage(entry: unknown): boolean {
+  if (!isJsonObject(entry) || entry.type !== "message" || !isJsonObject(entry.message)) return false;
+  if (entry.message.role === "toolResult") return true;
+  return entry.message.role === "assistant"
+    && entry.message.stopReason !== "error"
+    && entry.message.stopReason !== "aborted";
+}
+
+function decodeRecord(value: unknown, route: string): StageRecord | undefined {
+  if (!isJsonObject(value)
+    || value.schema !== 1
+    || value.route !== route
+    || (value.stage !== "bootstrap" && value.stage !== "execute")) return undefined;
+  return { schema: 1, route, stage: value.stage };
+}
+
+export function restoreStage(entries: readonly unknown[], route: string): LiveStage | undefined {
+  for (let index = entries.length - 1; index >= 0; index -= 1) {
+    const entry = entries[index];
+    if (!isJsonObject(entry)
+      || entry.type !== "custom"
+      || entry.customType !== STATE_ENTRY_TYPE) continue;
+    const record = decodeRecord(entry.data, route);
+    if (record === undefined) continue;
+    const stage = record.stage === "bootstrap" && entries.slice(index + 1).some(completedMessage)
+      ? "execute"
+      : record.stage;
+    return { ...record, stage, requestIssued: false };
+  }
+  return undefined;
+}