From bfae73aa258d4044a051f53d7a07564629ed1e34 Mon Sep 17 00:00:00 2001 From: Jalil Date: Fri, 13 Nov 2026 17:00:07 -0800 Subject: [PATCH] perf(headless-runner): reuse a turn's tools and memory across its steps A turn that runs step by step connected to OpenWork MCP, listed the tools and read memory/ again before every model call. A turn run in one go does each once. - The run's first step lists the tools and reads memory; later steps show the model the same list and memory, 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 the call, and closes when the step ends. Nothing stays open between steps. - When that connection can't be made, the step's MCP calls say they were not made and the next step connects again. The turn no longer ends mcp_unavailable there; the first step still does. - A restart or a fresh MCP token lists the tools again. Signed-off-by: Jalil --- ee/apps/headless-runner/OWNER-CELLS.md | 2 +- ee/apps/headless-runner/README.md | 8 +- ee/apps/headless-runner/scripts/steps.test.ts | 94 ++++++++++++++++++- ee/apps/headless-runner/src/memory.ts | 7 +- ee/apps/headless-runner/src/runner.ts | 85 ++++++++++++++--- 5 files changed, 179 insertions(+), 17 deletions(-) diff --git a/ee/apps/headless-runner/OWNER-CELLS.md b/ee/apps/headless-runner/OWNER-CELLS.md index 947dd878e62..643302b364a 100644 --- a/ee/apps/headless-runner/OWNER-CELLS.md +++ b/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 diff --git a/ee/apps/headless-runner/README.md b/ee/apps/headless-runner/README.md index 5de520fda37..f4108f1c5b8 100644 --- a/ee/apps/headless-runner/README.md +++ b/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_` - `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 diff --git a/ee/apps/headless-runner/scripts/steps.test.ts b/ee/apps/headless-runner/scripts/steps.test.ts index 48733800e12..a9b707b8df3 100644 --- a/ee/apps/headless-runner/scripts/steps.test.ts +++ b/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") + 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") +}) diff --git a/ee/apps/headless-runner/src/memory.ts b/ee/apps/headless-runner/src/memory.ts index 344a221c715..af0a491926f 100644 --- a/ee/apps/headless-runner/src/memory.ts +++ b/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 diff --git a/ee/apps/headless-runner/src/runner.ts b/ee/apps/headless-runner/src/runner.ts index 28db6c1653e..c3933a17c8b 100644 --- a/ee/apps/headless-runner/src/runner.ts +++ b/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((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 + /** + * 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 | 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 { + private async turnSetup(sessionId: string, messageId: string, kind: JobKind, tools: ToolSession | null, credentials: TurnCredentials, knownMemory?: MemoryFiles): Promise { 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 | 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 } : {}) }) }, })