@anthropic-ai/sdk 0.118.0 → 0.120.0

package/src/lib/environments/worker.ts

@@ -1,11 +1,14 @@
-import { AnthropicError } from '../../core/error';
+import { AnthropicError, APIError } from '../../core/error';
 import type { Anthropic } from '../../client';
-import type { BetaSelfHostedWork } from '../../resources/beta/environments/work';
+import type { BetaSelfHostedWork, BetaWorkSecret } from '../../resources/beta/environments/work';
 import { loggerFor, type Logger } from '../../internal/utils/log';
+import { fromBase64 } from '../../internal/utils/base64';
+import { decodeUTF8 } from '../../internal/utils/bytes';
 import { readEnv } from '../../internal/utils/env';
 import { sleep } from '../../internal/utils/sleep';
 import { isFatal4xx, isStatus } from '../../internal/utils/backoff';
 import { linkAbort } from '../../internal/utils/abort';
+import { isObj } from '../../internal/utils/values';
 import { buildHeaders } from '../../internal/headers';
 import type { BetaRunnableTool } from '../tools/BetaRunnableTool';
 import type { BetaToolRunnerRequestOptions } from '../tools/BetaToolRunner';
@@ -18,7 +21,9 @@ import { copyClientForHelper } from '../helper-client';
 // per-item handler. That keeps this file free of Node-only deps in the static
 // import graph, which is what lets `client.beta.environments.work.worker()`
 // exist as a resource method without pulling Node built-ins into the SDK core.
-import type { AgentToolContext } from '../../tools/agent-toolset/node';
+import type { AgentToolContext, MemoryDeleteMode, SessionMemoryStores } from '../../tools/agent-toolset/node';
+import { checkMemorySyncInterval } from '../../tools/agent-toolset/sync-interval';
+import type { BetaManagedAgentsSession } from '../../resources/beta/sessions/sessions';
 
 const HEARTBEAT_DEFAULT_MS = 30_000;
 const HEARTBEAT_TTL_DEFAULT_MS = 90_000;
