@drej/agent 0.1.1 → 0.2.0

package/README.md

@@ -0,0 +1,331 @@
+# @drej/agent
+
+Run [Pi](https://pi.ai) coding agents inside isolated [drej](https://drej.dev) sandbox containers. Pi can read and write files, run shell commands, and execute scripts — streamed back through a simple TypeScript API.
+
+```bash
+bun add @drej/agent
+```
+
+**[Full documentation →](https://docs.drej.dev/docs/agent)**
+
+---
+
+## Quickstart
+
+Create an agent spec (`agents/my-agent.json`):
+
+```json
+{
+  "$schema": "https://registry.drej.dev/schema/agent-item.json",
+  "name": "my-agent",
+  "cli": "pi",
+  "model": "gemini-flash-latest",
+  "packages": ["python3"],
+  "env": { "GEMINI_API_KEY": "${GEMINI_API_KEY}" },
+  "resources": { "cpu": "1000m", "memory": "2Gi" }
+}
+```
+
+```ts
+import { Agent, textOnly } from "@drej/agent";
+
+const agent = await Agent.load("./agents/my-agent.json");
+try {
+  for await (const chunk of textOnly(agent.prompt("Write and run a Python hello world script."))) {
+    process.stdout.write(chunk);
+  }
+} finally {
+  await agent.close();
+}
+```
+
+---
+
+## Agent spec
+
+The spec JSON controls the agent's environment, model, and workspace setup.
+
+| Field        | Type                     | Description                                                           |
+| ------------ | ------------------------ | --------------------------------------------------------------------- |
+| `name`       | `string`                 | Unique identifier, used as the sandbox session name                   |
+| `cli`        | `"pi"`                   | CLI to run (currently only `"pi"`)                                    |
+| `cliVersion` | `string?`                | Pin to a specific Pi version, e.g. `"0.80.2"`. Defaults to latest.    |
+| `model`      | `string?`                | Model ID passed to Pi via `--model`                                   |
+| `provider`   | `string?`                | AI provider passed via `--provider`. Omit for direct Google API key.  |
+| `packages`   | `string[]?`              | APT packages to install before Pi. e.g. `["git", "python3"]`          |
+| `env`        | `Record<string,string>?` | Env vars in the sandbox. Values may reference host env: `"${MY_KEY}"` |
+| `resources`  | `object?`                | CPU/memory limits: `{ cpu: "1000m", memory: "2Gi" }`                  |
+| `setup`      | `SetupStep[]?`           | Workspace setup steps (see below)                                     |
+
+### Setup steps
+
+`setup` runs bash commands after Pi CLI install, before the snapshot is taken. Changes to any step automatically invalidate the snapshot cache.
+
+```json
+{
+  "name": "my-agent",
+  "cli": "pi",
+  "setup": [
+    { "name": "Create workspace", "run": "mkdir -p /workspace" },
+    { "name": "Install deps", "run": "npm install", "cwd": "/workspace" },
+    { "name": "Seed data", "run": "node scripts/seed.js", "cwd": "/workspace" }
+  ]
+}
+```
+
+Each step:
+
+| Field  | Type      | Description                                                       |
+| ------ | --------- | ----------------------------------------------------------------- |
+| `name` | `string`  | Human-readable label shown in logs and included in the setup hash |
+| `run`  | `string`  | Bash command to execute                                           |
+| `cwd`  | `string?` | Working directory. Runs as `cd <cwd> && <run>`                    |
+
+---
+
+## Snapshotting
+
+On first load, `Agent.load()` installs the Pi CLI and any `setup` steps, then checkpoints the sandbox. Subsequent loads restore from that snapshot — skipping the install entirely.
+
+```
+Load 1 (cold):   sandbox → Pi install → setup steps → checkpoint → bridge   ~50s
+Load 2 (warm):   snapshot restore → bridge                                   ~5s
+```
+
+The snapshot is invalidated automatically when `cli`, `cliVersion`, `packages`, or `setup` change.
+
+```ts
+const agent = await Agent.load("./agents/my-agent.json");
+console.log(agent.fromSnapshot); // false on first load, true after
+
+// Force a full reinstall:
+const agent = await Agent.load("./agents/my-agent.json", { rebuild: true });
+```
+
+---
+
+## Streaming
+
+`agent.prompt()` and `agent.bash()` return an `AgentStream` — an `AsyncIterable<AgentEvent>`:
+
+```ts
+type AgentEvent =
+  | { type: "text"; text: string }
+  | { type: "tool_start"; toolCallId: string; toolName: string; args: unknown }
+  | { type: "tool_update"; toolCallId: string; toolName: string; partialResult: unknown }
+  | { type: "tool_end"; toolCallId: string; toolName: string; result: unknown; isError: boolean };
+```
+
+Use `textOnly()` to filter to just the text chunks (equivalent to the old `PromptStream` behavior):
+
+```ts
+import { Agent, textOnly } from "@drej/agent";
+
+for await (const chunk of textOnly(agent.prompt("Summarise this repo."))) {
+  process.stdout.write(chunk);
+}
+```
+
+### Tool call observability
+
+Iterate the raw stream to see every tool Pi uses:
+
+```ts
+for await (const ev of agent.prompt("Run /workspace/script.py with python3.")) {
+  switch (ev.type) {
+    case "text":
+      process.stdout.write(ev.text);
+      break;
+    case "tool_start":
+      console.log(`[tool] ${ev.toolName} args=${JSON.stringify(ev.args)}`);
+      break;
+    case "tool_end":
+      console.log(`[tool] ${ev.toolName} done  isError=${ev.isError}`);
+      break;
+  }
+}
+```
+
+---
+
+## API reference
+
+### Loading and lifecycle
+
+#### `Agent.load(specPath, opts?)`
+
+Load a spec, spin up a sandbox, install Pi, run setup steps, and return a ready `Agent`. Restores from snapshot on subsequent calls.
+
+```ts
+const agent = await Agent.load("./agents/my-agent.json");
+const agent = await Agent.load("./agents/my-agent.json", { rebuild: true });
+```
+
+#### `Agent.resume(sandboxId, opts?)`
+
+Reconnect to an existing sandbox after the host process has exited. Only restarts the bridge — Pi and the workspace are untouched.
+
+```ts
+// Original process saved agent.sandboxId somewhere...
+const agent = await Agent.resume(savedSandboxId);
+// Or provide the spec explicitly:
+const agent = await Agent.resume(savedSandboxId, { specPath: "./agents/my-agent.json" });
+```
+
+#### `agent.close()`
+
+Stop the sandbox container and release all resources. Always call in a `finally` block.
+
+---
+
+### Streaming
+
+#### `agent.prompt(message, opts?)`
+
+Send a message to Pi and stream the response as `AgentStream`.
+
+```ts
+for await (const chunk of textOnly(agent.prompt("Explain this file."))) {
+  process.stdout.write(chunk);
+}
+```
+
+#### `agent.bash(command)`
+
+Run a shell command inside Pi's working context and stream stdout as `AgentStream`.
+
+```ts
+for await (const chunk of textOnly(agent.bash("ls -la /workspace"))) {
+  process.stdout.write(chunk);
+}
+```
+
+---
+
+### Mid-flight control
+
+#### `agent.steer(message)`
+
+Redirect Pi's current response mid-flight. Pi acknowledges and adjusts.
+
+```ts
+const stream = textOnly(agent.prompt("Write an essay on every sorting algorithm..."));
+setTimeout(() => agent.steer("Stop — give me 3 bullet points instead."), 1500);
+for await (const chunk of stream) process.stdout.write(chunk);
+```
+
+#### `agent.followUp(message)`
+
+Queue a message for Pi to process after it finishes the current task.
+
+#### `agent.abort()`
+
+Interrupt the current in-progress response immediately.
+
+---
+
+### Session management
+
+#### `agent.newSession()`
+
+Start a fresh Pi conversation, clearing all context. Filesystem and workspace are unchanged.
+
+#### `agent.clone()`
+
+Branch the current Pi session at the current position. Returns `{ cancelled: boolean }`.
+
+#### `agent.fork(entryId)`
+
+Branch from a specific message entry in the conversation history. Returns `{ text, cancelled }`.
+
+#### `agent.switchSession(sessionPath)`
+
+Switch Pi to a different session file on disk.
+
+#### `agent.getMessages()`
+
+Retrieve the full conversation history for the current session.
+
+```ts
+const messages = await agent.getMessages();
+console.log(messages.length, "messages");
+```
+
+---
+
+### Model control
+
+#### `agent.setModel(provider, modelId)`
+
+Switch Pi to a specific model. Returns the activated `PiModel`.
+
+#### `agent.cycleModel()`
+
+Cycle to the next configured model. Returns `{ model, thinkingLevel, isScoped }` or `null` if only one model is configured.
+
+#### `agent.getAvailableModels()`
+
+List all models available to Pi under the current provider configuration.
+
+#### `agent.setThinkingLevel(level)`
+
+Set Pi's reasoning level (`"low" | "medium" | "high"`). Only effective on models that support extended thinking.
+
+#### `agent.cycleThinkingLevel()`
+
+Cycle Pi's thinking level. Returns `{ level }` or `null` if the current model doesn't support thinking.
+
+---
+
+### Context management
+
+#### `agent.setAutoCompaction(enabled)`
+
+Enable or disable Pi's automatic context compaction.
+
+#### `agent.compact(customInstructions?)`
+
+Manually trigger Pi's context compaction. Returns `{ tokensBefore, estimatedTokensAfter }`.
+
+---
+
+### Environment and debugging
+
+#### `agent.setEnv(vars)`
+
+Set or update env vars in the running container. Restarts Pi so it picks up the new env.
+
+```ts
+await agent.setEnv({ DATABASE_URL: "postgres://..." });
+```
+
+#### `agent.getLogs()`
+
+Retrieve recent bridge logs (ring-buffered, last 200 entries).
+
+#### `agent.sandbox`
+
+Direct access to the underlying `Sandbox` — run commands, read/write files, or inspect state independently of Pi.
+
+```ts
+await agent.sandbox.writeFile("/workspace/input.txt", data);
+const { stdout } = await agent.sandbox.exec("wc -l /workspace/input.txt");
+const result = await agent.sandbox.readFile("/workspace/output.txt");
+```
+
+---
+
+## Properties
+
+| Property             | Type      | Description                                    |
+| -------------------- | --------- | ---------------------------------------------- |
+| `agent.sandboxId`    | `string`  | OpenSandbox container ID                       |
+| `agent.name`         | `string`  | Agent name from the spec                       |
+| `agent.sandbox`      | `Sandbox` | Underlying drej `Sandbox` object               |
+| `agent.fromSnapshot` | `boolean` | `true` when restored from snapshot (fast path) |
+
+---
+
+## License
+
+Apache 2.0

package/dist/index.d.mts

@@ -0,0 +1,341 @@
+import { Sandbox } from "@drej/core";
+
+//#region src/schema.d.ts
+/**
+ * A single named setup step run inside the sandbox after Pi CLI install,
+ * before the snapshot is taken. The step is a bash shell command.
+ */
+interface SetupStep {
+  /** Human-readable label shown in logs and included in the setup hash. */
+  name: string;
+  /** Shell command to run (bash). */
+  run: string;
+  /** If set, the command runs as `cd <cwd> && <run>`. */
+  cwd?: string;
+}
+/**
+ * JSON spec for an agent — typically loaded from an `agent.json` file on disk.
+ * Pass the path to `Agent.load(specPath)`.
+ *
+ * Environment variable references in `env` values are interpolated from
+ * `process.env` at load time: `"${MY_API_KEY}"` → `process.env.MY_API_KEY`.
+ */
+interface AgentSpec {
+  $schema?: string;
+  /** Unique identifier for this agent. Used as the sandbox session name. */
+  name: string;
+  /** Human-readable display name. */
+  title?: string;
+  description?: string;
+  author?: string;
+  categories?: string[];
+  /** CLI to run inside the sandbox. Currently only `"pi"` is supported. */
+  cli: "pi";
+  /** Pin to a specific CLI version (e.g. `"0.80.2"`). Defaults to latest. */
+  cliVersion?: string;
+  /**
+   * AI provider passed to the CLI via `--provider`. For Pi with a direct Google
+   * API key, omit this (Pi defaults to the Google Generative AI endpoint).
+   */
+  provider?: string;
+  /**
+   * Model ID passed to the CLI via `--model`.
+   * For Pi 0.80.x with a free Google AI Studio key, use `"gemini-flash-latest"`
+   * (Pi's alias for gemini-3.5-flash via the direct Google Generative AI API).
+   */
+  model?: string;
+  /**
+   * APT packages to install in the sandbox before starting the CLI.
+   * Example: `["python3", "git"]`. `nodejs_22` and `nodejs` are silently ignored
+   * since the base image is `node:22`.
+   */
+  packages?: string[];
+  /**
+   * Environment variables available inside the sandbox.
+   * Values may reference host env vars: `{ GEMINI_API_KEY: "${GEMINI_API_KEY}" }`.
+   */
+  env?: Record<string, string>;
+  /**
+   * CPU/memory/GPU resource limits for the sandbox container.
+   * Falls back to defaults in `drej.config.json` if omitted.
+   */
+  resources?: {
+    cpu: string;
+    memory: string;
+    gpu?: string;
+  };
+  /** Arbitrary labels attached to the sandbox. */
+  metadata?: Record<string, string>;
+  /** Other agent specs to load as dependencies before this one. */
+  registryDependencies?: string[];
+  /**
+   * Setup steps run inside the sandbox after Pi CLI install, before the snapshot.
+   * Baked into the snapshot — any change to a step invalidates the cache automatically.
+   * Example: create directories, write seed files, install project dependencies.
+   */
+  setup?: SetupStep[];
+}
+declare function validateAgentSpec(data: unknown): AgentSpec;
+//#endregion
+//#region src/snapshots.d.ts
+interface AgentSnapshotRecord {
+  specName: string;
+  setupHash: string;
+  /** OpenSandbox snapshot ID returned by `sb.checkpoint()`. Pass to `client.restoreSnapshot()`. */
+  snapshotId: string;
+  createdAt: number;
+}
+/**
+ * Hash of the fields that require re-installing the agent CLI:
+ * cli, cliVersion, and packages. Excludes env (hot-reloadable),
+ * model/provider (CLI flags only), and cosmetic fields.
+ */
+declare function computeSetupHash(spec: AgentSpec): string;
+/** Derive the agent-snapshots.json path from the ledger adapter path. */
+declare function snapshotsPath(adapterPath: string): string;
+/** File-backed store for agent snapshot records, keyed by specName + setupHash. */
+declare class AgentSnapshotStore {
+  private readonly path;
+  constructor(path: string);
+  private _read;
+  private _write;
+  private _key;
+  get(specName: string, setupHash: string): Promise<AgentSnapshotRecord | null>;
+  save(record: AgentSnapshotRecord): Promise<void>;
+  /** Remove all snapshot records for the given spec name. */
+  delete(specName: string): Promise<void>;
+  list(): Promise<AgentSnapshotRecord[]>;
+}
+//#endregion
+//#region src/types.d.ts
+/**
+ * A single event emitted by an agent stream. Discriminated on `type`.
+ *
+ * - `text` — a chunk of Pi's text response (assistant prose)
+ * - `tool_start` — Pi started executing a tool (bash, file write, etc.)
+ * - `tool_update` — streaming progress from an in-flight tool execution
+ * - `tool_end` — a tool finished; `result` holds its output, `isError` whether it failed
+ */
+type AgentEvent = {
+  type: "text";
+  text: string;
+} | {
+  type: "tool_start";
+  toolCallId: string;
+  toolName: string;
+  args: unknown;
+} | {
+  type: "tool_update";
+  toolCallId: string;
+  toolName: string;
+  partialResult: unknown;
+} | {
+  type: "tool_end";
+  toolCallId: string;
+  toolName: string;
+  result: unknown;
+  isError: boolean;
+};
+/**
+ * Async iterable of structured agent events. Returned by `Agent.prompt()` and `Agent.bash()`.
+ *
+ * Use `textOnly()` to filter down to just the text chunks:
+ *
+ * ```ts
+ * for await (const chunk of textOnly(agent.prompt("Hello"))) {
+ *   process.stdout.write(chunk);
+ * }
+ * ```
+ *
+ * Or iterate the full stream to observe tool calls:
+ *
+ * ```ts
+ * for await (const ev of agent.prompt("Build a widget")) {
+ *   if (ev.type === "text") process.stdout.write(ev.text);
+ *   else if (ev.type === "tool_start") console.log("Pi is running:", ev.toolName, ev.args);
+ *   else if (ev.type === "tool_end") console.log("Tool done, isError:", ev.isError);
+ * }
+ * ```
+ */
+type AgentStream = AsyncIterable<AgentEvent>;
+/** Filter an `AgentStream` down to just text chunks. */
+declare function textOnly(stream: AgentStream): AsyncIterable<string>;
+/**
+ * @deprecated Use `AgentStream`. `PromptStream` will be removed in a future minor.
+ */
+type PromptStream = AsyncIterable<string>;
+/** A Pi model as returned by `getAvailableModels`, `setModel`, and `cycleModel`. */
+interface PiModel {
+  id: string;
+  api: string;
+  [key: string]: unknown;
+}
+/** Thinking (reasoning) level supported by models that have extended thinking. */
+type ThinkingLevel = "none" | "low" | "medium" | "high";
+/** A single message in Pi's conversation history, as returned by `getMessages`. */
+interface PiMessage {
+  role: "user" | "assistant";
+  [key: string]: unknown;
+}
+/** Result of a `compact` operation. */
+interface CompactResult {
+  summary: string;
+  firstKeptEntryId: string;
+  tokensBefore: number;
+  estimatedTokensAfter: number;
+}
+//#endregion
+//#region src/agent.d.ts
+/**
+ * A live AI coding agent running inside an OpenSandbox container.
+ *
+ * Wraps a Pi CLI process (`pi --mode rpc --approve`) in an HTTP bridge so the
+ * host can send prompts and receive streamed responses over a stable API,
+ * while Pi manages its own tool use, file writes, and code execution inside the
+ * sandbox.
+ *
+ * Create an agent with `Agent.load(specPath)`. Always call `close()` when done
+ * to release the underlying sandbox container.
+ *
+ * @example
+ * ```ts
+ * import { Agent } from "@drej/agent";
+ *
+ * const agent = await Agent.load("./agents/my-agent.json");
+ * try {
+ *   for await (const chunk of agent.prompt("Explain this codebase")) {
+ *     process.stdout.write(chunk);
+ *   }
+ * } finally {
+ *   await agent.close();
+ * }
+ * ```
+ */
+declare class Agent {
+  /** OpenSandbox container ID for this agent's sandbox. */
+  readonly sandboxId: string;
+  readonly name: string;
+  /**
+   * Direct access to the underlying `Sandbox` — full drej Sandbox API, bypasses Pi.
+   * Use this to read or write files, run shell commands, or inspect container state
+   * independently of the Pi conversation.
+   */
+  readonly sandbox: Sandbox;
+  /**
+   * `true` when this agent was loaded from a cached snapshot (fast path).
+   * `false` on the first load for a given spec, or after `{ rebuild: true }`.
+   */
+  readonly fromSnapshot: boolean;
+  private readonly _adapter;
+  private _env;
+  private constructor();
+  /**
+   * Load an agent spec from disk and return a fully initialised `Agent`.
+   *
+   * On first load the Pi CLI is installed inside a `node:22` sandbox, then
+   * the sandbox is checkpointed. Subsequent `load()` calls for the same spec
+   * restore from that snapshot — skipping the install and starting in ~3s instead
+   * of ~90s.
+   *
+   * Pass `{ rebuild: true }` to force a full reinstall (e.g. after changing
+   * the spec's `packages` or `cliVersion`).
+   *
+   * Logs timing for each phase to stdout via `[agent]` prefixed lines.
+   */
+  static load(specPath: string, opts?: {
+    rebuild?: boolean;
+  }): Promise<Agent>;
+  /**
+   * Reconnect to a previously-created agent whose host process has exited.
+   *
+   * The sandbox container must still be running. Pi and any installed packages
+   * are already present — only the bridge process needs to be restarted.
+   * Pi is started with `--continue` so it resumes the most recent session.
+   *
+   * @param sandboxId  The sandbox ID returned by the original `Agent.load()`.
+   * @param opts.specPath  Path to the agent spec JSON. If omitted, the ledger
+   *   is queried for the sandbox name and the spec is loaded from
+   *   `./agents/<name>.json`.
+   *
+   * @example
+   * ```ts
+   * // Original process:
+   * const agent = await Agent.load("./agents/hello-agent.json");
+   * console.log(agent.sandboxId); // save this
+   * // ... process exits ...
+   *
+   * // New process:
+   * const agent = await Agent.resume(savedSandboxId);
+   * for await (const chunk of agent.prompt("What did we discuss earlier?")) {
+   *   process.stdout.write(chunk);
+   * }
+   * await agent.close();
+   * ```
+   */
+  static resume(sandboxId: string, opts?: {
+    specPath?: string;
+  }): Promise<Agent>;
+  /** Send a prompt to Pi and stream the response. Pi manages its own session context. */
+  prompt(message: string, opts?: {
+    streamingBehavior?: "steer" | "followUp";
+  }): AgentStream;
+  /** Run a shell command inside Pi's working context and stream stdout. */
+  bash(command: string): AgentStream;
+  /** Steer Pi's current response mid-flight. Waits for Pi's RPC acknowledgment. */
+  steer(message: string): Promise<void>;
+  /** Abort Pi's current operation. */
+  abort(): Promise<void>;
+  /** Queue a message to be sent to Pi after it finishes its current task. */
+  followUp(message: string): Promise<void>;
+  /** Start a fresh Pi session, clearing all prior context. */
+  newSession(): Promise<void>;
+  /** Set Pi's reasoning level (for models that support extended thinking). */
+  setThinkingLevel(level: ThinkingLevel): Promise<void>;
+  /** Enable or disable Pi's automatic context compaction. */
+  setAutoCompaction(enabled: boolean): Promise<void>;
+  /**
+   * Fork Pi's session at the given entry ID, creating a new branch.
+   * Returns the text of the forked message and whether the fork was cancelled.
+   */
+  fork(entryId: string): Promise<{
+    text: string;
+    cancelled: boolean;
+  }>;
+  /** Clone the current Pi session into a new branch at the current position. */
+  clone(): Promise<{
+    cancelled: boolean;
+  }>;
+  /** Switch Pi to a different session file on disk. */
+  switchSession(sessionPath: string): Promise<{
+    cancelled: boolean;
+  }>;
+  /** Switch Pi to a specific model. Returns the activated model. */
+  setModel(provider: string, modelId: string): Promise<PiModel>;
+  /** Cycle Pi to the next available model. Returns null if only one model is configured. */
+  cycleModel(): Promise<{
+    model: PiModel;
+    thinkingLevel: ThinkingLevel;
+    isScoped: boolean;
+  } | null>;
+  /** Cycle Pi's thinking level. Returns null if the current model doesn't support thinking. */
+  cycleThinkingLevel(): Promise<{
+    level: ThinkingLevel;
+  } | null>;
+  /** Manually trigger Pi's context compaction. */
+  compact(customInstructions?: string): Promise<CompactResult>;
+  /** Retrieve Pi's full conversation history for the current session. */
+  getMessages(): Promise<PiMessage[]>;
+  /** List all models available to Pi under the current provider configuration. */
+  getAvailableModels(): Promise<PiModel[]>;
+  /**
+   * Set or update env vars in the running container. Writes to /etc/drej-env and restarts
+   * the Pi subprocess so it picks up the new env. Waits until Pi is ready before returning.
+   */
+  setEnv(vars: Record<string, string>): Promise<void>;
+  /** Retrieve recent bridge logs (ring-buffered, last 200 entries). */
+  getLogs(): Promise<string>;
+  /** Delete the sandbox container and release all resources. Always call in a `finally` block. */
+  close(): Promise<void>;
+}
+//#endregion
+export { Agent, type AgentEvent, type AgentSnapshotRecord, AgentSnapshotStore, type AgentSpec, type AgentStream, type CompactResult, type PiMessage, type PiModel, type PromptStream, type SetupStep, type ThinkingLevel, computeSetupHash, snapshotsPath, textOnly, validateAgentSpec };