experimental-a2 0.17.0 → 0.18.1

package/src/ai-pi.ts

@@ -0,0 +1,1322 @@
+/**
+ * experimental-a2/ai/pi — `pi()`: the Pi harness for `agent({ harness })`.
+ * Server-only; Pi comes from the optional `@earendil-works/pi-coding-agent`
+ * and `@earendil-works/pi-ai` peer dependencies.
+ */
+
+// oxlint-disable no-await-in-loop -- turns replay in order, each on the session the last one left.
+
+import { createHash } from 'node:crypto'
+import { readFile as readLocalFile, unlink } from 'node:fs/promises'
+import { tmpdir } from 'node:os'
+import { basename, dirname, join } from 'node:path'
+import type {
+  AssistantMessage,
+  AssistantMessageEvent,
+  AssistantMessageEventStream,
+  JsonObject,
+  TextContent,
+  ThinkingContent,
+  ToolCall,
+  Usage,
+} from '@earendil-works/pi-ai'
+import type {
+  AgentSession,
+  AgentToolResult,
+  BashOperations,
+  ModelRuntime,
+  ResourceLoader,
+  SessionEntry,
+  SessionHeader,
+  ToolDefinition,
+} from '@earendil-works/pi-coding-agent'
+import {
+  asSchema,
+  jsonSchema,
+  type Experimental_SandboxSession,
+  type FinishReason,
+  type JSONSchema7,
+  type LanguageModelUsage,
+  type Tool,
+  type ToolExecutionOptions,
+  type ToolSet,
+  type UIMessageChunk,
+} from 'ai'
+import {
+  type Deferred,
+  defineHarness,
+  deferred,
+  Divergence,
+  gatewayCredential,
+  type Harness,
+  type HarnessGeneration,
+  type HarnessRead,
+  type HarnessSandbox,
+  type HarnessStep,
+  type HarnessTurn,
+  pending,
+  readHarness,
+  type RecordedGeneration,
+  saveHarness,
+} from './ai-harness.ts'
+import { type SchemaInput, schemaInput } from './ai-tool-input.ts'
+import type { CompactionPartData, ToolResultPayload } from './ai.ts'
+import { boundedCache } from './bounded-cache.ts'
+import { ambientToolScopeStorage } from './internal.ts'
+
+export type {
+  Harness,
+  HarnessSandbox,
+  HarnessSandboxContext,
+} from './ai-harness.ts'
+
+/** What `pi()` takes. */
+export type PiOptions = {
+  /** The session's sandbox: Pi's tools, and the agent's, work in it. */
+  sandbox: HarnessSandbox
+}
+
+/**
+ * The Pi harness. Every model call Pi makes is one durable step; A2 rebuilds
+ * Pi's session for each step from its latest snapshot and the log since.
+ */
+export function pi(options: PiOptions): Harness<'pi', Builtin> {
+  if (typeof options?.sandbox !== 'function')
+    throw new TypeError('pi() requires sandbox')
+  return defineHarness<'pi', Builtin>('pi', {
+    sandbox: options.sandbox,
+    tools: builtinTools,
+    step: (step) =>
+      runStep(step, async () =>
+        options.sandbox({
+          session: { id: step.sessionId },
+          signal: step.signal,
+        }),
+      ),
+  })
+}
+
+const SNAPSHOT_VERSION = 1
+const GATEWAY = 'vercel-ai-gateway'
+// The most a file read or a command's output moves from the sandbox to the
+// server. Pi's tools show the model 50 KB at a time, after reading it all.
+const MAX_TRANSFER_BYTES = 4 * 1024 * 1024
+const MAX_TRANSFER = '4 MiB'
+// Pi's default tools.
+const BUILTINS = ['read', 'write', 'edit', 'bash'] as const
+type Builtin = (typeof BUILTINS)[number]
+
+type PiModules = {
+  agent: typeof import('@earendil-works/pi-coding-agent')
+  ai: typeof import('@earendil-works/pi-ai')
+}
+
+/**
+ * A content block without the text the transcript already holds. A tool call
+ * keeps its arguments only where the transcript's input is not the model's.
+ */
+type StoredBlock =
+  | (Omit<TextContent, 'text'> & { ref: string })
+  | (Omit<ThinkingContent, 'thinking'> & { ref: string })
+  | (Omit<ToolCall, 'arguments'> & Partial<Pick<ToolCall, 'arguments'>>)
+
+/** What a step's transcript parts do not hold of Pi's response. */
+type StoredResponse = Omit<AssistantMessage, 'content'> & {
+  content: StoredBlock[]
+}
+
+/** The transient record of one model call, for replay. */
+type PiRecord =
+  | { call: 'step'; response: StoredResponse }
+  | { call: 'compaction'; response: AssistantMessage }
+
+type RecordedCall = {
+  kind: 'step' | 'compaction'
+  response: AssistantMessage
+  /** Steering to queue while serving this call. */
+  steers: string[]
+}
+
+type Saved = {
+  cwd: string
+  header: SessionHeader
+  entries: SessionEntry[]
+}
+
+type Completion = { finishReason?: FinishReason; usage?: LanguageModelUsage }
+
+/** The turn being replayed or run live, as the model and tool hooks see it. */
+type TurnState = {
+  plan: HarnessTurn
+  calls: RecordedCall[]
+  next: number
+  liveStep: boolean
+  /** The live step is over; Pi's run is being stopped. */
+  over: boolean
+  failure?: Error
+  /** Resolves when the live step is over: Pi asked for its tools or another step. */
+  stepOver: Deferred
+  /** Live compaction responses, in request order. */
+  compactions: AssistantMessage[]
+}
+
+function endStep(state: TurnState): void {
+  state.over = true
+  state.stepOver.resolve()
+}
+
+const diverged = (message: string): Divergence => new Divergence('Pi', message)
+
+// ---------------------------------------------------------------- step
+
+async function runStep(
+  step: HarnessStep,
+  sandbox: () => Promise<Experimental_SandboxSession>,
+): Promise<HarnessGeneration> {
+  const modules = await loadPi()
+  const credential = gatewayCredential('pi')
+  const name = `${step.agentName}.a2.harness.pi`
+  const { replayAll, ...first } = await readHarness({
+    step,
+    name,
+    parse: savedSession,
+  })
+  let read: HarnessRead<Saved> = first
+  let results = read.results
+  const cwd =
+    read.saved?.cwd ??
+    (await sessionCwd({
+      key: cwdKey(step.agentName, step.sessionId),
+      sandbox,
+    }))
+
+  let completion: Completion = {}
+  let session: AgentSession | undefined
+  let compacting = false
+  let turn: TurnState | undefined
+
+  const fail = (state: TurnState, error: Error): never => {
+    state.failure ??= error
+    endStep(state)
+    void session?.abort()
+    throw state.failure
+  }
+
+  let controller: ReadableStreamDefaultController<UIMessageChunk> | undefined
+  let closed = false
+  const emit = (chunk: UIMessageChunk): void => {
+    if (!closed) controller!.enqueue(chunk)
+  }
+  const record = (data: PiRecord): void =>
+    emit({ type: 'data-a2-pi', data, transient: true })
+
+  // ------------------------------------------------------ model calls
+
+  const answer = async (
+    out: AssistantMessageEventStream,
+    model: Parameters<ModelRuntime['streamSimple']>[0],
+    context: Parameters<ModelRuntime['streamSimple']>[1],
+    options: Parameters<ModelRuntime['streamSimple']>[2],
+  ): Promise<void> => {
+    const state = turn!
+    if (state.failure) throw state.failure
+    // Stopping Pi runs its turn-end hooks; nothing they ask the model counts.
+    if (state.over) {
+      replay(out, interrupted(model))
+      return
+    }
+    const kind = compacting ? 'compaction' : 'step'
+    const recorded = state.calls[state.next]
+    if (recorded !== undefined) {
+      if (recorded.kind !== kind)
+        fail(state, diverged(`call ${state.next} is not a ${kind}`))
+      state.next++
+      // Steering queued while a response streams enters after its tools.
+      for (const text of recorded.steers) await session!.steer(text)
+      replay(out, recorded.response)
+      return
+    }
+    if (state.plan.end === 'stop')
+      fail(state, diverged('Pi continued a turn the log ended'))
+    if (state.plan.end === 'interrupted') {
+      replay(out, interrupted(model))
+      return
+    }
+    if (kind === 'step' && state.liveStep) {
+      // One generation is one step; the next step is A2's next generation.
+      endStep(state)
+      replay(out, interrupted(model))
+      return
+    }
+    if (kind === 'step') state.liveStep = true
+    const upstream = realStream(model, context, {
+      ...options,
+      headers: { ...options?.headers, ...gatewayHeaders(credential) },
+    })
+    const live = kind === 'step' ? liveStep() : undefined
+    for await (const event of upstream) {
+      await live?.(event)
+      if (event.type === 'done' || event.type === 'error') {
+        const response = event.type === 'done' ? event.message : event.error
+        if (kind === 'compaction') state.compactions.push(response)
+        // A2 retries the step; Pi's own recovery would run outside it.
+        if (event.type === 'error' && event.reason === 'error') {
+          state.failure ??= new Error(
+            response.errorMessage ?? 'the Pi model call failed',
+          )
+          endStep(state)
+        }
+      }
+      out.push(event)
+    }
+  }
+
+  /**
+   * A live call's input as Pi's loop, then the AI SDK step, read it. Pi
+   * answers a call to no tool, or whose arguments it cannot prepare or
+   * validate, itself. The agent's tool reads the arguments Pi accepts through
+   * its own schema; A2 runs Pi's own tools from the model's arguments, as
+   * Pi's loop does.
+   */
+  const callInput = async (block: ToolCall): Promise<SchemaInput> => {
+    const definition = tools.find((tool) => tool.name === block.name)
+    if (definition === undefined)
+      return { success: false, errorText: `Tool ${block.name} not found` }
+    let args: unknown
+    try {
+      const prepared = definition.prepareArguments?.(block.arguments)
+      args = modules.ai.validateToolArguments(
+        definition,
+        prepared === undefined
+          ? block
+          : { ...block, arguments: prepared as JsonObject },
+      )
+    } catch (error) {
+      return { success: false, errorText: asError(error).message }
+    }
+    if (builtinOf(step.tools, block.name) !== undefined)
+      return { success: true, value: block.arguments }
+    return schemaInput({ tools: step.tools, toolName: block.name, value: args })
+  }
+
+  /** Streams one live step response as UI chunks and records it. */
+  const liveStep = () => {
+    const ref = (index: number) => `pi:${step.requestIndex}:${index}`
+    const started = new Set<string>()
+    return async (event: AssistantMessageEvent): Promise<void> => {
+      switch (event.type) {
+        case 'start':
+          emit({ type: 'start-step' })
+          break
+        case 'text_start':
+          emit({ type: 'text-start', id: ref(event.contentIndex) })
+          break
+        case 'text_delta':
+          emit({
+            type: 'text-delta',
+            id: ref(event.contentIndex),
+            delta: event.delta,
+          })
+          break
+        case 'text_end':
+          emit({ type: 'text-end', id: ref(event.contentIndex) })
+          break
+        case 'thinking_start':
+          emit({ type: 'reasoning-start', id: ref(event.contentIndex) })
+          break
+        case 'thinking_delta':
+          emit({
+            type: 'reasoning-delta',
+            id: ref(event.contentIndex),
+            delta: event.delta,
+          })
+          break
+        case 'thinking_end':
+          emit({ type: 'reasoning-end', id: ref(event.contentIndex) })
+          break
+        case 'toolcall_start':
+        case 'toolcall_delta': {
+          const call = event.partial.content[event.contentIndex] as ToolCall
+          if (!call.id) break
+          if (!started.has(call.id)) {
+            started.add(call.id)
+            emit({
+              type: 'tool-input-start',
+              toolCallId: call.id,
+              toolName: call.name,
+            })
+          }
+          if (event.type === 'toolcall_delta')
+            emit({
+              type: 'tool-input-delta',
+              toolCallId: call.id,
+              inputTextDelta: event.delta,
+            })
+          break
+        }
+        case 'done':
+        case 'error': {
+          const response = event.type === 'done' ? event.message : event.error
+          // Calls whose transcript input is not the model's arguments.
+          const changed = new Set<string>()
+          if (event.type === 'done') {
+            for (const block of response.content) {
+              if (block.type !== 'toolCall') continue
+              if (!started.has(block.id))
+                emit({
+                  type: 'tool-input-start',
+                  toolCallId: block.id,
+                  toolName: block.name,
+                })
+              // Pi does not run the calls of a cut-off response.
+              const input: SchemaInput =
+                response.stopReason === 'length'
+                  ? {
+                      success: false,
+                      errorText:
+                        'The response hit the output token limit before this call was complete.',
+                    }
+                  : await callInput(block)
+              if (
+                input.success &&
+                JSON.stringify(input.value) !== JSON.stringify(block.arguments)
+              )
+                changed.add(block.id)
+              emit(
+                input.success
+                  ? {
+                      type: 'tool-input-available',
+                      toolCallId: block.id,
+                      toolName: block.name,
+                      input: input.value,
+                    }
+                  : {
+                      type: 'tool-input-error',
+                      toolCallId: block.id,
+                      toolName: block.name,
+                      input: block.arguments,
+                      errorText: input.errorText,
+                    },
+              )
+            }
+          }
+          record({
+            call: 'step',
+            response: storedResponse(response, ref, changed),
+          })
+          completion = {
+            finishReason: finishReason(response.stopReason),
+            usage: modelUsage(response.usage),
+          }
+          break
+        }
+      }
+    }
+  }
+
+  // ----------------------------------------------------- tool results
+
+  const toolResult = async (
+    toolCallId: string,
+    signal: AbortSignal | undefined,
+    shape: 'pi' | 'tool',
+  ): Promise<AgentToolResult<unknown>> => {
+    const state = turn!
+    if (state.failure) throw state.failure
+    const result = results.get(toolCallId)
+    if (result !== undefined) return answeredTool(result, shape)
+    if (state.plan.end === 'live' && state.liveStep) {
+      // Pi now waits on tools that A2's tool handlers run.
+      endStep(state)
+      return pending(signal)
+    }
+    if (state.plan.end === 'interrupted') {
+      void session!.abort()
+      return pending(signal)
+    }
+    return fail(state, diverged(`no recorded result for ${toolCallId}`))
+  }
+
+  // ------------------------------------------------------------ run
+
+  const runtime = await modules.agent.ModelRuntime.create({
+    credentials: new modules.ai.InMemoryCredentialStore(),
+    modelsPath: null,
+    allowModelNetwork: false,
+    refreshOnCreate: false,
+  })
+  await runtime.setRuntimeApiKey(GATEWAY, credential.token)
+  const model = runtime.getModel(GATEWAY, step.model)
+  if (model === undefined)
+    throw new Error(`Pi does not know the Gateway model ${step.model}`)
+  const realStream = runtime.streamSimple.bind(runtime)
+  runtime.streamSimple = (requested, context, options?) => {
+    const out = modules.ai.createAssistantMessageEventStream()
+    answer(out, requested, context, options).catch((error: unknown) => {
+      const state = turn
+      if (state) state.failure ??= asError(error)
+      replay(out, { ...interrupted(requested), stopReason: 'error' })
+    })
+    return out
+  }
+
+  const tools = await sessionTools(modules, step.tools, cwd, toolResult)
+
+  const createSession = async (from: Saved | undefined) => {
+    const sessionManager =
+      from === undefined
+        ? modules.agent.SessionManager.inMemory(cwd, {
+            id: sessionUuid(step.agentName, step.sessionId),
+          })
+        : modules.agent.SessionManager.inMemory(cwd, { id: from.header.id }, [
+            from.header,
+            ...from.entries,
+          ])
+    const created = await modules.agent.createAgentSession({
+      cwd,
+      // Pi reads nothing there: every resource and manager comes from A2.
+      agentDir: join(tmpdir(), 'a2-pi'),
+      model,
+      ...thinkingLevel(step),
+      modelRuntime: runtime,
+      resourceLoader: resources(modules, step.instructions),
+      sessionManager,
+      settingsManager: modules.agent.SettingsManager.inMemory({
+        // A2 retries failed steps.
+        retry: { enabled: false, maxRetries: 0, baseDelayMs: 0 },
+        cacheWarming: 'off',
+      }),
+      customTools: tools,
+      tools: tools.map((definition) => definition.name),
+    })
+    created.session.subscribe((event) => {
+      if (event.type === 'compaction_start') compacting = true
+      if (event.type !== 'compaction_end') return
+      compacting = false
+      const state = turn
+      if (state === undefined || state.over || state.compactions.length === 0)
+        return
+      for (const response of state.compactions.splice(0))
+        record({ call: 'compaction', response })
+      emit({
+        type: 'data-a2-compaction',
+        id: `pi-compaction:${step.requestIndex}:${compactions++}`,
+        data: compactionData(event),
+      })
+    })
+    return created.session
+  }
+  let compactions = 0
+
+  const saveSnapshot = async (): Promise<void> => {
+    const manager = session!.sessionManager
+    const header = manager.getHeader()
+    const entries = manager.getEntries()
+    const last = entries.findLast(
+      (entry) => entry.type === 'message' && entry.message.role === 'assistant',
+    )
+    const stop =
+      last?.type === 'message' && last.message.role === 'assistant'
+        ? last.message.stopReason
+        : undefined
+    // Only a turn this attempt finished describes the log.
+    if (
+      header === null ||
+      step.signal.aborted ||
+      stop === 'aborted' ||
+      stop === 'error'
+    )
+      return
+    await saveHarness({
+      step,
+      name,
+      state: { version: SNAPSHOT_VERSION, cwd, header, entries },
+    })
+  }
+
+  const run = async (): Promise<void> => {
+    emit({ type: 'start', messageId: step.responseMessageId })
+    try {
+      session = await createSession(read.saved)
+    } catch (error) {
+      if (read.saved === undefined) throw error
+      // An unusable snapshot: replay the whole log instead.
+      read = await replayAll()
+      results = read.results
+      session = await createSession(undefined)
+    }
+    for (const plan of read.turns) {
+      if (step.signal.aborted) return
+      const state: TurnState = {
+        plan,
+        calls: recordedTurn(plan),
+        next: 0,
+        liveStep: false,
+        over: false,
+        stepOver: deferred(),
+        compactions: [],
+      }
+      turn = state
+      const running = session.prompt(plan.prompt, {
+        expandPromptTemplates: false,
+      })
+      if (plan.end !== 'live') {
+        await running.catch((error: unknown) => {
+          throw state.failure ?? error
+        })
+        if (state.failure) throw state.failure
+        continue
+      }
+      const ended = await Promise.race([
+        running.then(() => true),
+        state.stepOver.promise.then(() => false),
+      ]).catch((error: unknown) => {
+        throw state.failure ?? error
+      })
+      if (!ended) {
+        await session.abort()
+        await running.catch(() => {})
+      }
+      if (state.failure) throw state.failure
+      if (ended) await saveSnapshot()
+      emit({ type: 'finish-step' })
+      emit({
+        type: 'finish',
+        ...(completion.finishReason === undefined
+          ? {}
+          : { finishReason: completion.finishReason }),
+      })
+    }
+  }
+
+  const stream = new ReadableStream<UIMessageChunk>({
+    start(streamController) {
+      controller = streamController
+      const abort = () => void session?.abort()
+      step.signal.addEventListener('abort', abort, { once: true })
+      void (async () => {
+        try {
+          await run()
+          closed = true
+          streamController.close()
+        } catch (error) {
+          closed = true
+          streamController.error(
+            step.signal.aborted ? step.signal.reason : error,
+          )
+        } finally {
+          step.signal.removeEventListener('abort', abort)
+          session?.dispose()
+        }
+      })()
+    },
+    cancel() {
+      closed = true
+      void session?.abort()
+    },
+  })
+
+  return { stream, completion: () => completion }
+}
+
+// ------------------------------------------------------------ the plan
+
+/** A turn's recorded model calls, in the order Pi made them. */
+function recordedTurn(turn: HarnessTurn): RecordedCall[] {
+  const calls: RecordedCall[] = []
+  turn.generations.forEach((generation, index) => {
+    calls.push(...recordedCalls(generation))
+    const steers = turn.steers.get(index)
+    const last = calls.findLast((call) => call.kind === 'step')
+    if (steers !== undefined && last !== undefined) last.steers.push(...steers)
+  })
+  return calls
+}
+
+function recordedCalls(generation: RecordedGeneration): RecordedCall[] {
+  const texts = new Map<string, string>()
+  const inputs = new Map<string, unknown>()
+  const calls: RecordedCall[] = []
+  for (const chunk of generation.chunks) {
+    switch (chunk.type) {
+      case 'text-delta':
+      case 'reasoning-delta':
+        texts.set(chunk.id, (texts.get(chunk.id) ?? '') + chunk.delta)
+        break
+      case 'tool-input-available':
+      case 'tool-input-error':
+        inputs.set(chunk.toolCallId, chunk.input)
+        break
+      case 'data-a2-pi': {
+        const record = chunk.data as PiRecord
+        calls.push({
+          kind: record.call,
+          response:
+            record.call === 'step'
+              ? restoredResponse(record.response, texts, inputs)
+              : record.response,
+          steers: [],
+        })
+        break
+      }
+    }
+  }
+  if (!calls.some((call) => call.kind === 'step'))
+    throw diverged('a completed generation has no recorded response')
+  return calls
+}
+
+function storedResponse(
+  response: AssistantMessage,
+  ref: (index: number) => string,
+  changed: ReadonlySet<string>,
+): StoredResponse {
+  return {
+    ...response,
+    content: response.content.map((block, index): StoredBlock => {
+      switch (block.type) {
+        case 'text': {
+          const { text: _, ...rest } = block
+          return { ...rest, ref: ref(index) }
+        }
+        case 'thinking': {
+          const { thinking: _, ...rest } = block
+          return { ...rest, ref: ref(index) }
+        }
+        case 'toolCall': {
+          if (changed.has(block.id)) return block
+          const { arguments: _, ...rest } = block
+          return rest
+        }
+      }
+    }),
+  }
+}
+
+function restoredResponse(
+  stored: StoredResponse,
+  texts: ReadonlyMap<string, string>,
+  inputs: ReadonlyMap<string, unknown>,
+): AssistantMessage {
+  return {
+    ...stored,
+    content: stored.content.map((block) => {
+      if (block.type === 'toolCall')
+        return {
+          ...block,
+          arguments:
+            block.arguments ?? ((inputs.get(block.id) ?? {}) as JsonObject),
+        }
+      const { ref, ...rest } = block
+      const text = texts.get(ref) ?? ''
+      return rest.type === 'text'
+        ? { ...rest, text }
+        : { ...rest, thinking: text }
+    }),
+  }
+}
+
+/** Hands a recorded response to Pi's loop. */
+function replay(out: AssistantMessageEventStream, response: AssistantMessage) {
+  out.push({ type: 'start', partial: response })
+  const reason = response.stopReason
+  if (reason === 'error' || reason === 'aborted')
+    out.push({ type: 'error', reason, error: response })
+  else if (reason === 'pending')
+    out.push({ type: 'done', reason: 'stop', message: response })
+  else out.push({ type: 'done', reason, message: response })
+}
+
+/** The response of a model call that A2 interrupted. */
+function interrupted(
+  model: Parameters<ModelRuntime['streamSimple']>[0],
+): AssistantMessage {
+  return {
+    role: 'assistant',
+    content: [],
+    api: model.api,
+    provider: model.provider,
+    model: model.id,
+    usage: {
+      input: 0,
+      output: 0,
+      cacheRead: 0,
+      cacheWrite: 0,
+      totalTokens: 0,
+      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+    },
+    stopReason: 'aborted',
+    timestamp: Date.now(),
+  }
+}
+
+// ----------------------------------------------------------- the tools
+
+const builtinTools: ToolSet = Object.fromEntries(
+  BUILTINS.map((name): [string, Tool] => [
+    name,
+    {
+      inputSchema: jsonSchema(
+        async () =>
+          builtinDefinition(await loadPi(), name, '/')
+            .parameters as JSONSchema7,
+      ),
+      execute: (input: unknown, options: ToolExecutionOptions<unknown>) =>
+        runBuiltin(name, input, options),
+    },
+  ]),
+)
+
+/** Runs one of Pi's tools in the session's sandbox, as Pi's loop would. */
+async function runBuiltin(
+  name: Builtin,
+  input: unknown,
+  options: ToolExecutionOptions<unknown>,
+): Promise<AgentToolResult<unknown>> {
+  const sandbox = options.experimental_sandbox
+  if (sandbox === undefined)
+    throw new Error(`pi() runs ${name} in the agent's sandbox`)
+  const modules = await loadPi()
+  const scope = ambientToolScopeStorage()?.getStore()
+  const event = (scope?.context as { event?: { sessionId?: unknown } })?.event
+  const cwd = await sessionCwd({
+    key:
+      scope?.contract !== undefined && typeof event?.sessionId === 'string'
+        ? cwdKey(scope.contract.name, event.sessionId)
+        : undefined,
+    sandbox: async () => sandbox,
+  })
+  const definition = builtinDefinition(modules, name, cwd, {
+    sandbox,
+    signal: options.abortSignal,
+  })
+  const args: unknown = modules.ai.validateToolArguments(definition, {
+    type: 'toolCall',
+    id: options.toolCallId,
+    name,
+    arguments: (definition.prepareArguments?.(input) ?? input) as JsonObject,
+  })
+  // Pi's bash keeps a long output's full text in a local file, named in its
+  // updates and its result. It is removed however the call ends.
+  let outputFile: string | undefined
+  try {
+    const result = await definition.execute(
+      options.toolCallId,
+      args,
+      options.abortSignal,
+      (update) => {
+        outputFile ??= fullOutputPath(update)
+      },
+      // Pi's tools take their context from the session, which A2 runs elsewhere.
+      undefined as never,
+    )
+    outputFile ??= fullOutputPath(result)
+    return await recordedResult(result, sandbox)
+  } catch (error) {
+    // Output whose last bytes tip it over Pi's limit opens the file only as
+    // it ends, after the last update; a failed call then names it in its error.
+    outputFile ??= namedOutputFile(error)
+    throw error
+  } finally {
+    if (outputFile !== undefined) await unlink(outputFile).catch(() => {})
+  }
+}
+
+type BuiltinDetails =
+  { truncation?: { content?: unknown }; fullOutputPath?: unknown } | undefined
+
+const fullOutputPath = (
+  result: AgentToolResult<unknown>,
+): string | undefined => {
+  const path = (result.details as BuiltinDetails)?.fullOutputPath
+  return typeof path === 'string' ? path : undefined
+}
+
+/** The full-output file an aborted or timed-out bash call's error names. */
+function namedOutputFile(error: unknown): string | undefined {
+  if (!(error instanceof Error)) return undefined
+  const path = /Full output: ([^\n]+)\]\n\nCommand [^\n]+$/.exec(
+    error.message,
+  )?.[1]
+  // The command's own output precedes Pi's note, and could name any file.
+  return path !== undefined &&
+    dirname(path) === tmpdir() &&
+    /^pi-bash-[0-9a-f]{16}\.log$/.test(basename(path))
+    ? path
+    : undefined
+}
+
+/**
+ * What the log keeps of a result: what Pi's tool-result message holds. A
+ * long output's full text moves to the path its result names in the sandbox,
+ * where the model's `read` finds it.
+ */
+async function recordedResult(
+  result: AgentToolResult<unknown>,
+  sandbox: Experimental_SandboxSession,
+): Promise<AgentToolResult<unknown>> {
+  // Pi's tool-result message does not carry structured content.
+  const { structuredContent: _, ...kept } =
+    result as AgentToolResult<unknown> & {
+      structuredContent?: unknown
+    }
+  const details = kept.details as BuiltinDetails
+  const path = fullOutputPath(kept)
+  if (path !== undefined) {
+    try {
+      await sandbox.writeBinaryFile({
+        path,
+        content: await readLocalFile(path),
+      })
+    } catch {
+      // The result stands without it; the model can rerun the command.
+    }
+  }
+  if (details?.truncation === undefined) return kept
+  // The truncated text is the result's content already.
+  const { content: __, ...truncation } = details.truncation
+  return { ...kept, details: { ...details, truncation } }
+}
+
+/** The Pi tool a name in the agent's tool set runs, unless the agent replaced it. */
+const builtinOf = (tools: ToolSet, name: string): Builtin | undefined =>
+  BUILTINS.find(
+    (candidate) =>
+      candidate === name &&
+      Object.hasOwn(tools, name) &&
+      tools[name] === builtinTools[name],
+  )
+
+// Each definition is typed by its parameters; a session takes any of them.
+const untyped = (definition: object) => definition as ToolDefinition
+
+function builtinDefinition(
+  modules: PiModules,
+  name: Builtin,
+  cwd: string,
+  call?: {
+    sandbox: Experimental_SandboxSession
+    signal: AbortSignal | undefined
+  },
+): ToolDefinition {
+  const sandbox = call?.sandbox
+  const files =
+    call === undefined ? undefined : { operations: sandboxFiles(call) }
+  switch (name) {
+    case 'read':
+      return untyped(modules.agent.createReadToolDefinition(cwd, files))
+    case 'write':
+      return untyped(modules.agent.createWriteToolDefinition(cwd, files))
+    case 'edit':
+      return untyped(modules.agent.createEditToolDefinition(cwd, files))
+    case 'bash':
+      return untyped(
+        modules.agent.createBashToolDefinition(cwd, {
+          // No PI_* variables reach the sandbox, so Pi's prompt names none.
+          exposeSessionEnvironment: false,
+          ...(sandbox === undefined
+            ? {}
+            : { operations: sandboxBash(sandbox) }),
+        }),
+      )
+  }
+}
+
+/**
+ * The tools Pi's session declares: Pi's own, and the agent's. Their results
+ * come from A2's log; A2's tool handlers ran them.
+ */
+async function sessionTools(
+  modules: PiModules,
+  tools: ToolSet,
+  cwd: string,
+  result: (
+    toolCallId: string,
+    signal: AbortSignal | undefined,
+    shape: 'pi' | 'tool',
+  ) => Promise<AgentToolResult<unknown>>,
+): Promise<ToolDefinition[]> {
+  return Promise.all(
+    Object.entries(tools)
+      .filter(([, tool]) => tool.type !== 'provider')
+      .map(async ([name, tool]): Promise<ToolDefinition> => {
+        const builtin = builtinOf(tools, name)
+        if (builtin !== undefined)
+          // Pi's own definition, for its prompt; A2 already ran the call.
+          return Object.assign(builtinDefinition(modules, builtin, cwd), {
+            execute: (
+              toolCallId: string,
+              _params: unknown,
+              signal?: AbortSignal,
+            ) => result(toolCallId, signal, 'pi'),
+          })
+        return {
+          name,
+          label: name,
+          description:
+            typeof tool.description === 'string' ? tool.description : '',
+          parameters: (await asSchema(tool.inputSchema)
+            .jsonSchema) as ToolDefinition['parameters'],
+          execute: (toolCallId, _params, signal) =>
+            result(toolCallId, signal, 'tool'),
+        }
+      }),
+  )
+}
+
+function answeredTool(
+  result: ToolResultPayload,
+  shape: 'pi' | 'tool',
+): AgentToolResult<unknown> {
+  if (result.error !== undefined) throw new Error(result.error)
+  if (result.denied === true) throw new Error('The tool call was denied.')
+  if (shape === 'pi') return result.output as AgentToolResult<unknown>
+  const output = result.output
+  return {
+    content: [
+      {
+        type: 'text',
+        text:
+          typeof output === 'string' ? output : (JSON.stringify(output) ?? ''),
+      },
+    ],
+    details: undefined,
+  }
+}
+
+// ------------------------------------------------------------ sandbox
+
+/** Pi's file operations, against the sandbox. */
+function sandboxFiles({
+  sandbox,
+  signal,
+}: {
+  sandbox: Experimental_SandboxSession
+  /** The tool call's: it stops a read or write in flight. */
+  signal: AbortSignal | undefined
+}) {
+  const abortSignal = signal === undefined ? {} : { abortSignal: signal }
+  let checked: { path: string; bytes: Buffer } | undefined
+  const readFile = async (path: string): Promise<Buffer> => {
+    if (checked?.path === path) {
+      const { bytes } = checked
+      checked = undefined
+      return bytes
+    }
+    const stream = await sandbox.readFile({ path, ...abortSignal })
+    if (stream === null)
+      throw Object.assign(
+        new Error(`ENOENT: no such file or directory, open '${path}'`),
+        { code: 'ENOENT' },
+      )
+    const chunks: Uint8Array[] = []
+    let size = 0
+    // Leaving the loop cancels the stream.
+    for await (const chunk of stream) {
+      size += chunk.byteLength
+      if (size > MAX_TRANSFER_BYTES)
+        throw new Error(
+          `${path} is larger than ${MAX_TRANSFER}. Read or change parts of it with bash, using sed, head, or grep.`,
+        )
+      chunks.push(chunk)
+    }
+    return Buffer.concat(chunks)
+  }
+  return {
+    readFile,
+    // Pi checks a file just before reading it; the check's read serves both.
+    access: async (path: string): Promise<void> => {
+      checked = { path, bytes: await readFile(path) }
+    },
+    writeFile: async (path: string, content: string): Promise<void> => {
+      checked = undefined
+      await sandbox.writeTextFile({ path, content, ...abortSignal })
+    },
+    // The sandbox creates parent directories as it writes.
+    mkdir: async (): Promise<void> => {},
+  }
+}
+
+/** Pi's shell operations, against the sandbox. */
+const sandboxBash = (sandbox: Experimental_SandboxSession): BashOperations => ({
+  // Pi hands over the server's environment; it never reaches the sandbox.
+  exec: (command, cwd, { onData, signal, timeout }) =>
+    sandboxExec({
+      sandbox,
+      command,
+      workingDirectory: cwd,
+      onData,
+      signal,
+      timeout,
+    }),
+})
+
+/**
+ * Runs a command in the sandbox with Pi's error contract. At most
+ * `MAX_TRANSFER_BYTES` of its output reach the server; past that, the
+ * command is stopped.
+ */
+async function sandboxExec({
+  sandbox,
+  command,
+  workingDirectory,
+  onData,
+  signal,
+  timeout,
+}: {
+  sandbox: Experimental_SandboxSession
+  command: string
+  /** The sandbox's own when absent. */
+  workingDirectory?: string
+  onData: (data: Buffer, stream: 'stdout' | 'stderr') => void
+  signal?: AbortSignal | undefined
+  /** In seconds. */
+  timeout?: number | undefined
+}): Promise<{ exitCode: number }> {
+  if (signal?.aborted) throw new Error('aborted')
+  const stop = new AbortController()
+  const abort = () => stop.abort()
+  signal?.addEventListener('abort', abort, { once: true })
+  let timedOut = false
+  const timer =
+    timeout === undefined || timeout <= 0
+      ? undefined
+      : setTimeout(() => {
+          timedOut = true
+          stop.abort()
+        }, timeout * 1000)
+  let size = 0
+  let overflowed = false
+  const overflow = (options?: ErrorOptions) =>
+    new Error(
+      `The command's output passed ${MAX_TRANSFER}, so it was stopped. Write long output to a file and read parts of it with head, tail, or grep.`,
+      options,
+    )
+  let exitCode: number
+  try {
+    const process = await sandbox.spawn({
+      command,
+      ...(workingDirectory === undefined ? {} : { workingDirectory }),
+      abortSignal: stop.signal,
+    })
+    const pipe = async (stream: 'stdout' | 'stderr') => {
+      for await (const chunk of process[stream]) {
+        size += chunk.byteLength
+        if (size > MAX_TRANSFER_BYTES) {
+          overflowed = true
+          stop.abort()
+          return
+        }
+        onData(Buffer.from(chunk), stream)
+      }
+    }
+    const [exited] = await Promise.all([
+      process.wait(),
+      pipe('stdout'),
+      pipe('stderr'),
+    ])
+    exitCode = exited.exitCode
+  } catch (error) {
+    if (signal?.aborted) throw new Error('aborted', { cause: error })
+    if (timedOut) throw new Error(`timeout:${timeout}`, { cause: error })
+    if (overflowed) throw overflow({ cause: error })
+    throw error
+  } finally {
+    clearTimeout(timer)
+    signal?.removeEventListener('abort', abort)
+  }
+  if (signal?.aborted) throw new Error('aborted')
+  if (timedOut) throw new Error(`timeout:${timeout}`)
+  if (overflowed) throw overflow()
+  return { exitCode }
+}
+
+// Sandbox resolvers may return a new object each time; sessions keep theirs.
+const cwds = boundedCache<string, Promise<string>>({
+  maxEntries: 1000,
+  maxBytes: 1000,
+})
+
+const cwdKey = (agentName: string, sessionId: string): string =>
+  `${agentName}\u001f${sessionId}`
+
+/** A session's working directory in its sandbox: Pi's cwd. */
+function sessionCwd({
+  key,
+  sandbox,
+}: {
+  /** The agent and session, when known; otherwise the sandbox is asked. */
+  key: string | undefined
+  sandbox: () => Promise<Experimental_SandboxSession>
+}): Promise<string> {
+  const cached = key === undefined ? undefined : cwds.get(key)
+  if (cached !== undefined) return cached
+  const cwd = sandbox().then(async (resolved) => {
+    // Bounded, unlike `run`: a login profile in the sandbox can print anything.
+    const output = { stdout: '', stderr: '' }
+    const decoder = { stdout: new TextDecoder(), stderr: new TextDecoder() }
+    const { exitCode } = await sandboxExec({
+      sandbox: resolved,
+      command: 'pwd',
+      onData: (data, stream) => {
+        output[stream] += decoder[stream].decode(data, { stream: true })
+      },
+    }).catch((error: unknown) => {
+      throw new Error("pi() could not read the sandbox's working directory", {
+        cause: error,
+      })
+    })
+    const path = output.stdout.trim()
+    if (exitCode !== 0 || !path.startsWith('/'))
+      throw new Error(
+        `pi() could not read the sandbox's working directory: ${(output.stderr || output.stdout).slice(0, 1000)}`,
+      )
+    return path
+  })
+  if (key !== undefined) {
+    cwds.set(key, cwd, 1)
+    cwd.catch(() => cwds.delete(key))
+  }
+  return cwd
+}
+
+// ------------------------------------------------------------- Pi
+
+let loading: Promise<PiModules> | undefined
+
+function loadPi(): Promise<PiModules> {
+  loading ??= Promise.all([
+    import('@earendil-works/pi-coding-agent'),
+    import('@earendil-works/pi-ai'),
+  ]).then(
+    ([agent, ai]) => ({ agent, ai }),
+    (error: unknown) => {
+      loading = undefined
+      throw new Error(
+        'pi() needs @earendil-works/pi-coding-agent and @earendil-works/pi-ai: install them',
+        { cause: error },
+      )
+    },
+  )
+  return loading
+}
+
+/** Pi's resources: nothing discovered from disk, the agent's instructions appended. */
+function resources(modules: PiModules, instructions: string[]): ResourceLoader {
+  const runtime = modules.agent.createExtensionRuntime()
+  return {
+    getExtensions: () => ({ extensions: [], errors: [], runtime }),
+    getSkills: () => ({ skills: [], diagnostics: [] }),
+    getPrompts: () => ({ prompts: [], diagnostics: [] }),
+    getThemes: () => ({ themes: [], diagnostics: [] }),
+    getAgentsFiles: () => ({ agentsFiles: [] }),
+    getSystemPrompt: () => undefined,
+    getSystemPromptSource: () => undefined,
+    getAppendSystemPrompt: () => instructions,
+    getAppendSystemPromptSources: () => [],
+    extendResources: () => {},
+    reload: async () => {},
+  }
+}
+
+function thinkingLevel(step: HarnessStep): {
+  thinkingLevel?: 'off' | 'minimal' | 'low' | 'medium' | 'high' | 'xhigh'
+} {
+  if (step.reasoning === undefined || step.reasoning === 'provider-default')
+    return {}
+  return { thinkingLevel: step.reasoning === 'none' ? 'off' : step.reasoning }
+}
+
+function gatewayHeaders(credential: {
+  token: string
+  method: string
+}): Record<string, string> {
+  return {
+    authorization: `Bearer ${credential.token}`,
+    'ai-gateway-auth-method': credential.method,
+  }
+}
+
+/** A stable Pi session ID for an A2 session, so Pi's cache affinity holds. */
+function sessionUuid(agentName: string, sessionId: string): string {
+  const hex = createHash('sha256')
+    .update(`${agentName}\u001f${sessionId}`)
+    .digest('hex')
+  return `${hex.slice(0, 8)}-${hex.slice(8, 12)}-8${hex.slice(13, 16)}-${hex.slice(16, 20)}-${hex.slice(20, 32)}`
+}
+
+function savedSession(state: unknown): Saved | undefined {
+  if (typeof state !== 'object' || state === null) return undefined
+  const { version, cwd, header, entries } = state as Record<string, unknown>
+  if (
+    version !== SNAPSHOT_VERSION ||
+    typeof cwd !== 'string' ||
+    typeof header !== 'object' ||
+    header === null ||
+    !Array.isArray(entries)
+  )
+    return undefined
+  return {
+    cwd,
+    header: header as SessionHeader,
+    entries: entries as SessionEntry[],
+  }
+}
+
+function compactionData(event: {
+  result: { summary: string; tokensBefore: number; usage?: Usage } | undefined
+  aborted: boolean
+  errorMessage?: string
+}): CompactionPartData {
+  if (event.result === undefined)
+    return { status: event.aborted ? 'interrupted' : 'failed' }
+  return {
+    status: 'completed',
+    inputTokens: event.result.tokensBefore,
+    summary: event.result.summary,
+    ...(event.result.usage === undefined
+      ? {}
+      : { usage: modelUsage(event.result.usage) }),
+  }
+}
+
+function finishReason(reason: AssistantMessage['stopReason']): FinishReason {
+  switch (reason) {
+    case 'stop':
+      return 'stop'
+    case 'length':
+      return 'length'
+    case 'toolUse':
+      return 'tool-calls'
+    case 'error':
+      return 'error'
+    default:
+      return 'other'
+  }
+}
+
+function modelUsage(usage: Usage): LanguageModelUsage {
+  return {
+    inputTokens: usage.input + usage.cacheRead + usage.cacheWrite,
+    inputTokenDetails: {
+      noCacheTokens: usage.input,
+      cacheReadTokens: usage.cacheRead,
+      cacheWriteTokens: usage.cacheWrite,
+    },
+    outputTokens: usage.output,
+    outputTokenDetails: {
+      textTokens:
+        usage.reasoning === undefined
+          ? undefined
+          : usage.output - usage.reasoning,
+      reasoningTokens: usage.reasoning,
+    },
+    totalTokens: usage.totalTokens,
+  }
+}
+
+// ------------------------------------------------------------- helpers
+
+function asError(error: unknown): Error {
+  return error instanceof Error ? error : new Error(String(error))
+}