feat(temporal): Kept a step's tools and seal on the worker that made its model call.

packages/temporal/README.md

@@ -60,7 +60,8 @@ session runs as the workflow `session-exec-<sessionID>`.
 | `OPENCODE_SESSION_IDLE_TIMEOUT` | | How long an idle session's workflow stays open. Local mode honors the same variable. |
 | `OPENCODE_TEMPORAL_STEPPED` | off | `1` runs each step as a model call, one activity per tool call, and a seal. |
 | `OPENCODE_TEMPORAL_WORKTREE_AFFINITY`, `OPENCODE_TEMPORAL_WORKTREE` | off | `1` routes a session to the workers serving its project tree, named by the second variable. |
-| `OPENCODE_TEMPORAL_SERIAL_TOOLS` | derived | Run a step's tool calls one at a time. On by default only where two hosts could write one step's tree: a shared store without affinity. |
+| `OPENCODE_TEMPORAL_STEP_AFFINITY` | on | Send a step's tool calls and seal back to the worker that made its model call. `0` turns it off. |
+| `OPENCODE_TEMPORAL_SERIAL_TOOLS` | derived | Run a step's tool calls one at a time. On by default only where two hosts could write one step's tree: a shared store with neither kind of affinity. |
 | `OPENCODE_EVENT_POLL_MS` | | How often a live subscriber re-reads the log for events another process appended. `0` turns the tick off. |
 | `OPENCODE_DB` | | One absolute path shared by every process on a host. |
 | `OPENCODE_DB_URL`, `OPENCODE_DB_AUTH_TOKEN` | | A libSQL URL for a store shared across hosts. Takes precedence over `OPENCODE_DB`. |
@@ -135,9 +136,12 @@ refusal apart from a tool that failed. The whole-step path still raises it as an
 the behaviour its own tests pin, so that drain says so at the boundary.
 
 Each tool call ships the tree from the host that ran it, because the seal can land on any worker and
-a capture there would miss what the tool wrote. Where two hosts could run one step's tools, they run
-one at a time (`OPENCODE_TEMPORAL_SERIAL_TOOLS`), or the second to ship would publish a tree without
-the first's work.
+a capture there would miss what the tool wrote. By default the tools and the seal of a step are sent
+back to the worker that made its model call, on a queue that worker polls on its own, so they write
+one tree through one filesystem and can run at once; a pin nobody answers falls back to the shared
+queue after a schedule-to-start bound. Where two hosts could still run one step's tools, they run one
+at a time (`OPENCODE_TEMPORAL_SERIAL_TOOLS`), or the second to ship would publish a tree without the
+first's work.
 
 What the split costs is the overlap between the model's stream and its own tools: a whole-step
 activity starts each tool the moment the model asks for it, while here the attempt returns first.

packages/temporal/src/config.ts

@@ -32,9 +32,15 @@ export interface Interface {
   /** Run a step's tool calls one at a time instead of together. Tools of one step write the same
    * tree and each ships from the host that ran it, so two on two hosts each publish a tree without
    * the other's work: the second is refused rather than reverting the first, which leaves its work
-   * stranded there. `OPENCODE_TEMPORAL_SERIAL_TOOLS=1` forces it on; it is on by default only where
-   * the store is shared and no affinity keeps a session's tools on one worker. */
+   * stranded there. `OPENCODE_TEMPORAL_SERIAL_TOOLS=1` forces it on; it is not needed while a step
+   * is pinned to one worker, which is the default. */
   readonly serialTools?: boolean
+  /** Send the tools and the seal of a step back to the worker that made its model call, on a queue
+   * that worker polls on its own. That worker is standing in the tree the tools are about to write,
+   * so the step's tools see each other's writes through the filesystem and can run at once. On by
+   * default: a pin nobody answers falls back to the shared queue after a schedule-to-start bound,
+   * so the worst it costs is that wait. `OPENCODE_TEMPORAL_STEP_AFFINITY=0` turns it off. */
+  readonly stepAffinity?: boolean
 }
 
 export class Service extends Context.Service<Service, Interface>()("@opencode/temporal/Config") {}
