char-slop/berry

git clone https://git.t4t.associates/char-slop/berry

Charlotte Somadd a /compact command to misc2c01ad92026-10-18

// Run one extension without a harness: state folds through its reducer, arguments are validated like the runtime does.

import { homedir } from "node:os";
import { resolve } from "node:path";
import { BACKGROUND_CONTEXT } from "@earendil-works/chord/context";
import { type Provider, validateToolArguments } from "@earendil-works/pi-ai";
import type { FileError, FileInfo, Result } from "@earendil-works/pi-durable/env";
import { NodeExecutionEnv } from "@earendil-works/pi-durable/env/node";
import {
  connectHere,
  directory,
  type Environment,
  localEnvironment,
  spawnHere,
} from "./environment.ts";
import { ptyHere } from "./pty.ts";
import {
  type Channel,
  type ChannelHandler,
  type ChannelMessage,
  type Command,
  type CompactionCap,
  type Configuration,
  configurationEvents,
  type ContextHook,
  type Extension,
  type ExtensionApi,
  type Handler,
  handlerArguments,
  type JsonValue,
  type PayloadHook,
  type PromptProse,
  type PromptRender,
  type SectionRender,
  tightestCap,
  type Tool,
  type ToolContext,
} from "./mod.ts";

export function testExtension<S extends JsonValue, E extends JsonValue>(
  extension: Extension<S, E>,
  {
    events = [],
    tid = "3mxd6epszps5d",
    env = localEnvironment(homedir()),
    auth = () => Promise.resolve(undefined),
  }: {
    events?: readonly E[];
    auth?: ExtensionApi<S, E>["auth"];
    tid?: string;
    /** Where its files and commands are; `cd` moves its `cwd`. Default: this machine's home. */
    env?: Environment;
  } = {},
) {
  let state = events.reduce((s, e) => extension.apply?.(s, e) ?? s, extension.initial());
  const messages: Parameters<ContextHook>[0][number][] = [];
  const emit = (event: E, options?: { content: string }) => {
    state = extension.apply?.(state, event) ?? state;
    if (options) messages.push({ role: "user", content: options.content, timestamp: Date.now() });
    return Promise.resolve();
  };
  const providers: Provider[] = [];
  const compactions: (string | undefined)[] = [];
  const compactionGetters: (() => CompactionCap | undefined)[] = [];
  const statics: Tool<S, E>[] = [];
  const handlers: Handler<S>[] = [];
  const channels: ChannelHandler<S>[] = [];
  const disposers: (() => void | Promise<void>)[] = [];
  const sections: SectionRender<S>[] = [];
  let configuration: Configuration<S, E> | undefined;
  let prose: PromptProse | undefined;
  const contribute = extension.setup({
    auth,
    provider: (provider) => void providers.push(provider),
    tool: (tool) => void statics.push(tool as Tool<S, E>),
    handle: (handler) => void handlers.push(handler as unknown as Handler<S>),
    channel: (handler) => void channels.push(handler as unknown as ChannelHandler<S>),
    prompt: (_name, render) => void sections.push(render),
    promptProse: (given) => void (prose = given),
    compactionCap: (getter) => void compactionGetters.push(getter),
    configuration: (given) => void (configuration = given as unknown as Configuration<S, E>),
    onDispose: (dispose) => void disposers.push(dispose),
  });

  function contribution(currentEnv = env) {
    const tools: Tool<S, E>[] = [...statics];
    const prompts: PromptRender[] = [];
    const payloadHooks: PayloadHook[] = [];
    const contextHooks: ContextHook[] = [];
    const commands: Command<S, E>[] = [];
    contribute?.({
      tid,
      state,
      env: currentEnv,
      tool: (tool) => void tools.push(tool as Tool<S, E>),
      prompt: (render) => void prompts.push(render),
      onPayload: (hook) => void payloadHooks.push(hook),
      onContext: (hook) => void contextHooks.push(hook),
      command: (command) => void commands.push(command),
    });
    return { tools, prompts, commands, contextHooks, payloadHooks };
  }

  /** Call a tool as codemode would: validated arguments, the value returned as is. */
  async function call(
    name: string,
    args: JsonValue,
    options: {
      signal?: AbortSignal;
      /** As though the conversation were on another machine for this one call. */
      env?: Environment;
      output?: (chunk: string) => void;
      progress?: (details: JsonValue) => void;
    } = {},
  ): Promise<JsonValue | undefined> {
    const { tools } = contribution(options.env ?? env);
    const tool = tools.findLast((tool) => tool.name === name);
    if (tool === undefined || (tool.enabled && !await tool.enabled({ env: options.env ?? env }))) {
      throw new Error(`No tool named ${name}`);
    }
    const validated = validateToolArguments(tool, {
      type: "toolCall",
      id: "test",
      name,
      arguments: args as never,
    });
    const ctx: ToolContext<S, E> = {
      tid,
      get state() {
        return state;
      },
      env: options.env ?? env,
      signal: options.signal ?? new AbortController().signal,
      emit,
      output: (chunk) => options.output?.(chunk),
      progress: (details) => options.progress?.(details),
    };
    return await (tool as Tool<S, E>).execute(validated, ctx);
  }

  return {
    get state() {
      return state;
    },
    get messages(): Parameters<ContextHook>[0] {
      return structuredClone(messages);
    },
    /** The instructions of each manual compaction a command started, in order; `undefined` for none given. */
    get compactions(): readonly (string | undefined)[] {
      return [...compactions];
    },
    /** The tightest compaction cap this load reports, as the harness would resolve it. */
    get compactionCap(): CompactionCap | undefined {
      const seen = compactionGetters.map((get) => {
        try {
          return get();
        } catch {
          return undefined;
        }
      });
      return tightestCap(seen);
    },
    env,
    providers: providers as readonly Provider[],
    call,
    async context(
      messages: Parameters<ContextHook>[0],
      model: Parameters<ContextHook>[1],
    ): Promise<Parameters<ContextHook>[0]> {
      messages = structuredClone(messages);
      for (const hook of contribution().contextHooks) {
        const next = await hook(messages, model);
        if (next !== undefined) messages = next;
      }
      return messages;
    },
    async payload(payload: unknown, model: Parameters<PayloadHook>[1]): Promise<unknown> {
      for (const hook of contribution().payloadHooks) {
        const next = await hook(payload, model);
        if (next !== undefined) payload = next;
      }
      return payload;
    },
    /** Ask a handler as the extension's frontend would: validated arguments, the value returned as is. */
    async request(
      name: string,
      args?: JsonValue,
      options: { signal?: AbortSignal; env?: Environment } = {},
    ): Promise<JsonValue | undefined> {
      const handler = handlers.findLast((handler) => handler.name === name);
      if (handler === undefined) throw new Error(`No handler named ${name}`);
      return await handler.execute(handlerArguments(handler.parameters, args), {
        tid,
        state,
        env: options.env ?? env,
        signal: options.signal ?? new AbortController().signal,
      });
    },
    /** Open a channel as the extension's frontend would, and get the frontend's end of it. */
    connect(name: string, args?: JsonValue, options: { env?: Environment } = {}): Channel {
      const handler = channels.findLast((handler) => handler.name === name);
      if (handler === undefined) throw new Error(`No channel named ${name}`);
      const { server, client, signal, fail } = channelPair();
      // a throw, even while validating, refuses the channel rather than this call
      void Promise.resolve().then(() =>
        handler.open(handlerArguments(handler.parameters, args), server, {
          tid,
          state,
          env: options.env ?? env,
          signal,
        })
      ).catch(fail);
      return client;
    },
    /** Run a command as the user; its events apply at once, as they do in an idle conversation. */
    async command(name: string, args = ""): Promise<JsonValue | undefined> {
      const { commands } = contribution();
      const command = commands.findLast((command) => command.name === name);
      if (command === undefined || (command.enabled && !await command.enabled({ env }))) {
        throw new Error(`No command /${name}`);
      }
      return await command.execute(args, {
        tid,
        state,
        env,
        emit,
        async cd(path) {
          env.cwd = await directory(env, path, BACKGROUND_CONTEXT);
          return env.cwd;
        },
        compact: (instructions) => {
          compactions.push(instructions);
          return Promise.resolve();
        },
      });
    },
    /** The events a conversation started with these values would begin with. */
    async seed(values: JsonValue = {}): Promise<readonly E[]> {
      if (!configuration) throw new Error(`${extension.name} has no configuration`);
      return await configurationEvents(configuration, values, extension.initial());
    },
    get configuration() {
      return configuration?.values(state);
    },
    /** Apply a configuration patch through the extension's events, as in an idle conversation. */
    async configure(values: Record<string, JsonValue>): Promise<void> {
      if (!configuration) throw new Error(`${extension.name} has no configuration`);
      for (const key of Object.keys(values)) {
        if (!Object.hasOwn(configuration.parameters.properties, key)) {
          throw new Error(`unknown configuration field: ${key}`);
        }
      }
      const events = await configurationEvents(configuration, {
        ...configuration.values(state),
        ...values,
      }, state);
      for (const event of events) await emit(event);
    },
    async prompt(): Promise<string> {
      // the harness replaces `$USER` at render; the helper keeps the default as entered
      const texts = await Promise.all([
        ...(prose && prose.default.trim() ? [prose.default] : []),
        ...contribution().prompts.map((render) => render({ env })),
        ...sections.map((render) => render({ tid, state, env })),
      ]);
      return texts.filter((text) => text !== undefined).join("\n\n");
    },
    async dispose() {
      for (const dispose of disposers) await dispose();
    },
  };
}

