experimental-a2 0.21.1 → 0.23.0

package/src/ai-fx.ts

@@ -34,18 +34,35 @@ import {
 } from './ai-harness.ts'
 import type { CompactionPartData, GenerationCompletedPayload } from './ai.ts'
 import { schemaInput } from './ai-tool-input.ts'
+import { HARNESS_PARKING } from './internal.ts'
 import { boundedCache } from './bounded-cache.ts'
 
 export type { Harness } from './ai-harness.ts'
 
 /**
  * The FX harness. Every model call FX makes is one durable step; A2 rebuilds
- * the kernel for each step from its latest checkpoint and the log since.
+ * the kernel for each step from its latest checkpoint and the log since,
+ * unless this process holds the kernel the step before left waiting on its
+ * tool calls.
  */
 export function fx(): Harness<'fx', never> {
+  // Per agent, as its saved states are: sessions of two agents never share one.
+  const lots = new WeakMap<object, ParkingLot>()
+  const lotOf = (scope: object): ParkingLot => {
+    let lot = lots.get(scope)
+    if (lot === undefined) {
+      lot = parkingLot()
+      lots.set(scope, lot)
+    }
+    return lot
+  }
   return defineHarness('fx', {
     tools: {},
-    step: runStep,
+    step: (step) => runStep(step, lotOf(step.scope)),
+    release: ({ scope, sessionId }) => {
+      const entry = lots.get(scope)?.take(sessionId)
+      if (entry !== undefined) closeParked(entry)
+    },
   })
 }
 
@@ -53,6 +70,8 @@ const SNAPSHOT_VERSION = 1
 const CATALOG_TTL_MS = 60 * 60 * 1000
 /** The Gateway model catalog libfx reads for each new kernel. */
 const CATALOG_PATH = '/coding-agent/v1/models'
+/** Times a step rebuilds its kernel for a stopped turn FX failed on replay. */
+const REPLAY_REBUILDS = 3
 
 type StreamPart = { type: string } & Record<string, unknown>
 /** A recorded tool call, with the input the kernel reads from its text. */
@@ -72,6 +91,8 @@ type TurnPlan = {
   /** Steering to queue while serving the call at that index. */
   steers: Map<number, string[]>
   end: 'stop' | 'interrupted' | 'live'
+  /** The attempt that made its last committed model call. */
+  generationId?: string
 }
 
 type Completion = { finishReason?: FinishReason; usage?: LanguageModelUsage }
@@ -92,6 +113,12 @@ type TurnState = {
   parsed?: Promise<void>
 }
 