@@ -42,10 +47,11 @@ export interface EnvironmentWorkerOptions {
    */
   environmentId?: string;
   /**
-   * The environment key — the single credential for the runner. It authenticates
-   * the work-poll calls and every per-session call (event stream, lease
-   * heartbeat, force-stop). Required by {@link EnvironmentWorker.run}; falls back
-   * to `ANTHROPIC_ENVIRONMENT_KEY` in {@link EnvironmentWorker.handleItem}.
+   * The environment key — the worker's standing credential: polling always
+   * uses it, and per-session calls fall back to it when a claimed item's
+   * `secret` doesn't yield a sessions token. Required by
+   * {@link EnvironmentWorker.run}; falls back to `ANTHROPIC_ENVIRONMENT_KEY` in
+   * {@link EnvironmentWorker.handleItem}.
    */
   environmentKey?: string;
   /**
@@ -60,12 +66,41 @@ export interface EnvironmentWorkerOptions {
   tools?: EnvironmentWorkerTools;
   /** Base directory for the per-session {@link AgentToolContext}. Defaults to `process.cwd()`. */
   workdir?: string;
-  /** Forwarded to the per-session {@link AgentToolContext}. */
+  /**
+   * @deprecated No longer accepted: the file tools are always confined to
+   * `workdir` plus `allowedRoots`, which is neither behavior this flag used to
+   * select, so passing either value throws. Remove it; list extra directories
+   * in {@link AgentToolContext.allowedRoots}. The property is removed in a
+   * future release.
+   */
   unrestrictedPaths?: boolean;
   /** Forwarded to the per-session {@link AgentToolContext} (`maxFileBytes`). */
   maxFileBytes?: number | null;
   /** Forwarded to {@link SessionToolRunner} (`maxIdleMs`). */
   maxIdleMs?: number;
+  /**
+   * How often (milliseconds) to sync the session's attached memory stores back
+   * while it runs — checked after each dispatched tool call, plus one final
+   * sync when the session ends cleanly. Defaults to
+   * `DEFAULT_MEMORY_SYNC_INTERVAL_MS` (15s); the constructor throws for
+   * values below `MIN_MEMORY_SYNC_INTERVAL_MS` (5s). Every teardown also runs
+   * a push-only flush of changed files; that flush and the final sync are
+   * each bounded by `MEMORY_FLUSH_TIMEOUT_MS`, and a warning is logged when
+   * either bound cuts work off. `null` disables memory download and sync
+   * entirely. Memory stores are only touched for work items whose `secret`
+   * carries a `sessions_token`; while sync is enabled, an item without one
+   * fails when its session has memory stores attached, because those stores
+   * cannot be mounted without the token. With `null` the same item runs,
+   * without memory, and nothing is logged — turning sync off is the
+   * operator's explicit choice.
+   */
+  memorySyncIntervalMs?: number | null;
+  /**
+   * Whether local file deletions may delete on the server — see
+   * {@link MemoryDeleteMode}. Uploads and pulls are unaffected.
+   * Defaults to `"enabled"`.
+   */
+  memorySyncDeletions?: MemoryDeleteMode;
   /** Forwarded to the {@link WorkPoller}. */
   workerId?: string;
   /** External abort signal; aborting it ends the run. */
@@ -100,12 +135,62 @@ export interface HandleItemOptions {
    * `ANTHROPIC_ENVIRONMENT_KEY`.
    */
   environmentKey?: string;
+  /**
+   * The work item's per-item `secret` payload from the poll response. Falls
+   * back to `ANTHROPIC_WORK_SECRET`. Unlike the others it is optional — when
+   * present, the sessions token extracted from it is preferred as the Bearer
+   * credential for this item's heartbeat / force-stop / skill-download /
+   * session calls; when absent (or undecodable) those calls use the
+   * environment key.
+   */
+  workSecret?: string;
   /** External abort signal; aborting it ends the run. Defaults to the constructor's signal. */
   signal?: AbortSignal;
 }
 
-/** The fields of {@link BetaSelfHostedWork} the per-item flow reads. */
-type ClaimedWork = Pick<BetaSelfHostedWork, 'id' | 'environment_id' | 'data'>;
+/**
+ * The fields of {@link BetaSelfHostedWork} the per-item flow reads. `secret` is
+ * declared here rather than picked because the generated model does not carry
+ * it — only the poll response populates it.
+ */
+type ClaimedWork = Pick<BetaSelfHostedWork, 'id' | 'environment_id' | 'data'> & {
+  secret?: string | null;
+};
+
+/** True when the session has at least one memory store attached. */
+function hasMemoryStore(session: BetaManagedAgentsSession): boolean {
+  return session.resources.some((r) => r.type === 'memory_store');
+}
+
+/**
+ * Extract the per-item sessions token from a work item's `secret` payload.
+ *
+ * The `secret` the poll response populates is not itself a credential: it is a
+ * URL-safe base64 JSON payload matching {@link BetaWorkSecret} — the
+ * `sessions_token` (the bearer for this item's work lifecycle and
+ * session-level calls) plus fields this worker does not consume. Returns the
+ * sessions token, or `null` (meaning: fall back to the environment key) when
+ * the payload is missing, doesn't decode, or carries no token. Never log the
+ * payload or anything extracted from it.
+ */
+export function sessionsTokenFromSecret(secret: string | null | undefined): string | null {
+  if (!secret) return null;
+  let parsed: unknown;
+  try {
+    // The payload may arrive URL-safe and without base64 padding; normalize
+    // both before decoding.
+    const normalized = secret.replace(/-/g, '+').replace(/_/g, '/');
+    const padded = normalized.padEnd(Math.ceil(normalized.length / 4) * 4, '=');
+    parsed = JSON.parse(decodeUTF8(fromBase64(padded)));
+  } catch {
+    return null;
+  }
+  if (parsed === null || typeof parsed !== 'object' || Array.isArray(parsed)) return null;
+  // The payload is untrusted input, so the token is still checked at runtime
+  // rather than trusted to match the schema.
+  const token = (parsed as Partial<BetaWorkSecret>).sessions_token;
+  return typeof token === 'string' && token !== '' ? token : null;
+}
 
 /**
  * The self-hosted environment runner, composed from the control-plane
@@ -115,9 +200,17 @@ type ClaimedWork = Pick<BetaSelfHostedWork, 'id' | 'environment_id' | 'data'>;
  * {@link AgentToolContext}, downloads the session agent's skills
  * (`setupSkills`), then runs a {@link SessionToolRunner} for the session
  * while heartbeating the work-item lease on the same event loop; on exit it
- * force-stops the work item, cleans up the downloaded skills, and loops to the
- * next one. The lease heartbeat reports `state === "stopping"` / a lost lease
- * back into the run by aborting the session runner.
+ * force-stops the work item (unless the lease was lost, in which case the item
+ * is left to whoever holds it now), cleans up the downloaded skills, and loops
+ * to the next one. The lease heartbeat reports `state === "stopping"` / a lost
+ * lease back into the run by aborting the session runner.
+ *
+ * The `environmentKey` is the worker's standing credential. When a claimed
+ * work item carries a per-item `secret` (a short-lived payload the poll
+ * response may populate), the sessions token extracted from it is preferred
+ * over the environment key for that item's heartbeat / force-stop /
+ * skill-download / session calls; polling itself always uses the environment
+ * key, and items without a usable secret fall back to it entirely.
  *
  * Use {@link EnvironmentWorker.handleItem} if you already hold a claimed work
  * item (e.g. a `worker poll --on-work` script handed one to a fresh process) and
@@ -144,22 +237,42 @@ export class EnvironmentWorker {
   readonly environmentKey: string | undefined;
   readonly tools: EnvironmentWorkerTools | undefined;
   readonly workdir: string;
+  /** @deprecated Never set; see {@link EnvironmentWorkerOptions.unrestrictedPaths}. */
   readonly unrestrictedPaths: boolean | undefined;
   readonly maxFileBytes: number | null | undefined;
   readonly maxIdleMs: number | undefined;
+  readonly memorySyncIntervalMs: number | null | undefined;
+  readonly memorySyncDeletions: MemoryDeleteMode;
   readonly workerId: string | undefined;
   readonly requestOptions: BetaToolRunnerRequestOptions | undefined;
   readonly #signal: AbortSignal | undefined;
 
+  constructor(opts: Omit<EnvironmentWorkerOptions, 'unrestrictedPaths'>);
+  /** @deprecated `unrestrictedPaths` is no longer accepted — see {@link EnvironmentWorkerOptions.unrestrictedPaths}. */
+  constructor(opts: Omit<EnvironmentWorkerOptions, 'unrestrictedPaths'> & { unrestrictedPaths: boolean }); // help language servers see deprecation
   constructor(opts: EnvironmentWorkerOptions) {
+    if (opts.unrestrictedPaths !== undefined) {
+      throw new AnthropicError(
+        'The `unrestrictedPaths` option you passed to EnvironmentWorker (or ' +
+          'client.beta.environments.work.worker()) is no longer supported. ' +
+          "The worker's file tools (read, write, edit, glob, grep) are now always confined to `workdir` " +
+          "plus the session's memory folders. Remove `unrestrictedPaths` from your options; to let the " +
+          'file tools reach any other directory, add it to `AgentToolContext.allowedRoots` from a ' +
+          '`tools` factory.',
+      );
+    }
     this.client = opts.client;
     this.environmentId = opts.environmentId;
     this.environmentKey = opts.environmentKey;
     this.tools = opts.tools;
     this.workdir = opts.workdir ?? process.cwd();
-    this.unrestrictedPaths = opts.unrestrictedPaths;
     this.maxFileBytes = opts.maxFileBytes;
     this.maxIdleMs = opts.maxIdleMs;
+    if (opts.memorySyncIntervalMs != null) {
+      checkMemorySyncInterval(opts.memorySyncIntervalMs, 'memorySyncIntervalMs');
+    }
+    this.memorySyncIntervalMs = opts.memorySyncIntervalMs;
+    this.memorySyncDeletions = opts.memorySyncDeletions ?? 'enabled';
     this.workerId = opts.workerId;
     this.requestOptions = opts.requestOptions;
     this.#signal = opts.signal;
@@ -185,13 +298,23 @@ export class EnvironmentWorker {
       ...(this.workerId !== undefined ? { workerId: this.workerId } : {}),
       ...(externalSignal ? { signal: externalSignal } : {}),
       ...(this.requestOptions !== undefined ? { requestOptions: this.requestOptions } : {}),
-      // The per-item handler force-stops every work item on exit; let it be the
-      // single owner of `work.stop` rather than double-posting from the poller.
+      // The per-item handler stops or releases every work item on exit; let it
+      // be the single owner of `work.stop` rather than double-posting from the
+      // poller.
       autoStop: false,
     });
 
     for await (const work of poller) {
-      await this.#handleItem(work, environmentKey, poller.signal);
+      try {
+        await this.#handleItem(work, environmentKey, poller.signal);
+      } catch (e) {
+        // One bad item fails that item, not the worker: the handler's teardown
+        // already stopped or released it, so the next poll claims the next
+        // item. A store directory left behind by a killed worker would
+        // otherwise crashloop this process forever.
+        if (poller.signal?.aborted) throw e;
+        loggerFor(this.client).error('work item failed', { work_id: work.id, error: String(e) });
+      }
     }
   }
 