/** Two ends of a channel in memory, as a websocket would join them; what crosses is copied, as it would be. */
function channelPair() {
  const closed = new AbortController();
  const end = () => {
    let controller!: ReadableStreamDefaultController<ChannelMessage>;
    return {
      messages: new ReadableStream<ChannelMessage>({ start: (c) => void (controller = c) }),
      controller,
    };
  };
  const [a, b] = [end(), end()];
  const close = () => {
    if (closed.signal.aborted) return;
    closed.abort();
    for (const { controller } of [a, b]) controller.close();
  };
  const channel = (own: typeof a, other: typeof a): Channel => ({
    send: (message) =>
      void (!closed.signal.aborted && other.controller.enqueue(structuredClone(message))),
    messages: own.messages,
    close,
  });
  return {
    server: channel(a, b),
    client: channel(b, a),
    signal: closed.signal,
    /** Refuse the channel, as the server would: the frontend's messages fail with why. */
    fail(error: unknown) {
      if (closed.signal.aborted) return;
      closed.abort();
      a.controller.close();
      b.controller.error(error);
    },
  };
}

/**
 * An environment whose `/` is `root` on this machine, so a test can tell the files an extension reaches through it
 * from this machine's own: a path that bypasses it names nothing. Commands still run here, from the mapped `cwd`.
 */
