perf(headless-runner): reuse a turn's tools and memory across its steps

ee/apps/headless-runner/OWNER-CELLS.md

@@ -49,7 +49,7 @@ Workbot + owner ─ router                      legacy chats not yet split, forw
   in the same shape as before; like legacy chats, only chats with a message are listed.
 - **Shared memory.** `memoryScope: "owner"` (strict enum; needs an owner) puts the chat's `memory/` in the owner's
   memory (`owner_memory`), shared by all their chats. A session cell reads it from the owner cell
-  (`/internal/memory`) once per turn (once per step in step mode), when it builds the system prompt, and writes through
+  (`/internal/memory`) once per run of a turn (again when a restart continues it), when it builds the system prompt, and writes through
   on each file tool call; the files API shows it too. If the owner cell doesn't answer (3 tries), the turn fails
   `memory_unavailable` (resend it) or, in step mode, the step runs again in 15 s; a file tool returns an error. The first
   time an owner's memory is used, the memory their legacy chats share becomes it (the chat most others name in

ee/apps/headless-runner/README.md

@@ -69,7 +69,7 @@ return credentials or URLs. Public `/health` remains a liveness check. See
 Turn status is one of `queued`, `running`, `completed`, `failed`, `interrupted` or `aborted`. `failed` and `interrupted` can be resumed by re-sending the same `messageId` with fresh credentials; `aborted` cannot. Resuming continues the same turn (Workbot's "Try again" does this), so the conversation never gets a second copy of the message. A failed turn's `errorDetail` says what the model provider answered, with its request id, and the turn's `turn ended` log line carries it too; it is for support, never for the person. The `error` field holds a stable code:
 
 - `model_credentials_missing`
-- `mcp_unavailable` (connecting to `HEADLESS_MCP_URL` failed three times in a row, 0.5 s and 2 s apart)
+- `mcp_unavailable` (connecting to `HEADLESS_MCP_URL` failed three times in a row, 0.5 s and 2 s apart, before the turn's tools were listed; see [Step mode](#step-mode) for later steps)
 - `model_http_<status>`
 - `turn_timeout`
 - `max_steps_exceeded` (only with `HEADLESS_MAX_STEPS` set: the turn made that many model calls)
@@ -219,6 +219,12 @@ interrupt, the repeat check and `credentials_refresh` work the same way.
 - **A step** is: build the context, one streaming model call, that call's tool calls, then one SQLite transaction
   that writes the step's messages and tool results, the turn's next step (`step_schedule`) and the attempt reset. A
   crash leaves the whole step or none of it.
+- **Tools and memory are read once per run.** The first step lists the OpenWork MCP tools and reads `memory/`;
+  later steps show the model the same list and memory, as a turn run in one go does, so every request starts the
+  same and no later step waits on either before its model call. A later step connects only when the model calls
+  an MCP tool, starting as soon as the model begins writing that call. If that connection fails (three tries),
+  the step's MCP calls get an error saying they were not made and the next step connects again; the turn does not
+  end `mcp_unavailable`. A restart or a fresh MCP token lists the tools again.
 - **Between steps nothing runs.** On celld, steps run from the cell's durable alarm (the same alarm as self-resume,
   set for whichever is due first), each to completion, lanes in parallel. Between two alarms the cell is idle, so a
   runner deploy switches it to the new code there instead of cancelling a step after `CELLD_DEPLOY_MAX_AGE_S`. On

ee/apps/headless-runner/scripts/steps.test.ts

@@ -10,11 +10,12 @@ import { denResumeClient, RESUME_EXCHANGE_PATH } from "../src/den-resume.js"
 import type { ErrorReport, ErrorReporter } from "../src/error-reporting.js"
 import { readPart } from "../src/files.js"
 import type { McpConnector } from "../src/mcp.js"
+import type { MemoryResolver } from "../src/memory.js"
 import type { ModelClient, ModelRequest } from "../src/model.js"
 import { nodeSqlite } from "../src/node-sqlite.js"
 import type { SqlDriver } from "../src/sql.js"
 import { grantVault } from "../src/resume-grants.js"
-import { CUT_OFF_INSTRUCTION, Runner, type RunnerOptions } from "../src/runner.js"
+import { APPS_UNREACHABLE, CUT_OFF_INSTRUCTION, Runner, type RunnerOptions } from "../src/runner.js"
 import { SelfResume } from "../src/self-resume.js"
 import { FRESH_REPEAT, type StepState } from "../src/steps.js"
 import { Store } from "../src/store.js"
@@ -54,6 +55,8 @@ function recordingErrors() {
 type WorldInput = {
   model: ModelClient
   mcp?: McpConnector
+  memory?: MemoryResolver
+  sleep?: RunnerOptions["sleep"]
   den?: string
   alarm?: boolean
   clock?: () => number
@@ -80,6 +83,8 @@ function world(directory: string, input: WorldInput) {
     defaultModelApiKey: "test-key",
     limits: { ...unlimited, ...input.limits },
     mcp: input.mcp,
+    memory: input.memory,
+    sleep: input.sleep,
     resume,
     errors: input.errors,
     now: input.clock,
@@ -678,3 +683,90 @@ test("a reply cut off at the output cap never ends the turn: it is dropped, the
     assert.deepEqual([turn?.status, turn?.error], ["failed", "output_limit"], `${execution ?? "one go"}`)
   }
 })
+
+/** Notes something, looks something up in the person's apps (twice when the first try fails), then answers. */
+const lookupModel = (requests: ModelRequest[]): ModelClient => ({
+  async complete(request) {
+    requests.push(request)
+    const results = resultsOf(request)
+    if (results.length === 0) return { text: "", toolCalls: [{ id: "call_note", name: "write_file", input: { path: "notes.md", content: "ask about tea" } }], usage }
+    const last = results.at(-1)
+    if (last?.name === "lookup" && !last.isError) return { text: `They like ${last.output}.`, toolCalls: [], usage }
+    request.onTool?.("lookup")
+    return { text: "", toolCalls: [{ id: `call_lookup_${results.length}`, name: "lookup", input: {} }], usage }
+  },
+})
+
+test("a step-mode run lists the person's apps' tools and reads memory once: later steps connect only for a call", async () => {
+  using temp = tempDir()
+  const requests: ModelRequest[] = []
+  let connects = 0
+  let closes = 0
+  let lists = 0
+  const mcp: McpConnector = async () => {
+    connects += 1
+    return {
+      tools: [{ name: "lookup", description: "Looks something up", inputSchema: { type: "object" } }],
+      call: async () => ({ output: "tea", isError: false }),
+      async close() {
+        closes += 1
+      },
+    }
+  }
+  const memory: MemoryResolver = () => ({
+    async list() {
+      lists += 1
+      return [{ path: "memory/me.md", content: `read ${lists}` }]
+    },
+    read: async () => null,
+    write: async () => ({ ok: true, size: 0 }),
+    delete: async () => false,
+  })
+  const { store, runner } = world(temp.directory, { model: lookupModel(requests), mcp, memory })
+  const session = store.createSession({})
+  runner.send({ sessionId: session.id, messageId: "m1", prompt: "What do they like?", credentials: { mcpToken: "token" }, execution: "steps" })
+  await runner.idle()
+  assert.equal(store.getTurn(session.id, "m1")?.status, "completed")
+  assert.equal(requests.length, 3, "a note, a lookup, the answer")
+  // The first step lists the tools; of the two later steps, only the one that calls an app tool connects.
+  assert.equal(connects, 2)
+  assert.equal(closes, 2, "nothing stays open between steps")
+  assert.equal(lists, 1)
+  for (const request of requests) {
+    assert.deepEqual(request.tools.map((tool) => tool.name), requests[0]?.tools.map((tool) => tool.name))
+    assert.equal(request.system, requests[0]?.system, "every step's request starts the same, so the prompt cache keeps hitting")
+  }
+  assert.ok(requests[0]?.tools.some((tool) => tool.name === "lookup"))
+  assert.ok(requests[0]?.system.includes("read 1"))
+  assert.deepEqual(transcript(store, session.id, "m1").filter((message) => message.role === "tool").map((message) => message.output).slice(1), ["tea"])
+})
+
+test("a later step that can't connect to the person's apps: its call was never sent and says so, and the next step connects again", async () => {
+  using temp = tempDir()
+  const requests: ModelRequest[] = []
+  let connects = 0
+  let calls = 0
+  // The first step lists the tools; then the apps are out of reach for one step's three tries.
+  const mcp: McpConnector = async () => {
+    connects += 1
+    if (connects >= 2 && connects <= 4) throw new Error("connect failed")
+    return {
+      tools: [{ name: "lookup", description: "Looks something up", inputSchema: { type: "object" } }],
+      async call() {
+        calls += 1
+        return { output: "tea", isError: false }
+      },
+      async close() {},
+    }
+  }
+  const { store, runner } = world(temp.directory, { model: lookupModel(requests), mcp, sleep: async () => undefined })
+  const session = store.createSession({})
+  runner.send({ sessionId: session.id, messageId: "m1", prompt: "What do they like?", credentials: { mcpToken: "token" }, execution: "steps" })
+  await runner.idle()
+  const turn = store.getTurn(session.id, "m1")
+  assert.deepEqual([turn?.status, turn?.error], ["completed", null])
+  const results = transcript(store, session.id, "m1").filter((message): message is Extract<Message, { role: "tool" }> => message.role === "tool")
+  assert.deepEqual(results.slice(1).map((result) => [result.output, result.isError]), [[APPS_UNREACHABLE, true], ["tea", false]])
+  assert.equal(connects, 5, "listed once, three failed tries in one step, connected again in the next")
+  assert.equal(calls, 1, "the call that couldn't be sent ran once, when it was tried again")
+})

ee/apps/headless-runner/src/memory.ts

@@ -3,15 +3,16 @@ import { FILE_LIMITS } from "./files.js"
 import type { Session, Store } from "./store.js"
 
 /**
- * Where a conversation's memory/ lives, for the model's context (all files, read once per turn or step) and for the
+ * Where a conversation's memory/ lives, for the model's context (all files, read once per run of a turn) and for the
  * file tools (one file at a time). Everything else in the scratch workspace stays the conversation's own.
  *
  * - No access (null): memory/ is the conversation's own files, as always.
  * - `memoryOf` (legacy): another conversation's files in the same database.
  * - `memoryScope: "owner"`: the owner's memory. Local where it is stored (Node, an owner cell); in a session cell, calls
- *   to the owner cell (see src/worker), one read per context build and each write through.
+ *   to the owner cell (see src/worker), one read per run and each write through.
  *
- * `list` is called once per turn (once per step in step mode) and its files are shown in the system prompt.
+ * `list` is called once per run of a turn (again after a restart) and its files are shown in the system prompt of
+ * every step, so a turn's requests start the same and the provider's prompt cache keeps hitting.
  */
 export type MemoryAccess = {
   list(): Promise<MemoryFile[]>

ee/apps/headless-runner/src/runner.ts

@@ -101,6 +101,8 @@ function stepSystem(base: string, askedToStop: boolean, cutOffs: number) {
 
 /** A failed connection to the person's apps is tried again after each of these waits; then the turn fails. */
 export const MCP_CONNECT_RETRY_MS = [500, 2_000]
+/** The result of a call that was never sent: a step-mode run's later step couldn't connect (see laterStepTools). */
+export const APPS_UNREACHABLE = "The person's apps couldn't be reached, so this call was not made. Try it again."
 
 const abortableSleep = (ms: number, signal: AbortSignal) =>
   new Promise<void>((resolve, reject) => {
@@ -259,6 +261,13 @@ type StepLane = {
   finished: boolean
   /** The current turn's files, read once and shown to the model on every step. */
   expanded: Map<string, Message>
+  /**
+   * What the run's first step read over the network, which its later steps use as it is, as a turn run in one go
+   * does: the person's apps' tools (for the token they were listed with) and the memory files. Every step's request
+   * so starts the same, and a later step's model call waits for neither (see laterStepTools).
+   */
+  tools?: { token: string; specs: ToolSpec[] }
+  memoryFiles?: MemoryFiles
   startedAt: number
   steps: number
   toolCalls: number
@@ -268,6 +277,7 @@ type StepLane = {
   /** What a step that failed unexpectedly threw, for the error report. */
   thrown?: unknown
 }
+type MemoryFiles = Array<{ path: string; content: string }>
 /** What a turn may use and is told, the same for every step of one run. */
 type TurnSetup = {
   turn: Turn | null
@@ -280,6 +290,8 @@ type TurnSetup = {
   taskTools: boolean
   /** Where memory/ lives when it isn't this conversation's own (memory.ts). */
   memory: MemoryAccess | null
+  /** The memory files in `baseSystem`, for a step-mode run's later steps (StepLane.memoryFiles). */
+  memoryFiles: MemoryFiles
   baseSystem: string
   toolSpecs: ToolSpec[]
   /** Tools whose calls only read: the local read tools, and app tools whose server says so (ToolSpec.readOnly). */
@@ -609,6 +621,47 @@ export class Runner {
     }
   }
 
+  /**
+   * A step-mode run's tools after its first step: the list that step got, so the model call doesn't wait for the
+   * person's apps, and a connection only if the model calls one of them (`prepare`: as soon as it starts writing such
+   * a call). When that connection can't be made, the step's calls to them were never sent and say so; the next step
+   * connects again.
+   */
+  private laterStepTools(mcp: McpConnector, token: string, specs: ToolSpec[], signal: AbortSignal): ToolSession & { prepare(name: string): void } {
+    const names = new Set(specs.map((tool) => tool.name))
+    // Ends a connection attempt still under way when the step is over.
+    const closing = new AbortController()
+    let connecting: Promise<ToolSession> | null = null
+    const connect = () => {
+      if (!connecting) {
+        connecting = this.connectTools(mcp, token, AbortSignal.any([signal, closing.signal]))
+        // Each call awaits it; a failure before any call does must not go unhandled.
+        connecting.catch(() => undefined)
+      }
+      return connecting
+    }
+    return {
+      tools: specs,
+      prepare: (name) => {
+        if (names.has(name)) void connect()
+      },
+      call: async (name, input, callSignal) => {
+        let session: ToolSession
+        try {
+          session = await connect()
+        } catch (error) {
+          if (signal.aborted) throw error
+          return { output: APPS_UNREACHABLE, isError: true }
+        }
+        return session.call(name, input, callSignal)
+      },
+      close: async () => {
+        closing.abort(new Error("step_ended"))
+        await connecting?.then((session) => session.close(), () => undefined)
+      },
+    }
+  }
+
   private scopedModel(access: HeadlessModelAccess) {
     return this.options.modelForAccess?.(access) ?? createModelAdapter(access, { maxOutputTokens: 16_384 })
   }
@@ -941,7 +994,7 @@ export class Runner {
    * What the turn may use and is told: its tools, files, computer and system prompt. The same for every step of a run,
    * so the provider's prompt cache keeps hitting.
    */
-  private async turnSetup(sessionId: string, messageId: string, kind: JobKind, tools: ToolSession | null, credentials: TurnCredentials): Promise<TurnSetup> {
+  private async turnSetup(sessionId: string, messageId: string, kind: JobKind, tools: ToolSession | null, credentials: TurnCredentials, knownMemory?: MemoryFiles): Promise<TurnSetup> {
     const { store } = this.options
     const session = store.getSession(sessionId)
     const turn = store.getTurn(sessionId, messageId)
@@ -957,9 +1010,9 @@ export class Runner {
     const sentFiles = store.turnMessages(sessionId, messageId).some(({ message }) => message.role === "user" && (message.attachments?.length ?? 0) > 0)
     if (computer && (sentFiles || computer.known(sessionId))) computer.prewarm(sessionId)
     // A person's conversations share one memory: the owner's (memoryScope), or another conversation's (memoryOf, while
-    // it exists). Read once here, for every step of this run (once per step in step mode).
+    // it exists). Read once, for every step of this run: a step-mode run's later steps pass what its first step read.
     const memory = session ? this.memory(session) : null
-    const memoryFiles = memory ? await memory.list() : store.memoryFiles(sessionId)
+    const memoryFiles = knownMemory ?? (memory ? await memory.list() : store.memoryFiles(sessionId))
     // Everything here stays the same from one turn to the next, so the provider's prompt cache keeps hitting.
     const baseSystem = [
       this.options.systemPrompt ?? DEFAULT_SYSTEM_PROMPT,
@@ -983,7 +1036,7 @@ export class Runner {
       if (tool.readOnly && !named.has(tool.name)) readOnlyTools.add(tool.name)
       named.add(tool.name)
     }
-    return { turn, session, repeatLimits, files, readOnly, computer, reactions, taskTools, memory, baseSystem, toolSpecs, readOnlyTools }
+    return { turn, session, repeatLimits, files, readOnly, computer, reactions, taskTools, memory, memoryFiles, baseSystem, toolSpecs, readOnlyTools }
   }
 
   /** The current turn's files as the model sees them, read once (into `expanded`) and shown on every step. */
@@ -1832,16 +1885,24 @@ export class Runner {
         // Again after a refresh: Den may have issued fresh scoped access with the token.
         const chosen = this.turnModel(sessionId, messageId, credentials)
         if (!chosen) return end("failed", "model_credentials_missing")
+        // Between steps nothing is held open. The run's first step (and the first with a fresh token) connects to
+        // list the tools; later steps show the model that list and connect only for a call.
+        let later: ReturnType<Runner["laterStepTools"]> | null = null
         if (needsTools && this.options.mcp && credentials.mcpToken) {
-          try {
-            // Every step connects again (between steps nothing is held open), so one failed connect is retried.
-            tools = await this.connectTools(this.options.mcp, credentials.mcpToken, signal)
-          } catch (error) {
-            if (signal.aborted) throw error
-            return end("failed", "mcp_unavailable")
+          if (lane.tools?.token === credentials.mcpToken) {
+            tools = later = this.laterStepTools(this.options.mcp, credentials.mcpToken, lane.tools.specs, signal)
+          } else {
+            try {
+              tools = await this.connectTools(this.options.mcp, credentials.mcpToken, signal)
+            } catch (error) {
+              if (signal.aborted) throw error
+              return end("failed", "mcp_unavailable")
+            }
+            lane.tools = { token: credentials.mcpToken, specs: tools.tools }
           }
         }
-        const setup = await this.turnSetup(sessionId, messageId, kind, tools, credentials)
+        const setup = await this.turnSetup(sessionId, messageId, kind, tools, credentials, lane.memoryFiles)
+        lane.memoryFiles = setup.memoryFiles
         reactions = setup.reactions
         const modelStep = turnMessages().filter((message) => message.role === "assistant").length
         const { messages, cacheAfter } = await this.stepContext(sessionId, messageId, kind, setup, lane.expanded)
@@ -1894,6 +1955,8 @@ export class Runner {
             },
             onTool: (tool: string, description?: string, view?: DesktopView) => {
               lane.timing.output()
+              // Connecting while the model still writes the call, so the call itself doesn't wait for it.
+              later?.prepare(tool)
               events?.emit(sessionId, { type: "tool", messageId, step: modelStep, tool, ...(description ? { description } : {}), ...(view ? { view } : {}) })
             },
           })