@@ -201,7 +324,9 @@ export class EnvironmentWorker {
    * download the session agent's skills (`setupSkills`), run a
    * {@link SessionToolRunner} for the session while heartbeating the work-item
    * lease, and force-stop the work item on exit (whether the runner finishes
-   * normally, throws, or the heartbeat loop signals shutdown).
+   * normally, throws, or the control plane signals shutdown). The one
+   * exception is a lost lease: the item then belongs to the queue or another
+   * worker and is left alone.
    *
    * Use this when something else does the claiming — e.g. a `worker poll
    * --on-work` script that hands an already-claimed item to a fresh process. The
@@ -211,7 +336,16 @@ export class EnvironmentWorker {
    * option, then the worker's own `environmentKey`, then
    * `ANTHROPIC_ENVIRONMENT_KEY`. With no arguments inside that command it just
    * works. Throws a clear error naming the first of the four required values
-   * still missing after resolution.
+   * still missing after resolution. Throws `SessionMemoryError` when the
+   * session has memory stores attached but they cannot be mounted — the work
+   * item carried no sessions token (unless `memorySyncIntervalMs` turned
+   * memory off), or a store failed to download.
+   *
+   * `workSecret` is the work item's per-item `secret` payload from the poll
+   * response, falling back to `ANTHROPIC_WORK_SECRET`; unlike the others it is
+   * optional — when present, the sessions token extracted from it is preferred
+   * as the Bearer credential for this item's heartbeat / force-stop / session
+   * calls; when absent (or undecodable) those calls use the environment key.
    */
   async handleItem(opts?: HandleItemOptions): Promise<void> {
     const workId = opts?.workId ?? readEnv('ANTHROPIC_WORK_ID');
@@ -219,6 +353,9 @@ export class EnvironmentWorker {
     const sessionId = opts?.sessionId ?? readEnv('ANTHROPIC_SESSION_ID');
     const environmentKey =
       opts?.environmentKey ?? this.environmentKey ?? readEnv('ANTHROPIC_ENVIRONMENT_KEY');
+    // `||` rather than `??` so an empty option still falls through to the env
+    // var and then to null (matching how `readEnv` treats empty values).
+    const workSecret = opts?.workSecret || readEnv('ANTHROPIC_WORK_SECRET') || null;
 
     if (!workId) {
       throw new AnthropicError('handleItem: workId is required — pass it or set ANTHROPIC_WORK_ID');
@@ -240,6 +377,7 @@ export class EnvironmentWorker {
     const work: ClaimedWork = {
       id: workId,
       environment_id: environmentId,
+      secret: workSecret,
       data: { type: 'session', id: sessionId },
     };
     await this.#handleItem(work, environmentKey, opts?.signal ?? this.#signal);
@@ -248,8 +386,14 @@ export class EnvironmentWorker {
   /**
    * The per-item body shared by {@link EnvironmentWorker.run}'s poll loop and
    * {@link EnvironmentWorker.handleItem}: run a {@link SessionToolRunner} for the
-   * work item's session while heartbeating its lease, force-stopping on exit.
-   * Non-session work items are ignored.
+   * work item's session while heartbeating its lease, force-stopping on exit
+   * unless the lease was lost. Non-session work items are ignored.
+   *
+   * When the poll response carried a per-item `secret` (a short-lived payload
+   * scoped to this work item), the sessions token extracted from it is
+   * preferred over `environmentKey` as the Bearer credential for those
+   * per-item calls; a missing/undecodable secret falls back to
+   * `environmentKey` unchanged.
    */
   async #handleItem(
     work: ClaimedWork,
@@ -257,14 +401,27 @@ export class EnvironmentWorker {
     externalSignal: AbortSignal | undefined,
   ): Promise<void> {
     const log = loggerFor(this.client);
+    // The per-item credential: the sessions token carried inside the work
+    // item's secret payload when the server issued one, otherwise the
+    // environment key. Never log this value.
+    const sessionsToken = sessionsTokenFromSecret(work.secret);
+    if (work.secret && sessionsToken === null) {
+      log.warn(
+        'work item carried a secret payload but no sessions token could be extracted; ' +
+          'falling back to the environment key',
+        { work_id: work.id },
+      );
+    }
+    const itemCredential = sessionsToken ?? environmentKey;
     // Every per-session call — the SessionToolRunner event stream/list/send, the
-    // lease heartbeat, and the work force-stop — authenticates with the
-    // environment key. Scope a client to it once and thread that through.
-    // `copyClientForHelper` also clears the parent's `apiKey`, so the sub-client
-    // emits *only* the bearer credential on the wire (a plain
-    // `withOptions({authToken})` would leave `X-Api-Key` set as well).
+    // lease heartbeat, the skill download, and the work force-stop —
+    // authenticates with the per-item credential. Scope a client to it once and
+    // thread that through. `copyClientForHelper` also clears the parent's
+    // `apiKey`, so the sub-client emits *only* the bearer credential on the
+    // wire (a plain `withOptions({authToken})` would leave `X-Api-Key` set as
+    // well).
     const sessionClient = copyClientForHelper(this.client, {
-      authToken: environmentKey,
+      authToken: itemCredential,
       helper: 'environments-worker',
     });
 
@@ -272,68 +429,191 @@ export class EnvironmentWorker {
     // single owner of `work.stop` for every claimed item.
     const sessionId = work.data.id;
 
-    const ctx: AgentToolContext = {
-      workdir: this.workdir,
-      client: this.client,
-      sessionId,
-      ...(this.unrestrictedPaths !== undefined ? { unrestrictedPaths: this.unrestrictedPaths } : {}),
-      ...(this.maxFileBytes !== undefined ? { maxFileBytes: this.maxFileBytes } : {}),
-    };
-    // Lazily load the Node-only toolset module — see the import note at the top.
-    const agentToolset = await import('../../tools/agent-toolset/node');
-    let cleanupSkills: () => Promise<void> = async () => {};
-    try {
-      cleanupSkills = await agentToolset.setupSkills(ctx);
-    } catch (e) {
-      log.warn('skill setup failed', { session_id: sessionId, work_id: work.id, error: String(e) });
-    }
-    const tools =
-      typeof this.tools === 'function' ?
-        this.tools(ctx)
-      : this.tools ?? agentToolset.betaAgentToolset20260401(ctx);
-
     // A per-session controller: aborts when the supplied signal aborts, when the
     // session runner finishes, or when the lease heartbeat says to stop.
     const ctrl = new AbortController();
     const detachExternal = linkAbort(externalSignal, ctrl);
+    const lease = new Lease(ctrl);
 
-    // Inert until iterated below.
-    const runner = new SessionToolRunner(sessionId, {
-      client: sessionClient,
-      tools,
-      ...(this.maxIdleMs !== undefined ? { maxIdleMs: this.maxIdleMs } : {}),
-      ...(this.requestOptions !== undefined ? { requestOptions: this.requestOptions } : {}),
-      signal: ctrl.signal,
-    });
+    // Lazily load the Node-only toolset module — see the import note at the top.
+    const agentToolset = await import('../../tools/agent-toolset/node');
 
+    // Start the lease heartbeat BEFORE the session fetch and the skill /
+    // memory downloads: those can take longer than the lease TTL, and an
+    // unheartbeated lease lapsing mid-download would let another worker
+    // reclaim the item and serve the same session (split-brain).
+    //
     // Each heartbeat reports the lease TTL the server is enforcing; it becomes
     // the runner's tool-result send retry window so a send keeps retrying
-    // exactly as long as the lease could still be live.
-    const heartbeatPromise = heartbeatLoop(sessionClient, work, ctrl, log, this.requestOptions, (ttlMs) =>
-      runner._setSendRetryWindow(ttlMs),
-    ).catch((e) => {
+    // exactly as long as the lease could still be live. The runner is only
+    // built after the downloads, so hold the latest TTL until then.
+    let leaseTtlMs: number | undefined;
+    let runner: SessionToolRunner | undefined;
+    const heartbeatPromise = heartbeatLoop(sessionClient, work, lease, log, this.requestOptions, (ttlMs) => {
+      leaseTtlMs = ttlMs;
+      runner?._setSendRetryWindow(ttlMs);
+    }).catch((e) => {
       if (!ctrl.signal.aborted) log.error('heartbeat loop failed', { work_id: work.id, error: String(e) });
       ctrl.abort();
     });
 
+    let cleanupSkills: () => Promise<void> = async () => {};
+    let stores: SessionMemoryStores | undefined;
+    let cleanEnd = false;
     try {
+      if (work.data.type !== 'session') {
+        log.debug('skipping non-session work item', { work_id: work.id, type: work.data.type });
+        return;
+      }
+      // One session fetch, shared by the skills download and the memory-store
+      // download — two fetches could disagree about the attached resources.
+      // A failed fetch fails the work item (the teardown below still stops
+      // or releases it).
+      const session: BetaManagedAgentsSession = await sessionClient.beta.sessions.retrieve(sessionId);
+      // Only with the session in hand can we tell one that simply has no
+      // memory from one whose memory we cannot mount. Turning memory off
+      // with the interval knob is a deliberate opt-out and stays quiet.
+      if (sessionsToken === null && this.memorySyncIntervalMs !== null && hasMemoryStore(session)) {
+        throw new agentToolset.SessionMemoryError(
+          `cannot mount the session's memories: the work item carried no sessions token ` +
+            `(work_id=${work.id}, session_id=${sessionId}); ` +
+            'the memory endpoints reject the environment key, so the poller must issue a per-item ' +
+            '`secret` carrying `sessions_token`, or set `memorySyncIntervalMs: null` to run without memory',
+        );
+      }
+
+      const ctx: AgentToolContext = {
+        workdir: this.workdir,
+        // The scoped sub-client, not the parent: the skill download
+        // `setupSkills` performs for this session rides the same per-item
+        // credential as every other per-item call.
+        client: sessionClient,
+        session,
+        ...(this.maxFileBytes !== undefined ? { maxFileBytes: this.maxFileBytes } : {}),
+      };
+      try {
+        cleanupSkills = await agentToolset.setupSkills(ctx);
+      } catch (e) {
+        log.warn('skill setup failed', { session_id: sessionId, work_id: work.id, error: String(e) });
+      }
+
+      // Memory stores: the memory_stores endpoints accept the per-item sessions
+      // token but reject the environment key, so download and sync only run when
+      // the item carried a usable secret (and the interval is set).
+      // `sessionClient` is already scoped to that token then, so the memory
+      // calls ride the same sub-client. A store that cannot be materialised
+      // throws `SessionMemoryError` out of `download` and fails the item.
+      if (sessionsToken !== null && this.memorySyncIntervalMs !== null) {
+        stores = new agentToolset.SessionMemoryStores(sessionClient, {
+          workdir: this.workdir,
+          ...(this.memorySyncIntervalMs !== undefined ? { syncIntervalMs: this.memorySyncIntervalMs } : {}),
+          syncDeletions: this.memorySyncDeletions,
+        });
+        await stores.download(session);
+        // A store mounted outside the workdir must stay reachable by the file
+        // tools; read-only stores still refuse writes.
+        ctx.allowedRoots = stores.roots;
+        ctx.readOnlyRoots = stores.readOnlyRoots;
+      } else {
+        log.debug('memory stores disabled for this item', { work_id: work.id });
+      }
+
+      const tools =
+        typeof this.tools === 'function' ?
+          this.tools(ctx)
+        : this.tools ?? agentToolset.betaAgentToolset20260401(ctx);
+
+      runner = new SessionToolRunner(sessionId, {
+        client: sessionClient,
+        tools,
+        ...(this.maxIdleMs !== undefined ? { maxIdleMs: this.maxIdleMs } : {}),
+        ...(this.requestOptions !== undefined ? { requestOptions: this.requestOptions } : {}),
+        signal: ctrl.signal,
+      });
+      if (leaseTtlMs !== undefined) runner._setSendRetryWindow(leaseTtlMs);
       for await (const _ of runner) {
         // Drive the runner to completion; per-call observability is not part
         // of this composition's surface — use `SessionToolRunner` directly
         // (via `client.beta.sessions.events.toolRunner`) if you want it.
+        if (stores) await stores.syncIfDue();
       }
+      // Only a clean stream end earns the last full sync; it runs in the
+      // teardown below.
+      cleanEnd = !ctrl.signal.aborted;
     } finally {
-      ctrl.abort();
+      // The heartbeat keeps the lease alive until this teardown is done.
+      try {
+        // cleanupSkills first, so its failure cannot skip the memory flush.
+        await cleanupSkills().catch((e) => {
+          log.warn('skill cleanup failed', { session_id: sessionId, work_id: work.id, error: String(e) });
+        });
+      } finally {
+        if (stores) {
+          const boundMs = agentToolset.MEMORY_FLUSH_TIMEOUT_MS;
+          if (cleanEnd) {
+            const finishCutOff = await withTimeout(stores.finish(), boundMs);
+            if (finishCutOff) {
+              log.warn(
+                `final memory sync cut off after ${boundMs}ms; the flush that follows still uploads changed files`,
+                { session_id: sessionId, work_id: work.id },
+              );
+            }
+          }
+          // Also after finish(): it swallows its own failures, and a
+          // clean flush is a no-op.
+          const flushBound = new AbortController();
+          const flushCutOff = await withTimeout(stores.flushWrites(flushBound.signal), boundMs);
+          if (flushCutOff) {
+            flushBound.abort();
+            log.warn(
+              `memory flush cut off after ${boundMs}ms; changed files it had not uploaded yet are not saved`,
+              { session_id: sessionId, work_id: work.id },
+            );
+          }
+          await stores.dispose().catch((e) => {
+            log.warn('memory store cleanup failed', {
+              session_id: sessionId,
+              work_id: work.id,
+              error: String(e),
+            });
+          });
+        }
+      }
+      lease.finish('runner_done');
       detachExternal();
       await heartbeatPromise;
-      await cleanupSkills().catch((e) => {
-        log.warn('skill cleanup failed', { session_id: sessionId, work_id: work.id, error: String(e) });
-      });
-      await forceStop(sessionClient, work, log, this.requestOptions);
+      // Stop only an item this worker still holds — after a lost lease it
+      // belongs to the queue or another worker.
+      if (lease.lost) {
+        log.info('lease lost; released without stopping it', { session_id: sessionId, work_id: work.id });
+      } else {
+        await forceStop(sessionClient, work, log, this.requestOptions);
+      }
     }
   }
 }
 
+/**
+ * Resolve when `p` settles or `ms` elapses — `true` when `ms` elapsed
+ * first. A timed-out `p` keeps running — JS cannot cancel a promise.
+ */
+async function withTimeout(p: Promise<void>, ms: number): Promise<boolean> {
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  try {
+    return await Promise.race([
+      p.then(
+        () => false,
+        () => false,
+      ),
+      new Promise<boolean>((resolve) => {
+        timer = setTimeout(() => resolve(true), ms);
+      }),
+    ]);
+  } finally {
+    if (timer !== undefined) clearTimeout(timer);
+  }
+}
+
 /** Force-stop a claimed work item, swallowing the 409 that means it's already stopped. */
 async function forceStop(
   client: Anthropic,
@@ -357,18 +637,67 @@ async function forceStop(
   }
 }
 
+/** Why heartbeating of a work item ended, as recorded on its {@link Lease}. */
+type LeaseEndReason =
+  | 'runner_done'
+  | 'control_plane_stop'
+  | 'lease_lost'
+  | 'heartbeat_rejected'
+  | 'assumed_lost';
+
 /**
- * Keep the work-item lease alive while a session is being served. Aborts `ctrl`
- * when the control plane reports the work is `stopping`/`stopped`, when the
- * lease is no longer extended, on a permanent heartbeat failure, or when no
+ * This worker's view of one work-item lease: the per-item abort signal plus
+ * why heartbeating ended. The first recorded reason wins, so a run aborted
+ * *because* the lease was lost still reads as lost afterwards; an abort with
+ * no recorded reason (the external signal) is not lost.
+ */
+class Lease {
+  readonly #ctrl: AbortController;
+  #endReason: LeaseEndReason | undefined;
+
+  constructor(ctrl: AbortController) {
+    this.#ctrl = ctrl;
+  }
+
+  get signal(): AbortSignal {
+    return this.#ctrl.signal;
+  }
+
+  finish(reason: LeaseEndReason): void {
+    this.#endReason ??= reason;
+    this.#ctrl.abort();
+  }
+
+  /** True once the item belongs to the queue or another worker. */
+  get lost(): boolean {
+    return this.#endReason === 'lease_lost' || this.#endReason === 'assumed_lost';
+  }
+}
+
+/** The server's view of the lease carried by a 412 heartbeat response, or empty if absent. */
+function serverLeaseState(e: unknown): Record<string, unknown> {
+  let node: unknown = e instanceof APIError ? e.error : undefined;
+  for (const key of ['error', 'details', 'current_state']) {
+    if (!isObj(node)) return {};
+    node = node[key];
+  }
+  return isObj(node) ? node : {};
+}
+
+/**
+ * Keep the work-item lease alive while a session is being served. Runs until
+ * `lease` ends, and ends it itself when the control plane reports the work is
+ * `stopping`/`stopped` or no longer extends the lease, when a heartbeat is
+ * rejected (a 412 means the lease already belongs to someone else), or when no
  * heartbeat has succeeded for longer than the lease ttl (the lease is assumed
- * lost). Each heartbeat call is cut off after the current beat interval so a
- * hung request cannot outlive the lease it is meant to renew.
+ * lost, so two runners don't end up serving the same work). Each heartbeat
+ * call is cut off after the current beat interval so a hung request cannot
+ * outlive the lease it is meant to renew.
  */
 async function heartbeatLoop(
   client: Anthropic,
   work: Pick<BetaSelfHostedWork, 'id' | 'environment_id'>,
-  ctrl: AbortController,
+  lease: Lease,
   logger: Logger,
   requestOptions?: BetaToolRunnerRequestOptions,
   /** Called with the server-reported lease TTL after every successful beat. */
@@ -382,7 +711,7 @@ async function heartbeatLoop(
     // Not the request `timeout` option: the core client retries timeouts, so
     // it would not bound the call as a whole.
     const beatCtrl = new AbortController();
-    const detach = linkAbort(ctrl.signal, beatCtrl);
+    const detach = linkAbort(lease.signal, beatCtrl);
     const cutoff = setTimeout(() => beatCtrl.abort(), intervalMs);
     try {
       const resp = await client.beta.environments.work.heartbeat(
@@ -399,19 +728,30 @@ async function heartbeatLoop(
       }
       if (resp.state === 'stopping' || resp.state === 'stopped') {
         logger.info('heartbeat signals shutdown', { work_id: work.id, state: resp.state });
-        ctrl.abort();
+        lease.finish('control_plane_stop');
       }
       if (!resp.lease_extended) {
         logger.warn('lease not extended, shutting down', { work_id: work.id });
-        ctrl.abort();
+        lease.finish('control_plane_stop');
       }
     } catch (e) {
       // An abort throws to unwind the caller (the `heartbeatLoop(...).catch`
       // in `#handleItem`) rather than returning early.
-      ctrl.signal.throwIfAborted();
+      lease.signal.throwIfAborted();
+      if (isStatus(e, 412)) {
+        const server = serverLeaseState(e);
+        logger.error('lease lost: heartbeat precondition failed', {
+          work_id: work.id,
+          server_state: server['state'],
+          server_ttl_seconds: server['ttl_seconds'],
+          server_last_heartbeat: server['last_heartbeat'],
+        });
+        lease.finish('lease_lost');
+        return;
+      }
       if (isFatal4xx(e)) {
         logger.error('permanent heartbeat failure', { work_id: work.id, error: String(e) });
-        ctrl.abort();
+        lease.finish('heartbeat_rejected');
         throw e;
       }
       if (Date.now() - lastSuccessMs > ttlMs) {
@@ -420,7 +760,7 @@ async function heartbeatLoop(
           ttl_ms: ttlMs,
           error: String(e),
         });
-        ctrl.abort();
+        lease.finish('assumed_lost');
         return;
       }
       logger.warn('transient heartbeat failure', { work_id: work.id, error: String(e) });
@@ -431,9 +771,9 @@ async function heartbeatLoop(
   };
 
   await beat();
-  while (!ctrl.signal.aborted) {
-    await sleep(intervalMs, ctrl.signal);
-    ctrl.signal.throwIfAborted();
+  while (!lease.signal.aborted) {
+    await sleep(intervalMs, lease.signal);
+    lease.signal.throwIfAborted();
     await beat();
   }
 }