export function rootedEnvironment(
  root: string,
  { cwd = "/", home = "/home", variables = {} }: {
    cwd?: string;
    home?: string;
    variables?: Record<string, string>;
  } = {},
): Environment {
  const inside = (path: string) => {
    const expanded = path === "~" ? home : path.startsWith("~/") ? home + path.slice(1) : path;
    return expanded.startsWith("/") ? root + expanded : expanded;
  };
  const outside = (path: string) =>
    path === root ? "/" : path.startsWith(`${root}/`) ? path.slice(root.length) : path;
  const mapped = <T, E>(result: Result<T, E>, map: (value: T) => T): Result<T, E> =>
    result.ok ? { ok: true, value: map(result.value) } : result;
  const info = (file: FileInfo) => ({ ...file, path: outside(file.path) });
  const inner = new NodeExecutionEnv({ cwd: inside(cwd) });
  const temporary = async (prefix: string | undefined): Promise<Result<string, FileError>> => {
    const created = await inner.createDir(`${root}/tmp`, { recursive: true }, BACKGROUND_CONTEXT);
    if (!created.ok) return created;
    const dir = await Deno.makeTempDir({ dir: `${root}/tmp`, prefix: prefix ?? "tmp-" });
    return { ok: true, value: outside(dir) };
  };
  return {
    id: `rooted:${root}`,
    name: "rooted",
    platform: "linux",
    home,
    variables,
    get cwd() {
      return outside(inner.cwd);
    },
    set cwd(path: string) {
      inner.cwd = inside(path);
    },
    absolutePath: async (path, context) =>
      mapped(await inner.absolutePath(inside(path), context), outside),
    joinPath: (parts, context) => inner.joinPath(parts, context),
    readTextFile: (path, context) => inner.readTextFile(inside(path), context),
    openTextLineReader: (path, context) => inner.openTextLineReader(inside(path), context),
    readTextLines: (path, options, context) => inner.readTextLines(inside(path), options, context),
    readBinaryFile: (path, context) => inner.readBinaryFile(inside(path), context),
    openBinaryReader: async (path, options, context) =>
      mapped(await inner.openBinaryReader(inside(path), options, context), (reader) => ({
        ...reader,
        info: async (context) => mapped(await reader.info(context), info),
        read: (offset, length, context) => reader.read(offset, length, context),
        scanLines: (options, context) => reader.scanLines(options, context),
        close: (context) => reader.close(context),
      })),
    writeFile: (path, content, context) => inner.writeFile(inside(path), content, context),
    appendFile: (path, content, context) => inner.appendFile(inside(path), content, context),
    truncateFile: (path, size, context) => inner.truncateFile(inside(path), size, context),
    flushFile: (path, context) => inner.flushFile(inside(path), context),
    renameFile: (from, to, context) => inner.renameFile(inside(from), inside(to), context),
    fileInfo: async (path, context) => mapped(await inner.fileInfo(inside(path), context), info),
    listDir: async (path, context) =>
      mapped(await inner.listDir(inside(path), context), (files) => files.map(info)),
    openDirReader: async (path, context) =>
      mapped(await inner.openDirReader(inside(path), context), (reader) => ({
        next: async (count, context) =>
          mapped(await reader.next(count, context), (page) => ({
            ...page,
            entries: page.entries.map(info),
          })),
        close: (context) => reader.close(context),
      })),
    watch: (targets, onChange, context) =>
      inner.watch(
        targets.map((target) => ({ ...target, path: inside(target.path) })),
        (change) => onChange("paths" in change ? { paths: change.paths.map(outside) } : change),
        context,
      ),
    canonicalPath: async (path, context) =>
      mapped(await inner.canonicalPath(inside(path), context), outside),
    exists: (path, context) => inner.exists(inside(path), context),
    createDir: (path, options, context) => inner.createDir(inside(path), options, context),
    remove: (path, options, context) => inner.remove(inside(path), options, context),
    createTempDir: (prefix) => temporary(prefix),
    async createTempFile(options, context) {
      const dir = await temporary("tmp-");
      if (!dir.ok) return dir;
      const path = `${dir.value}/${options?.prefix ?? ""}file${options?.suffix ?? ""}`;
      const written = await inner.writeFile(inside(path), "", context);
      return written.ok ? { ok: true, value: path } : written;
    },
    exec: async (command, options, context) =>
      await inner.exec(
        command,
        options?.cwd === undefined ? options : { ...options, cwd: inside(options.cwd) },
        context,
      ),
    cleanup: (context) => inner.cleanup(context),
    spawn: (argv, options, context) =>
      Promise.resolve(spawnHere(argv, {
        ...options,
        cwd: options.cwd === undefined ? inner.cwd : resolve(inner.cwd, inside(options.cwd)),
      }, context)),
    connect: (hostname, port, context) => connectHere(hostname, port, context),
    fetch: (input, init) => fetch(input, init),
    pty: (argv, options, context) =>
      new Promise((done) =>
        done(ptyHere(argv, {
          ...options,
          cwd: options.cwd === undefined ? inner.cwd : resolve(inner.cwd, inside(options.cwd)),
        }, context))
      ),
  };
}