+/** A model request as the kernel sends it to the Gateway. */
+type GatewayRequest = {
+  toolChoice?: { type?: string }
+  providerOptions?: Record<string, Record<string, unknown> | undefined>
+}
+
 function fail(state: TurnState, error: Error): never {
   state.failure ??= error
   state.fxTurn.cancel()
@@ -100,12 +127,108 @@ function fail(state: TurnState, error: Error): never {
 
 const diverged = (message: string): Divergence => new Divergence('FX', message)
 
+// ------------------------------------------------------------- kernels
+
+/** A tool call the kernel waits on while A2's tool consumer runs it. */
+type Waiting = {
+  toolName: string
+  input: unknown
+  resolve: (output: unknown) => void
+  reject: (error: unknown) => void
+}
+
+/** What a kernel's fetch and tools call: the step driving it now. */
+type KernelHooks = {
+  fetch: (
+    input: string | URL | Request,
+    init?: RequestInit,
+  ) => Promise<Response>
+  tool: (toolName: string, input: unknown) => unknown
+}
+
+type Kernel = {
+  agent: FxAgent
+  hooks: KernelHooks
+  waiting: Waiting[]
+  /** Whether FX answered the kernel's latest prompt with an error. */
+  failed: boolean
+}
+
+/**
+ * A kernel a step left waiting on its tool calls. The session's next step in
+ * this process, in the same response or one that steered it, continues it
+ * with their results instead of rebuilding one from the log.
+ */
+type Parked = {
+  kernel: Kernel
+  fxTurn: FxTurn
+  drained: Promise<void>
+  /** The attempt whose model call the kernel made. */
+  generationId: string
+  settings: string
+}
+
+type ParkingLot = {
+  /** Removes the session's parked kernel; the caller then owns it. */
+  take(sessionId: string): Parked | undefined
+  park(sessionId: string, entry: Parked): void
+}
+
+function closeParked(entry: Parked): void {
+  entry.fxTurn.cancel()
+  void entry.kernel.agent.close().catch(() => {})
+}
+
+/**
+ * One agent's parked kernels: one per session, at most `maxKernels`. Each is
+ * closed when its session has nothing left to run in this process (the
+ * agent's release), when its turn ends, or when another takes its place.
+ */
+function parkingLot(): ParkingLot {
+  const held = new Map<string, Parked>()
+  const take = (sessionId: string): Parked | undefined => {
+    const entry = held.get(sessionId)
+    held.delete(sessionId)
+    return entry
+  }
+  return {
+    take,
+    park(sessionId, entry) {
+      const replaced = take(sessionId)
+      if (replaced !== undefined) closeParked(replaced)
+      held.set(sessionId, entry)
+      // A turn that ends while it waits, as a cancelled one does, is done.
+      const ended = () => {
+        if (held.get(sessionId) !== entry) return
+        held.delete(sessionId)
+        closeParked(entry)
+      }
+      entry.fxTurn.result.then(ended, ended)
+      const oldest = held.keys().next().value
+      if (held.size > HARNESS_PARKING.maxKernels && oldest !== undefined) {
+        const evicted = take(oldest)
+        if (evicted !== undefined) closeParked(evicted)
+      }
+    },
+  }
+}
+
 // ---------------------------------------------------------------- step
 
-async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
+async function runStep(
+  step: HarnessStep,
+  lot: ParkingLot,
+): Promise<HarnessGeneration> {
   const { createFxAgent } = await loadLibfx()
   const credential = gatewayCredential('fx')
   const name = `${step.agentName}.a2.harness.fx`
+  const definitions = await toolDefinitions(step.tools)
+  // What the kernel is made with: one made otherwise can't continue the step.
+  const settings = JSON.stringify([
+    kernelModel(step),
+    step.instructions,
+    definitions,
+  ])
   const { replayAll, ...first } = await readHarness({
     step,
     name,
@@ -119,17 +242,29 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
   // has read a response's finish, while A2 may still be reading its copy.
   const released = new AbortController()
   const network = AbortSignal.any([step.signal, released.signal])
-  let agent: FxAgent | undefined
+  let kernel: Kernel | undefined
+  /** Whether this step left its kernel for the next one. */
+  let parked = false
   let turn: TurnState | undefined
 
   const run = async (
     controller: ReadableStreamDefaultController<UIMessageChunk>,
     turns: TurnPlan[],
+    adopted: Parked | undefined,
   ): Promise<void> => {
     controller.enqueue({ type: 'start', messageId: step.responseMessageId })
-    for (const plan of turns) {
-      if (stopped) return
-      const fxTurn = agent!.prompt(plan.prompt, { signal: step.signal })
+    if (adopted !== undefined) {
+      await resume(controller, turns.at(-1)!, adopted)
+      return
+    }
+    let rebuilds = 0
+    let next = 0
+    while (next < turns.length) {
+      if (stopped || step.signal.aborted) return
+      const plan = turns[next]!
+      next += 1
+      // The step's abort cancels the turn; the turn outlives it when parked.
+      const fxTurn = kernel!.agent.prompt(plan.prompt)
       const state: TurnState = {
         plan,
         fxTurn,
@@ -159,39 +294,137 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
           throw diverged(
             `a recorded turn ended with ${result.stopReason}, not ${expected}`,
           )
+        if (kernel!.failed) {
+          // libfx's cancel aborts the request FX waits on as it sends the
+          // stop, and FX fails the turn when it reads the abort first. That
+          // kernel holds no interruption where the log has one, so the turns
+          // replay on a new kernel.
+          if (rebuilds === REPLAY_REBUILDS)
+            throw diverged('FX failed a stopped turn instead of cancelling it')
+          rebuilds += 1
+          await kernel!.agent.close().catch(() => {})
+          kernel = await create(read.saved)
+          next = 0
+        }
         continue
       }
-      const outcome = await Promise.race([
-        fxTurn.result.then((result) => ({ ended: true, result })),
-        state.toolsRequested.promise.then(() => ({ ended: false as const })),
-      ]).catch((error: unknown) => {
-        throw state.failure ?? error
-      })
-      if (state.failure) throw state.failure
-      await state.parsed
-      if (outcome.ended) {
-        await drained
-        // Only a turn this attempt finished describes the log.
-        if (outcome.result.stopReason === 'end_turn' && !step.signal.aborted)
-          await saveCheckpoint()
-      } else {
-        // FX now waits on tools that A2's tool consumer run.
-        fxTurn.cancel()
+      await finishLive(controller, state, drained)
+    }
+  }
+
+  /** Continues a parked kernel: the results it waits on are in the log. */
+  const resume = async (
+    controller: ReadableStreamDefaultController<UIMessageChunk>,
+    plan: TurnPlan,
+    adopted: Parked,
+  ): Promise<void> => {
+    const state: TurnState = {
+      plan,
+      fxTurn: adopted.fxTurn,
+      next: plan.calls.length,
+      served: plan.calls.at(-1)!.toolCalls,
+      answered: new Set(),
+      live: false,
+      liveStep: false,
+      toolsRequested: deferred(),
+    }
+    turn = state
+    // Steering goes in before any tool result, so FX's next model call can't
+    // start without it, as replay queues it before the recorded call ends.
+    for (const text of plan.steers.get(plan.calls.length - 1) ?? [])
+      await adopted.fxTurn.steer(text)
+    for (const waiting of kernel!.waiting.splice(0)) {
+      try {
+        const answered = answer(state, waiting.toolName, waiting.input)
+        if (answered === undefined)
+          fail(state, diverged(`no recorded result for ${waiting.toolName}`))
+        waiting.resolve(answered.output)
+      } catch (error) {
+        waiting.reject(error)
       }
-      controller.enqueue({ type: 'finish-step' })
-      controller.enqueue({
-        type: 'finish',
-        ...(completion.finishReason === undefined
-          ? {}
-          : { finishReason: completion.finishReason }),
+    }
+    await finishLive(controller, state, adopted.drained)
+  }
+
+  const finishLive = async (
+    controller: ReadableStreamDefaultController<UIMessageChunk>,
+    state: TurnState,
+    drained: Promise<void>,
+  ): Promise<void> => {
+    const outcome = await Promise.race([
+      state.fxTurn.result.then((result) => ({ ended: true, result })),
+      state.toolsRequested.promise.then(() => ({ ended: false as const })),
+    ]).catch((error: unknown) => {
+      throw state.failure ?? error
+    })
+    if (state.failure) throw state.failure
+    await state.parsed
+    if (outcome.ended) {
+      if (
+        outcome.result.stopReason === 'cancelled' &&
+        !stopped &&
+        !step.signal.aborted
+      )
+        throw diverged('the turn was cancelled while this step ran it')
+      await drained
+      // Only a turn this attempt finished describes the log.
+      if (outcome.result.stopReason === 'end_turn' && !step.signal.aborted)
+        await saveCheckpoint()
+    } else if (kernel!.waiting.length > 0 && !stopped && !closed) {
+      // FX now waits on tools that A2's tool consumer runs; the next step in
+      // this process can answer it with their results.
+      lot.park(step.sessionId, {
+        kernel: kernel!,
+        fxTurn: state.fxTurn,
+        drained,
+        generationId: step.generationId,
+        settings,
       })
+      parked = true
+    } else {
+      state.fxTurn.cancel()
     }
+    controller.enqueue({ type: 'finish-step' })
+    controller.enqueue({
+      type: 'finish',
+      ...(completion.finishReason === undefined
+        ? {}
+        : { finishReason: completion.finishReason }),
+    })
+  }
+
+  /**
+   * Whether a parked kernel stands where this step's log does: made the same
+   * way, its model call the last one the log's live turn committed, and a
+   * recorded result for each tool call it waits on. A steered response
+   * continues the same turn, so its steps qualify too.
+   */
+  const adoptable = (entry: Parked, turns: TurnPlan[]): boolean => {
+    const plan = turns.at(-1)
+    if (
+      entry.settings !== settings ||
+      plan?.end !== 'live' ||
+      plan.generationId !== entry.generationId
+    )
+      return false
+    const unanswered = [...(plan.calls.at(-1)?.toolCalls ?? [])]
+    return entry.kernel.waiting.every((waiting) => {
+      const at = unanswered.findIndex(
+        (call) =>
+          call.toolName === waiting.toolName &&
+          sameJson(call.input, waiting.input) &&
+          results.has(call.toolCallId),
+      )
+      if (at === -1) return false
+      unanswered.splice(at, 1)
+      return true
+    })
   }
 
   const saveCheckpoint = async (): Promise<void> => {
     let checkpoint: Uint8Array
     try {
-      checkpoint = await agent!.checkpoint()
+      checkpoint = await kernel!.agent.checkpoint()
     } catch {
       // Later steps replay from the previous snapshot.
       return
@@ -218,7 +451,8 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
     const state = turn!
     if (state.failure) throw state.failure
     const body = decodeBody(init.body)
-    const compaction = isCompaction(body)
+    const request = JSON.parse(body) as GatewayRequest
+    const compaction = request.toolChoice?.type === 'none'
     const recorded = state.plan.calls[state.next]
     if (recorded !== undefined) {
       if ((recorded.kind === 'compaction') !== compaction)
@@ -245,7 +479,12 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
       return pending(init.signal)
     }
     state.live = true
-    return live(state, url, { ...init, body }, compaction)
+    return live(
+      state,
+      url,
+      { ...init, body: withGatewayCaching(body, request) },
+      compaction,
+    )
   }
 
   const affinity = gatewaySessionAffinity({
@@ -256,18 +495,13 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
   const live = async (
     state: TurnState,
     url: string,
-    init: RequestInit & { body: string },
+    init: RequestInit,
     compaction: boolean,
   ): Promise<Response> => {
     const headers = withAuth(init.headers, credential)
     // The kernel's own affinity is new with every rebuild; A2's lasts the session.
     headers.set('x-session-affinity', affinity)
-    const response = await fetch(url, {
-      ...init,
-      body: withGatewayCaching(init.body),
-      headers,
-      signal: network,
-    })
+    const response = await fetch(url, { ...init, headers, signal: network })
     // Abandoned before its headers came: the kernel has moved on without it.
     if (init.signal?.aborted) {
       void response.body?.cancel()
@@ -316,9 +550,12 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
     })
   }
 
-  const hostTools = await kernelTools(step.tools, (toolName, toolInput) => {
-    const state = turn!
-    if (state.failure) throw state.failure
+  /** The recorded result of a call from the step the kernel was served. */
+  const answer = (
+    state: TurnState,
+    toolName: string,
+    toolInput: unknown,
+  ): { output: unknown } | undefined => {
     const call = state.served.find(
       (candidate) =>
         candidate.toolName === toolName &&
@@ -326,24 +563,33 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
         sameJson(candidate.input, toolInput),
     )
     const result = call === undefined ? undefined : results.get(call.toolCallId)
-    if (call !== undefined && result !== undefined) {
-      state.answered.add(call.toolCallId)
-      if (result.error !== undefined) throw new Error(result.error)
-      if (result.denied === true) throw new Error('The tool call was denied.')
-      return result.output ?? null
-    }
+    if (call === undefined || result === undefined) return undefined
+    state.answered.add(call.toolCallId)
+    if (result.error !== undefined) throw new Error(result.error)
+    if (result.denied === true) throw new Error('The tool call was denied.')
+    return { output: result.output ?? null }
+  }
+
+  const hostTool = (toolName: string, toolInput: unknown): unknown => {
+    const state = turn!
+    if (state.failure) throw state.failure
+    const answered = answer(state, toolName, toolInput)
+    if (answered !== undefined) return answered.output
     if (state.live) {
       state.toolsRequested.resolve()
-      return pending(undefined)
+      return new Promise((resolve, reject) => {
+        kernel!.waiting.push({ toolName, input: toolInput, resolve, reject })
+      })
     }
     if (state.plan.end === 'interrupted') {
       state.fxTurn.cancel()
       return pending(undefined)
     }
     return fail(state, diverged(`no recorded result for ${toolName}`))
-  })
+  }
 
-  const plans = (turns: HarnessTurn[]): Promise<TurnPlan[]> =>
+  const hooks: KernelHooks = { fetch: hostFetch, tool: hostTool }
+  const planTurns = (turns: HarnessTurn[]): Promise<TurnPlan[]> =>
     Promise.all(turns.map((recorded) => turnPlan(recorded, step.signal)))
   let results = read.results
   let controller: ReadableStreamDefaultController<UIMessageChunk> | undefined
@@ -352,17 +598,45 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
     if (!closed) controller!.enqueue(chunk)
   }
 
-  const create = (checkpoint: Uint8Array | undefined) =>
-    createFxAgent({
+  const create = async (
+    checkpoint: Uint8Array | undefined,
+  ): Promise<Kernel> => {
+    const created = { hooks, waiting: [] as Waiting[], failed: false }
+    let prompt: unknown
+    const agent = await createFxAgent({
       apiKey: credential.token,
       model: kernelModel(step),
       ...(step.instructions.length === 0
         ? {}
         : { instructions: step.instructions }),
-      tools: hostTools,
+      tools: definitions.map((definition): FxHostTool => ({
+        ...definition,
+        execute: (toolInput: unknown) =>
+          created.hooks.tool(definition.name, toolInput),
+      })),
       ...(checkpoint === undefined ? {} : { checkpoint }),
-      fetch: hostFetch,
+      fetch: (input, init) => created.hooks.fetch(input, init),
+      // libfx reports a turn FX failed after its cancel as cancelled; only
+      // the ACP response to the turn's prompt tells the two apart.
+      onEvent: ({ type, message }) => {
+        const rpc = (message ?? {}) as {
+          id?: unknown
+          method?: unknown
+          error?: unknown
+        }
+        if (type === 'acp.send' && rpc.method === 'session/prompt') {
+          prompt = rpc.id
+          created.failed = false
+        } else if (
+          type === 'acp.receive' &&
+          rpc.method === undefined &&
+          rpc.id === prompt
+        )
+          created.failed = rpc.error !== undefined
+      },
     })
+    return Object.assign(created, { agent })
+  }
 
   const stream = new ReadableStream<UIMessageChunk>({
     start(streamController) {
@@ -374,16 +648,33 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
       step.signal.addEventListener('abort', abort, { once: true })
       void (async () => {
         try {
-          try {
-            agent = await create(read.saved)
-          } catch (error) {
-            if (read.saved === undefined) throw error
-            // An unusable checkpoint: replay the whole log instead.
-            read = await replayAll()
-            results = read.results
-            agent = await create(undefined)
+          let turns = await planTurns(read.turns)
+          // An aborted step leaves a parked kernel for a later attempt.
+          const candidate = step.signal.aborted
+            ? undefined
+            : lot.take(step.sessionId)
+          const adopted =
+            candidate !== undefined && adoptable(candidate, turns)
+              ? candidate
+              : undefined
+          if (candidate !== undefined && adopted === undefined)
+            closeParked(candidate)
+          if (adopted !== undefined) {
+            kernel = adopted.kernel
+            kernel.hooks = hooks
+          } else {
+            try {
+              kernel = await create(read.saved)
+            } catch (error) {
+              if (read.saved === undefined) throw error
+              // An unusable checkpoint: replay the whole log instead.
+              read = await replayAll()
+              results = read.results
+              kernel = await create(undefined)
+              turns = await planTurns(read.turns)
+            }
           }
-          await run(streamController, await plans(read.turns))
+          await run(streamController, turns, adopted)
           closed = true
           streamController.close()
         } catch (error) {
@@ -394,7 +685,7 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
         } finally {
           released.abort()
           step.signal.removeEventListener('abort', abort)
-          await agent?.close().catch(() => {})
+          if (!parked) await kernel?.agent.close().catch(() => {})
         }
       })()
     },
@@ -429,7 +720,9 @@ async function turnPlan(
     input: turn.prompt,
     signal,
   })
+  const generationId = turn.generations.at(-1)?.completed.generationId
   return {
+    ...(generationId === undefined ? {} : { generationId }),
     prompt:
       images.length === 0
         ? turn.prompt.text
@@ -786,10 +1079,7 @@ async function loadLibfx(): Promise<typeof import('libfx')> {
  * asks for it; options the request already has take precedence. Every other
  * byte stays as the kernel wrote it, a tool input's spacing included.
  */
-function withGatewayCaching(body: string): string {
-  const request = JSON.parse(body) as {
-    providerOptions?: Record<string, Record<string, unknown> | undefined>
-  }
+function withGatewayCaching(body: string, request: GatewayRequest): string {
   const options = request.providerOptions
   if (options === undefined)
     return body.replace(
@@ -812,6 +1102,9 @@ function withAuth(
 ): Headers {
   const result = new Headers(headers)
   result.set('ai-gateway-auth-method', credential.method)
+  // A parked kernel keeps the credential it was made with; the step's is current.
+  if (result.has('authorization'))
+    result.set('authorization', `Bearer ${credential.token}`)
   return result
 }
 
@@ -831,17 +1124,16 @@ function kernelModel(
   }
 }
 
-async function kernelTools(
+/** The agent's tools as the kernel takes them, without their execution. */
+async function toolDefinitions(
   tools: ToolSet,
-  execute: (toolName: string, input: unknown) => unknown,
-): Promise<FxHostTool[]> {
+): Promise<{ name: string; description: string; inputSchema: object }[]> {
   return Promise.all(
     Object.entries(tools).map(async ([toolName, tool]) => ({
       name: toolName,
       description:
         typeof tool.description === 'string' ? tool.description : toolName,
       inputSchema: (await asSchema(tool.inputSchema).jsonSchema) as object,
-      execute: (input: unknown) => execute(toolName, input),
     })),
   )
 }
@@ -889,11 +1181,6 @@ function savedCheckpoint(state: unknown): Uint8Array | undefined {
   return new Uint8Array(Buffer.from(checkpoint, 'base64'))
 }
 
-function isCompaction(body: string): boolean {
-  const request = JSON.parse(body) as { toolChoice?: { type?: string } }
-  return request.toolChoice?.type === 'none'
-}
-
 function decodeBody(body: BodyInit | null | undefined): string {
   if (typeof body === 'string') return body
   if (body instanceof Uint8Array || body instanceof ArrayBuffer)

package/src/ai-harness.ts

@@ -69,6 +69,8 @@ export type HarnessLog = {
 
 /** Everything one durable generation hands its harness. */
 export type HarnessStep = {
+  /** The agent's own identity: what a harness keeps in memory is per scope. */
+  scope: object
   agentName: string
   sessionId: string
   requestIndex: number
@@ -95,6 +97,11 @@ export type HarnessDriver = {
   /** Tools the harness brings; A2 executes them like the agent's own. */
   readonly tools: ToolSet
   step(step: HarnessStep): Promise<HarnessGeneration>
+  /**
+   * The session has nothing left to run in this process. Lets go of anything
+   * kept in memory for its next step; it holds nothing the log needs.
+   */
+  release?(session: { scope: object; sessionId: string }): void
 }
 
 // A global symbol, so harness values work across separately bundled copies.