fix(workbot): keep a computer command's result when its wait is cut

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

@@ -19,7 +19,7 @@ import { APPS_UNREACHABLE, CUT_OFF_INSTRUCTION, Runner, type RunnerOptions } fro
 import { SelfResume } from "../src/self-resume.js"
 import { FRESH_REPEAT, type StepState } from "../src/steps.js"
 import { Store } from "../src/store.js"
-import type { Message } from "../src/types.js"
+import type { Message, SessionComputer, ToolResult } from "../src/types.js"
 
 const usage = { inputTokens: 3, cachedInputTokens: 1, outputTokens: 2 }
 const unlimited: RunnerOptions["limits"] = { maxConcurrentTurns: 4, maxSteps: 20, turnTimeoutMs: Number.POSITIVE_INFINITY, credentialRefreshMs: Number.POSITIVE_INFINITY, contextCharBudget: 100_000 }
@@ -63,6 +63,7 @@ type WorldInput = {
   errors?: ErrorReporter
   limits?: Partial<RunnerOptions["limits"]>
   driver?: (db: SqlDriver) => SqlDriver
+  computer?: SessionComputer
 }
 
 /**
@@ -89,6 +90,7 @@ function world(directory: string, input: WorldInput) {
     errors: input.errors,
     now: input.clock,
     stepAlarm: input.alarm,
+    computer: input.computer,
   })
   const steps = runner.recoverSteps()
   return { store, runner, recovered, steps }
@@ -323,6 +325,59 @@ test("a tool call whose result wasn't recorded is never run again: the model get
   assert.deepEqual(tokens, ["caller-token", "fresh-token-1"])
 })
 
+test("a computer command cut off with its step is picked up from the computer, never run again; one it never started stays interrupted", async () => {
+  using temp = tempDir()
+  const requests: ModelRequest[] = []
+  const model: ModelClient = {
+    async complete(request) {
+      requests.push(request)
+      if (resultsOf(request).length === 0) {
+        return { text: "Building it.", toolCalls: [{ id: "call_build", name: "bash", input: { command: "make report" } }, { id: "call_send", name: "bash", input: { command: "send report" } }], usage }
+      }
+      return { text: "The report is built; I couldn't confirm it was sent.", toolCalls: [], usage }
+    },
+  }
+  // The conversation's computer, which outlives the runner: it keeps what it started, by the call's key.
+  const started: Array<{ command: unknown; key: string | undefined }> = []
+  const asked: string[] = []
+  const computer = (results: Record<string, ToolResult>): SessionComputer => ({
+    prompt: "",
+    tools: [{ name: "bash", description: "Runs a command", inputSchema: { type: "object" } }],
+    toolNames: new Set(["bash"]),
+    image: "test",
+    async run(_sessionId, _name, input, callKey) {
+      started.push({ command: input.command, key: callKey })
+      return never()
+    },
+    async pickUp(_sessionId, _name, input, callKey) {
+      asked.push(callKey)
+      return results[String(input.command)] ?? null
+    },
+    known: () => false,
+    prewarm: () => undefined,
+    release: () => undefined,
+    delete: async () => undefined,
+    previewSavedFile: async () => undefined,
+  })
+  const first = world(temp.directory, { model, computer: computer({}) })
+  const session = first.store.createSession({ computer: true })
+  first.runner.send({ sessionId: session.id, messageId: "m1", prompt: "Build and send the report", credentials: { modelApiKey: "caller-key" }, execution: "steps" })
+  await until(() => started.length === 1, "the build starting")
+  // The runner dies while the build runs. On the computer it finishes; the send never started.
+  const second = world(temp.directory, { model, computer: computer({ "make report": { output: "[exit 0]\nreport.pdf written", isError: false } }) })
+  await second.runner.idle()
+  assert.equal(second.store.getTurn(session.id, "m1")?.status, "completed")
+  assert.deepEqual(started.map((call) => call.command), ["make report"], "no command ran a second time")
+  assert.equal(asked.length, 2)
+  assert.equal(asked[0], started[0]?.key, "the computer is asked by the key the command was started under")
+  assert.notEqual(asked[1], asked[0])
+  const results = transcript(second.store, session.id, "m1").filter((message) => message.role === "tool")
+  assert.deepEqual(results.map((message) => [message.callId, message.isError]), [["call_build", false], ["call_send", true]])
+  assert.equal(results[0]?.output, "[exit 0]\nreport.pdf written")
+  assert.ok(results[1]?.output.includes("not retried"))
+  assert.equal(requests.length, 2, "the step finished from the kept answer, without asking the model again")
+})
+
 test("without a resume grant, a turn that lost its credentials is interrupted as on any restart, and resumes when re-sent", async () => {
   using temp = tempDir()
   let calls = 0

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

@@ -152,7 +152,9 @@ export function routeComputers(input: {
     desktopViewer: async (sessionId, view) => (await desktop(sessionId, false)).desktopViewer(sessionId, view),
     browserDemonstration: async (sessionId, command) => (await desktop(sessionId, true)).browserDemonstration(sessionId, command),
     desktopAction: async (sessionId, action) => (await desktop(sessionId, false)).desktopAction(sessionId, action),
-    run: async (sessionId, name, args) => (await routed(sessionId, true)).run(sessionId, name, args),
+    run: async (sessionId, name, args, callKey) => (await routed(sessionId, true)).run(sessionId, name, args, callKey),
+    // Asked where the conversation's computer is, never placing one: a call no computer started has no result.
+    pickUp: async (sessionId, name, args, callKey) => (await (await route(sessionId, false))?.pickUp?.(sessionId, name, args, callKey)) ?? null,
     known: (sessionId) => current(sessionId)?.known(sessionId) ?? false,
     prewarm: (sessionId) => void route(sessionId, true).then((computer) => computer?.prewarm(sessionId)).catch(() => undefined),
     // Both sides: a no-op where the conversation has no computer, and one switched away from still pauses.

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

@@ -540,6 +540,12 @@ function toolMessage(call: ToolCall, outcome: ToolResult): Message {
   }
 }
 
+/**
+ * Names one call of one step for good: the same when the step starts again, never shared with another step or run (a
+ * model may reuse its call ids). The computer keeps a command's result under it.
+ */
+const stepCallKey = (messageId: string, state: StepState, call: ToolCall) => `${messageId}/${state.stretchStartedAt}/${state.steps}/${call.id}`
+
 /** A step's kept answer as messages: each call with its recorded result, or `missing` as an error for the rest. */
 function pendingOutcomes(pending: PendingStep, missing: string): ToolResult[] {
   return pending.toolCalls.map((call) => pending.results.find((entry) => entry.callId === call.id)?.result ?? { output: missing, isError: true })
@@ -1109,7 +1115,7 @@ export class Runner {
   private async runCall(
     setup: TurnSetup,
     call: ToolCall,
-    context: { sessionId: string; kind: JobKind; tools: ToolSession | null; signal: AbortSignal; held: HeldCredentials },
+    context: { sessionId: string; kind: JobKind; tools: ToolSession | null; signal: AbortSignal; held: HeldCredentials; callKey?: string },
   ): Promise<ToolResult> {
     const { store } = this.options
     const { readOnly, files, reactions, taskTools, turn, computer, memory } = setup
@@ -1139,7 +1145,7 @@ export class Runner {
           ? desktopViewRefused(store.getSession(sessionId)?.desktopViews ?? null, call.input)
             ?? await computer.desktopRun(sessionId, call.input, signal, { browserRecovery: store.getSession(sessionId)?.browserRecovery === true }).catch((error: unknown) => ({ output: error instanceof Error ? error.message : "Desktop unavailable", isError: true }))
         : computer?.toolNames.has(call.name)
-          ? await computer.run(sessionId, call.name, call.input).catch((error: unknown) => ({
+          ? await computer.run(sessionId, call.name, call.input, context.callKey).catch((error: unknown) => ({
               output: `The computer didn't respond: ${error instanceof Error ? error.message : "unknown error"}. Its outcome is unknown; check before repeating it.`,
               isError: true,
             }))
@@ -1845,8 +1851,23 @@ export class Runner {
         // The step stopped while its tool calls ran: it finishes from the kept answer. A call without a recorded
         // result is not run again (it may have had side effects); the model gets an error and checks.
         answer = state.pending
+        const readOnly = (lane.held?.credentials.readOnly ?? state.readOnly) === true
+        // Except a command on the conversation's computer, which went on without the runner and kept its result
+        // there: the computer is asked for it. A command it never started stays an interrupted call.
+        const computer = session?.computer && !readOnly && kind !== "report" && this.options.computer?.available?.(sessionId) !== false ? this.options.computer : undefined
+        for (const call of answer.toolCalls) {
+          const pickUp = computer?.pickUp
+          if (!pickUp || call.inputError || !computer.toolNames.has(call.name) || answer.results.some((entry) => entry.callId === call.id)) continue
+          signal.throwIfAborted()
+          const key = stepCallKey(messageId, state, call)
+          const found = await lane.timing.tool(call.name, () => pickUp.call(computer, sessionId, call.name, call.input, key).catch(() => null))
+          if (!found) continue
+          answer.results.push({ callId: call.id, result: found })
+          state = { ...state, pending: answer }
+          store.setStepState(sessionId, messageId, state)
+        }
         outcomes = pendingOutcomes(answer, INTERRUPTED_CALL)
-        reactions = session?.reactions === true && kind === "message" && (lane.held?.credentials.readOnly ?? state.readOnly) !== true
+        reactions = session?.reactions === true && kind === "message" && !readOnly
       } else {
         const turnMessages = () => store.turnMessages(sessionId, messageId).map((entry) => entry.message)
         // History from a run in one go (or an import) may hold calls without results: never re-run.
@@ -1998,7 +2019,7 @@ export class Runner {
           store.setStepState(sessionId, messageId, state)
         }
         const runCall = (call: ToolCall) =>
-          lane.timing.tool(call.name, () => this.runCall(setup, call, { sessionId, kind, tools, signal, held: { credentials: lane.held?.credentials ?? {}, at: lane.issuedAt } }))
+          lane.timing.tool(call.name, () => this.runCall(setup, call, { sessionId, kind, tools, signal, held: { credentials: lane.held?.credentials ?? {}, at: lane.issuedAt }, callKey: stepCallKey(messageId, state, call) }))
         if (runsTogether(setup, answer.toolCalls)) {
           // Calls that can't change anything (lookups in the person's apps) run at the same time.
           signal.throwIfAborted()

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

@@ -14,7 +14,8 @@ import { toolAppSchema, toolCallSchema, type RepeatLimits, type ToolCall, type T
  *   minutes, then its turn fails with `step_crash_loop`.
  * - **No repeated writes.** Once the model has answered with tool calls, the answer is kept (`pending`) before any
  *   call runs, and each result as it returns. A step that dies there finishes from it: a call without a recorded
- *   result is never run again, it gets an error result so the model can check its effect.
+ *   result is never run again, it gets an error result so the model can check its effect. A command on the
+ *   conversation's computer went on without the runner and kept its result there, so the computer is asked for it.
  * - **Stream checkpoints.** A long model call keeps its text so far (at most every 2 s). A step that restarts after a
  *   cancel continues that text where the provider accepts a partial answer (Anthropic Messages without extended
  *   thinking) and otherwise starts over. Nothing but text is kept, so a half-written tool call is never run.

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

@@ -152,7 +152,13 @@ export type SessionComputer = {
   desktopViewer?(sessionId: string, input: { view: "browser" | "computer"; control: boolean; origin: string }): Promise<{ state: "paused" } | { state: "connected"; url: string; expiresAt: number; paused: boolean }>
   browserDemonstration?(sessionId: string, command: import("@openwork/types/browser-demonstration").DemonstrationCommand): Promise<import("@openwork/types/browser-demonstration").DemonstrationReply>
   desktopAction?(sessionId: string, action: "continue" | "pause" | "resume" | "rebuild"): Promise<{ ok: boolean }>
-  run(sessionId: string, name: string, input: Record<string, unknown>): Promise<ToolResult>
+  /** `callKey` names this one call for good, so a command started under it can be picked up later (`pickUp`). */
+  run(sessionId: string, name: string, input: Record<string, unknown>, callKey?: string): Promise<ToolResult>
+  /**
+   * The result of a command started under `callKey` whose result never reached the runner: its real result once it
+   * ends, or that it is still running. Null when the computer never started it. Never starts a command or a computer.
+   */
+  pickUp?(sessionId: string, name: string, input: Record<string, unknown>, callKey: string): Promise<ToolResult | null>
   /** Whether this conversation has a computer here at all (see computer-route.ts); unset means every conversation does. */
   available?(sessionId: string): boolean
   known(sessionId: string): boolean