@@ -43,6 +49,7 @@ const given = (name: string) => process.env[name] !== undefined && process.env[n
 const onOff = (name: string, fallback: boolean) => (given(name) ? process.env[name] === "1" : fallback)
 
 export const fromEnv = (): Interface => {
+  const stepAffinity = process.env.OPENCODE_TEMPORAL_STEP_AFFINITY !== "0"
   const worktreeAffinity = process.env.OPENCODE_TEMPORAL_WORKTREE_AFFINITY === "1"
   const sharedStore = !!process.env.OPENCODE_DB_URL
   return {
@@ -54,8 +61,9 @@ export const fromEnv = (): Interface => {
     stepped: onOff("OPENCODE_TEMPORAL_STEPPED", false),
     worktreeAffinity,
     worktree: process.env.OPENCODE_TEMPORAL_WORKTREE,
-    // Serial by default only where two hosts can end up writing one step's tree: a shared store,
-    // with nothing keeping the step's tools on one worker.
-    serialTools: onOff("OPENCODE_TEMPORAL_SERIAL_TOOLS", !worktreeAffinity && sharedStore),
+    stepAffinity,
+    // Serial by default only where two hosts can end up writing one step's tree: a shared store with
+    // neither kind of affinity keeping the step's tools on one worker.
+    serialTools: onOff("OPENCODE_TEMPORAL_SERIAL_TOOLS", !stepAffinity && !worktreeAffinity && sharedStore),
   }
 }

packages/temporal/src/executor.ts

@@ -1,6 +1,7 @@
 export * as SessionExecutionTemporal from "./executor"
 
 import { fileURLToPath } from "node:url"
+import { hostname } from "node:os"
 import { Effect, Layer, Option } from "effect"
 import { Client, Connection, WithStartWorkflowOperation } from "@temporalio/client"
 // Imported lazily inside the worker branch: the worker package drags webpack and swc (it bundles
@@ -16,7 +17,7 @@ import { SessionExecution } from "@opencode-ai/core/session/execution"
 import { makeStepActivities, makeSteppedTurnActivities } from "./activities"
 import { makeDrains } from "./drain"
 import { makeSteppedDrains } from "./stepped-drain"
-import { queueForWorktree } from "./queue"
+import { queueForWorktree, queueForWorker } from "./queue"
 import { Database } from "@opencode-ai/core/database/database"
 import { ProjectTable } from "@opencode-ai/core/project/sql"
 import { eq } from "drizzle-orm"
@@ -75,6 +76,11 @@ const layer = Layer.effect(
     // Which queue a worker polls. With affinity off this is the one shared queue and any worker can
     // draw any session, rebuilding the tree if it has to.
     const POLL_QUEUE = AFFINITY ? queueForWorktree(TASK_QUEUE, SERVED_WORKTREE) : TASK_QUEUE
+    // The queue this worker polls on its own, so a step can be sent back to it. Keyed by host as
+    // well as directory: two containers serve `/project` and share none of it. Only workers have
+    // one, and only they report it, so a client-only process never pins a step to itself.
+    const STEP_QUEUE =
+      HOST_WORKER && config.stepAffinity !== false ? queueForWorker(TASK_QUEUE, hostname(), SERVED_WORKTREE) : undefined
     // Which queue a session's workflow runs on. Keyed on the PROJECT worktree, not the session's
     // directory: `worktrees.ensure` rebuilds the project tree, so keying on the directory a session
     // happened to start in would split one physical tree across a queue per subdirectory, and a
@@ -101,7 +107,7 @@ const layer = Layer.effect(
     const drains = makeDrains({ store, locations, ctx, events, worktrees })
     // The stepped mode's three drains. Registered unconditionally: which mode a session runs is a
     // property of its workflow input, so a worker has to be able to serve either.
-    const stepped = makeSteppedDrains({ inSession: drains.inSession })
+    const stepped = makeSteppedDrains({ inSession: drains.inSession, stepQueue: STEP_QUEUE })
 
     // Worker connection (native) hosts the runTurnStep activity + the workflow. Skipped in
     // client-only role so serve can run without an embedded worker.
@@ -135,6 +141,29 @@ const layer = Layer.effect(
           await runHandle.catch(() => {})
         }),
       )
+
+      // A second poller, on this worker's own queue, for the steps pinned to it. Activities only:
+      // the workflow runs wherever it was started, and only the work that has to come back here is
+      // addressed here. Without it a pin has nobody to answer it and every step pays the
+      // schedule-to-start wait before falling back.
+      if (STEP_QUEUE) {
+        const pinnedWorker = yield* Effect.promise(() =>
+          Worker.create({
+            connection: nativeConn,
+            namespace: NAMESPACE,
+            taskQueue: STEP_QUEUE,
+            activities: { ...makeStepActivities(drains.stepDrain), ...makeSteppedTurnActivities(stepped) },
+          }),
+        )
+        const pinnedHandle = pinnedWorker.run()
+        pinnedHandle.catch(() => {})
+        yield* Effect.addFinalizer(() =>
+          Effect.promise(async () => {
+            pinnedWorker.shutdown()
+            await pinnedHandle.catch(() => {})
+          }),
+        )
+      }
     }
 
     // Worker-only process: it hosts activities but drives no workflows, so the client methods are

packages/temporal/src/queue.ts

@@ -40,3 +40,22 @@ export const queueForWorktree = (base: string, directory: string): string => {
   const digest = createHash("sha256").update(canonical).digest("hex").slice(0, DIGEST_LENGTH)
   return `${base}-wt-${digest}`
 }
+
+/**
+ * The queue a worker polls on its own, alongside the shared one.
+ *
+ * This is what a step is pinned to: the tools of one step write the tree the worker that made the
+ * model call is standing in, so they are sent back to it rather than to whoever is free. Keyed by
+ * host as well as directory, unlike the worktree queue above, because two hosts can serve the same
+ * path without sharing a byte of it: a key on the path alone routes a step to a host whose tree is
+ * a different tree.
+ *
+ * Stable across a restart, so a worker that comes back keeps serving what it served. A pin nobody
+ * answers can release unstarted dispatches after their schedule-to-start bound, subject to
+ * the other pinned attempts' outcomes.
+ */
+export const queueForWorker = (base: string, host: string, directory: string): string => {
+  const canonical = resolve(directory).replace(/[/\\]+$/, "")
+  const digest = createHash("sha256").update(`${host}\0${canonical}`).digest("hex").slice(0, DIGEST_LENGTH)
+  return `${base}-w-${digest}`
+}

packages/temporal/src/stepped-drain.ts

@@ -30,6 +30,11 @@ export type ModelCallDrainResult =
       /** The event-log token this attempt claimed. The tool and seal activities of this step must
        * publish under it, so it travels with the calls instead of being minted again. */
       readonly owner: string
+      /** The queue this worker polls on its own, when it has one. The tools of this step write the
+       * tree this worker is standing in, so sending them here keeps them on it. Absent when the
+       * worker was not given a queue of its own, and never required: the step falls back to the
+       * shared queue and the tree is rebuilt there. */
+      readonly queue?: string
     }
 
 export interface ToolCallDrainInput {
@@ -53,9 +58,12 @@ export interface SealDrainInput {
 
 export interface SteppedDrainDeps {
   readonly inSession: InSession
+  /** The queue this worker polls on its own, reported by the model call so the rest of the step can
+   * be sent back to it. Absent when the worker has none. */
+  readonly stepQueue?: string
 }
 
-export const makeSteppedDrains = ({ inSession }: SteppedDrainDeps) => {
+export const makeSteppedDrains = ({ inSession, stepQueue }: SteppedDrainDeps) => {
   // Only the model call claims the log: it is the writer that supersedes a previous attempt, and
   // the rest of the step rides its token.
   const modelCallDrain = async (
@@ -78,7 +86,7 @@ export const makeSteppedDrains = ({ inSession }: SteppedDrainDeps) => {
           (result): ModelCallDrainResult =>
             result === undefined || result.kind === "settled"
               ? { kind: "settled", result: toStepResult(input.step, result?.result) }
-              : { ...result, owner: input.owner },
+              : { ...result, owner: input.owner, ...(stepQueue === undefined ? {} : { queue: stepQueue }) },
         ),
       ),
     )

packages/temporal/src/stepped-turn.ts

@@ -10,7 +10,7 @@
 // because a workflow cannot consume a stream. The tools of one step still run concurrently with
 // each other; what is lost is the overlap between the model and its own tools.
 
-import { ActivityFailure, type ApplicationFailure } from "@temporalio/workflow"
+import { ActivityFailure, type ApplicationFailure, TimeoutFailure } from "@temporalio/workflow"
 import { HALTED_FAILURE_TYPE } from "./protocol"
 import type { StepDrainInput, StepDrainResult } from "./drain"
 import type {
@@ -32,6 +32,13 @@ import type {
 export const isHaltFailure = (error: unknown) =>
   error instanceof ActivityFailure && (error.cause as ApplicationFailure | undefined)?.type === HALTED_FAILURE_TYPE
 
+// This permits migration only when pinned dispatches have no automatic retries.
+// A later attempt can time out in the queue after an earlier attempt took effect.
+export const isUnclaimedFailure = (error: unknown) =>
+  error instanceof ActivityFailure &&
+  error.cause instanceof TimeoutFailure &&
+  error.cause.timeoutType === "SCHEDULE_TO_START"
+
 /** The three activities a stepped turn drives. */
 export interface SteppedActivities {
   readonly runModelCall: (input: ModelCallDrainInput) => Promise<ModelCallDrainResult>
@@ -54,8 +61,14 @@ export interface SteppedTurnDeps {
   readonly log?: (message: string, attributes: Record<string, unknown>) => void
   /** Run the calls one at a time. Each tool ships the tree from the host that ran it, so two on two
    * hosts each publish a tree without the other's work and the second is refused, leaving its work
-   * stranded there. Serial is what moving files between hosts costs. */
+   * stranded there. Serial is what moving files between hosts costs, and it is what pinning a
+   * step's tools to one worker buys back. */
   readonly serial?: boolean
+  /** The same activities, addressed to one worker's own queue. A step's tools write the tree the
+   * model call's worker is standing in, so keeping them there is what lets them run at once: they
+   * see each other's writes through the filesystem rather than through the store. Only offered a
+   * queue the model call reported. Only a dispatch that queue never started may move off it. */
+  readonly pinnedTo?: (queue: string) => Pick<SteppedActivities, "runToolCall" | "sealStep">
   /** Run something where the driver's cancellation cannot reach it. An interrupt landing during the
    * tool phase otherwise leaves the step with no ending published at all, so a follower waiting on
    * the turn never hears it stop. */
@@ -67,18 +80,81 @@ export interface SteppedTurnDeps {
  * of a whole-step activity.
  */
 export const makeSteppedTurn =
-  ({ activities, isCancellation, isHalt, log, serial, nonCancellable }: SteppedTurnDeps) =>
+  ({ activities, isCancellation, isHalt, log, serial, nonCancellable, pinnedTo }: SteppedTurnDeps) =>
   async (input: StepDrainInput): Promise<StepDrainResult> => {
     const model = await activities.runModelCall(input)
     // A crashed step finalized from the log, or the recovery gate finding no work: the step is over
     // and there is nothing to dispatch or seal.
     if (model.kind === "settled") return model.result
 
+    // The worker that made the model call, when it offered its own queue. Everything else in this
+    // step goes to it first, because it is the host holding the tree the tools are about to write.
+    const pinned = model.queue && pinnedTo ? pinnedTo(model.queue) : undefined
+    let unclaimed = false
+    let uncertain: { error: unknown } | undefined
+    let stopped: { error: unknown } | undefined
+    const observeStop = (error: unknown): never => {
+      if (isCancellation(error) || isHalt(error)) stopped = { error }
+      throw error
+    }
+    // Shared dispatches must wait for the pinned batch because the hosts do not share a worktree.
+    let shared: Promise<unknown> = Promise.resolve()
+    const pendingPins = new Set<Promise<unknown>>()
+    const onShared = <A>(run: (on: SteppedActivities) => Promise<A>, allowStopped = false): Promise<A> => {
+      const next = shared.then(async () => {
+        // Queue saturation can leave a sibling running on the pinned host.
+        await Promise.allSettled(pendingPins)
+        if (uncertain) throw uncertain.error
+        if (stopped && !allowStopped) throw stopped.error
+        return run(activities).catch(observeStop)
+      })
+      shared = next.then(
+        () => undefined,
+        () => undefined,
+      )
+      return next
+    }
+    // A timeout settles the workflow promise; it does not stop the tool process behind it. So a
+    // pinned attempt that started is not evidence that its directory is free, and the rest of the
+    // step stays off that host until the process is known to have stopped or its workspace is its
+    // own. Only a dispatch nobody started moves, and only after every pinned sibling has settled.
+    const viaPinned = async <A>(
+      run: (on: Pick<SteppedActivities, "runToolCall" | "sealStep">) => Promise<A>,
+      allowStopped = false,
+    ): Promise<A> => {
+      if (stopped && !allowStopped) throw stopped.error
+      if (uncertain) throw uncertain.error
+      if (!pinned) return run(activities).catch(observeStop)
+      if (unclaimed) return onShared(run, allowStopped)
+      const attempt = run(pinned)
+      pendingPins.add(attempt)
+      try {
+        return await attempt
+      } catch (error) {
+        if (isCancellation(error) || isHalt(error)) {
+          stopped = { error }
+          throw error
+        }
+        if (!isUnclaimedFailure(error)) {
+          uncertain = { error }
+          throw error
+        }
+        unclaimed = true
+        log?.("the pinned activity did not start; remaining calls move to the shared queue", {
+          sessionID: input.sessionID,
+          step: model.step,
+        })
+        return onShared(run, allowStopped)
+      } finally {
+        pendingPins.delete(attempt)
+      }
+    }
+
     // Each call is its own unit of work. A tool that fails outright does not take the turn with it:
     // the seal closes its call as an error and the model gets to react, which is better than losing
     // the step. A cancel and a user halt are different, and both have to propagate.
     const dispatch = (call: (typeof model.calls)[number]) =>
-      activities.runToolCall({ sessionID: input.sessionID, call, owner: model.owner })
+      viaPinned((on) => on.runToolCall({ sessionID: input.sessionID, call, owner: model.owner }))
     const dispatched: PromiseSettledResult<ToolCallDrainResult>[] = []
     if (serial) {
       // One at a time, and still settled rather than thrown, so a tool that fails does not take the
@@ -108,7 +184,7 @@ export const makeSteppedTurn =
       needsContinuation: stopped ? false : model.needsContinuation,
       owner: model.owner,
     })
-    const seal = (stopped: boolean) => activities.sealStep(sealing(stopped))
+    const seal = (stopped: boolean) => viaPinned((on) => on.sealStep(sealing(stopped)), stopped)
 
     for (const outcome of dispatched) {
       if (outcome.status !== "rejected") continue