@anthropic-ai/sdk 0.95.2 → 0.97.0

package/src/lib/environments/worker.ts

@@ -0,0 +1,395 @@
+import { AnthropicError } from '../../core/error';
+import type { Anthropic } from '../../client';
+import type { BetaSelfHostedWork } from '../../resources/beta/environments/work';
+import { loggerFor, type Logger } from '../../internal/utils/log';
+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 { buildHeaders } from '../../internal/headers';
+import type { BetaRunnableTool } from '../tools/BetaRunnableTool';
+import type { BetaToolRunnerRequestOptions } from '../tools/BetaToolRunner';
+import { SessionToolRunner } from '../tools/SessionToolRunner';
+import { WorkPoller } from './poller';
+import { copyClientForHelper } from '../helper-client';
+// `tools/agent-toolset/node` is Node-only (node:child_process, node:fs, …).
+// Only the type is imported statically (erased at build); the module's values
+// (`setupSkills`, `betaAgentToolset20260401`) are loaded lazily inside the
+// 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';
+
+const HEARTBEAT_DEFAULT_MS = 30_000;
+const NO_HEARTBEAT_SENTINEL = 'NO_HEARTBEAT';
+
+/**
+ * Either a fixed tool array or a factory invoked once per claimed session with
+ * that session's {@link AgentToolContext} — use the factory form to bind
+ * `betaAgentToolset20260401` (or any tool that needs the workdir / session
+ * id) to the right session.
+ */
+export type EnvironmentWorkerTools =
+  | Array<BetaRunnableTool>
+  | ((ctx: AgentToolContext) => Array<BetaRunnableTool>);
+
+export interface EnvironmentWorkerOptions {
+  client: Anthropic;
+  /**
+   * The self-hosted environment to poll for work. Required by
+   * {@link EnvironmentWorker.run}; not used by {@link EnvironmentWorker.handleItem}.
+   */
+  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}.
+   */
+  environmentKey?: string;
+  /**
+   * Tools to expose to each claimed session. Defaults to
+   * `betaAgentToolset20260401(ctx)` (the standard `agent_toolset_20260401` set
+   * bound to the per-session {@link AgentToolContext}).
+   */
+  tools?: EnvironmentWorkerTools;
+  /** Base directory for the per-session {@link AgentToolContext}. Defaults to `process.cwd()`. */
+  workdir?: string;
+  /** Forwarded to the per-session {@link AgentToolContext}. */
+  unrestrictedPaths?: boolean;
+  /** Forwarded to {@link SessionToolRunner} (`maxIdleMs`). */
+  maxIdleMs?: number;
+  /** Forwarded to the {@link WorkPoller}. */
+  workerId?: string;
+  /** External abort signal; aborting it ends the run. */
+  signal?: AbortSignal;
+  /**
+   * Extra per-request options merged into every call this worker issues — the
+   * work poll/ack/heartbeat/stop control-plane calls and the per-session
+   * SessionToolRunner's stream/list/send. Mirrors what
+   * `client.beta.messages.toolRunner` accepts: custom `headers` (e.g. a proxy's
+   * auth/routing headers) reach all of them. The worker owns the abort signals,
+   * so a `signal` here is ignored — use {@link EnvironmentWorkerOptions.signal}.
+   */
+  requestOptions?: BetaToolRunnerRequestOptions;
+}
+
+/**
+ * Options for {@link EnvironmentWorker.handleItem}. Every field falls back to the
+ * matching `ANTHROPIC_*` environment variable — the ones the
+ * `ant worker poll --on-work` command sets for the process it spawns — when not
+ * passed explicitly.
+ */
+export interface HandleItemOptions {
+  /** Work item id. Falls back to `ANTHROPIC_WORK_ID`. */
+  workId?: string;
+  /** Self-hosted environment id. Falls back to `ANTHROPIC_ENVIRONMENT_ID`. */
+  environmentId?: string;
+  /** Session id. Falls back to `ANTHROPIC_SESSION_ID`. */
+  sessionId?: string;
+  /**
+   * The environment key used to authenticate every per-session call. Resolution
+   * order: this option, then the worker's own `environmentKey`, then
+   * `ANTHROPIC_ENVIRONMENT_KEY`.
+   */
+  environmentKey?: 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 self-hosted environment runner, composed from the control-plane
+ * {@link WorkPoller} and the per-session {@link SessionToolRunner}.
+ *
+ * For each claimed `session` work item it: builds the per-session
+ * {@link AgentToolContext}, downloads the session agent's skills
+ * (`setupSkills`), then runs a {@link SessionToolRunner} for the session
+ * *while* heartbeating the work-item lease in parallel; 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.
+ *
+ * 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
+ * just want the per-item flow without the poll loop — with no arguments it reads
+ * the `ANTHROPIC_*` env vars that command sets.
+ *
+ * Construct it via `client.beta.environments.work.worker({ ... })` (or
+ * `new EnvironmentWorker({ client, ... })` directly).
+ *
+ * @example
+ * ```ts
+ * // Long-running daemon: poll for work, serve each session, loop.
+ * await client.beta.environments.work
+ *   .worker({ environmentId, environmentKey, workdir: '/workspace' })
+ *   .run(AbortSignal.timeout(60 * 60_000));
+ *
+ * // Already-claimed item (e.g. inside `ant worker poll --on-work ...`):
+ * await client.beta.environments.work.worker({ workdir: '/workspace' }).handleItem();
+ * ```
+ */
+export class EnvironmentWorker {
+  readonly client: Anthropic;
+  readonly environmentId: string | undefined;
+  readonly environmentKey: string | undefined;
+  readonly tools: EnvironmentWorkerTools | undefined;
+  readonly workdir: string;
+  readonly unrestrictedPaths: boolean | undefined;
+  readonly maxIdleMs: number | undefined;
+  readonly workerId: string | undefined;
+  readonly requestOptions: BetaToolRunnerRequestOptions | undefined;
+  readonly #signal: AbortSignal | undefined;
+
+  constructor(opts: EnvironmentWorkerOptions) {
+    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.maxIdleMs = opts.maxIdleMs;
+    this.workerId = opts.workerId;
+    this.requestOptions = opts.requestOptions;
+    this.#signal = opts.signal;
+  }
+
+  /**
+   * Poll the environment and service each claimed session until the supplied
+   * signal (or the one passed to the constructor) aborts. Throws if
+   * `environmentId` / `environmentKey` were not provided to the constructor.
+   */
+  async run(signal?: AbortSignal): Promise<void> {
+    const { environmentId, environmentKey } = this;
+    if (environmentId === undefined || environmentKey === undefined) {
+      throw new AnthropicError(
+        'EnvironmentWorker.run: environmentId and environmentKey are required to poll for work',
+      );
+    }
+    const externalSignal = signal ?? this.#signal;
+    const poller = new WorkPoller({
+      client: this.client,
+      environmentId,
+      environmentKey,
+      ...(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.
+      autoStop: false,
+    });
+
+    for await (const work of poller) {
+      await this.#handleItem(work, environmentKey, poller.signal);
+    }
+  }
+
+  /**
+   * Service a single, already-claimed work item without the poll loop: build the
+   * per-session {@link AgentToolContext} (workdir from this worker's options),
+   * download the session agent's skills (`setupSkills`), run a
+   * {@link SessionToolRunner} for the session while heartbeating the work-item
+   * lease in parallel, and force-stop the work item on exit (whether the runner
+   * finishes normally, throws, or the heartbeat loop signals shutdown).
+   *
+   * 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
+   * work id / environment id / session id each fall back to `ANTHROPIC_WORK_ID` /
+   * `ANTHROPIC_ENVIRONMENT_ID` / `ANTHROPIC_SESSION_ID` (the env vars that
+   * command sets) when not passed; the environment key resolves from this
+   * 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.
+   */
+  async handleItem(opts?: HandleItemOptions): Promise<void> {
+    const workId = opts?.workId ?? readEnv('ANTHROPIC_WORK_ID');
+    const environmentId = opts?.environmentId ?? readEnv('ANTHROPIC_ENVIRONMENT_ID');
+    const sessionId = opts?.sessionId ?? readEnv('ANTHROPIC_SESSION_ID');
+    const environmentKey =
+      opts?.environmentKey ?? this.environmentKey ?? readEnv('ANTHROPIC_ENVIRONMENT_KEY');
+
+    if (!workId) {
+      throw new AnthropicError('handleItem: workId is required — pass it or set ANTHROPIC_WORK_ID');
+    }
+    if (!environmentId) {
+      throw new AnthropicError(
+        'handleItem: environmentId is required — pass it or set ANTHROPIC_ENVIRONMENT_ID',
+      );
+    }
+    if (!sessionId) {
+      throw new AnthropicError('handleItem: sessionId is required — pass it or set ANTHROPIC_SESSION_ID');
+    }
+    if (!environmentKey) {
+      throw new AnthropicError(
+        'handleItem: environmentKey is required — pass it, construct the worker with it, or set ANTHROPIC_ENVIRONMENT_KEY',
+      );
+    }
+
+    const work: ClaimedWork = {
+      id: workId,
+      environment_id: environmentId,
+      data: { type: 'session', id: sessionId },
+    };
+    await this.#handleItem(work, environmentKey, opts?.signal ?? this.#signal);
+  }
+
+  /**
+   * 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.
+   */
+  async #handleItem(
+    work: ClaimedWork,
+    environmentKey: string,
+    externalSignal: AbortSignal | undefined,
+  ): Promise<void> {
+    const log = loggerFor(this.client);
+    // 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).
+    const sessionClient = copyClientForHelper(this.client, {
+      authToken: environmentKey,
+      helper: 'environments-worker',
+    });
+
+    // The poller runs with `autoStop: false`, so the per-item handler is the
+    // 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 } : {}),
+    };
+    // 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 heartbeatPromise = heartbeatLoop(sessionClient, work, ctrl, log, this.requestOptions).catch((e) => {
+      if (!ctrl.signal.aborted) log.error('heartbeat loop failed', { work_id: work.id, error: String(e) });
+      ctrl.abort();
+    });
+
+    try {
+      const runner = new SessionToolRunner(sessionId, {
+        client: sessionClient,
+        tools,
+        ...(this.maxIdleMs !== undefined ? { maxIdleMs: this.maxIdleMs } : {}),
+        ...(this.requestOptions !== undefined ? { requestOptions: this.requestOptions } : {}),
+        signal: ctrl.signal,
+      });
+      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.
+      }
+    } finally {
+      ctrl.abort();
+      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);
+    }
+  }
+}
+
+/** Force-stop a claimed work item, swallowing the 409 that means it's already stopped. */
+async function forceStop(
+  client: Anthropic,
+  work: Pick<BetaSelfHostedWork, 'id' | 'environment_id'>,
+  log: Logger,
+  requestOptions?: BetaToolRunnerRequestOptions,
+): Promise<void> {
+  try {
+    await client.beta.environments.work.stop(
+      work.id,
+      { environment_id: work.environment_id, force: true },
+      // Caller's headers pass through; the helper-tag header is on the scoped
+      // sub-client's default_headers via copyClientForHelper, so no per-call
+      // re-stamping needed.
+      { ...requestOptions, headers: buildHeaders([requestOptions?.headers]) },
+    );
+  } catch (e) {
+    if (!isStatus(e, 409)) {
+      log.error('force-stop on exit failed', { work_id: work.id, error: String(e) });
+    }
+  }
+}
+
+/**
+ * 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, or on a permanent heartbeat failure.
+ */
+async function heartbeatLoop(
+  client: Anthropic,
+  work: Pick<BetaSelfHostedWork, 'id' | 'environment_id'>,
+  ctrl: AbortController,
+  logger: Logger,
+  requestOptions?: BetaToolRunnerRequestOptions,
+): Promise<void> {
+  let intervalMs = HEARTBEAT_DEFAULT_MS;
+  let last = NO_HEARTBEAT_SENTINEL;
+  const beat = async (): Promise<void> => {
+    try {
+      const resp = await client.beta.environments.work.heartbeat(
+        work.id,
+        { environment_id: work.environment_id, expected_last_heartbeat: last },
+        { ...requestOptions, headers: buildHeaders([requestOptions?.headers]), signal: ctrl.signal },
+      );
+      last = resp.last_heartbeat;
+      if (resp.ttl_seconds > 0) {
+        intervalMs = Math.max(1_000, Math.min((resp.ttl_seconds * 1000) / 2, HEARTBEAT_DEFAULT_MS));
+      }
+      if (resp.state === 'stopping' || resp.state === 'stopped') {
+        logger.info('heartbeat signals shutdown', { work_id: work.id, state: resp.state });
+        ctrl.abort();
+      }
+      if (!resp.lease_extended) {
+        logger.warn('lease not extended, shutting down', { work_id: work.id });
+        ctrl.abort();
+      }
+    } catch (e) {
+      // An abort throws to unwind the caller (the `heartbeatLoop(...).catch`
+      // in `#handleItem`) rather than returning early.
+      ctrl.signal.throwIfAborted();
+      if (isFatal4xx(e)) {
+        logger.error('permanent heartbeat failure', { work_id: work.id, error: String(e) });
+        ctrl.abort();
+        throw e;
+      }
+      logger.warn('transient heartbeat failure', { work_id: work.id, error: String(e) });
+    }
+  };
+
+  await beat();
+  while (!ctrl.signal.aborted) {
+    await sleep(intervalMs, ctrl.signal);
+    ctrl.signal.throwIfAborted();
+    await beat();
+  }
+}

package/src/lib/helper-client.ts

@@ -0,0 +1,89 @@
+import { AnthropicError } from '../core/error';
+import type { Anthropic } from '../client';
+import { buildHeaders, type HeadersLike, type NullableHeaders } from '../internal/headers';
+
+/**
+ * Shared util for building a runner-helper-bound sub-client.
+ *
+ * The work poller, the environment worker, and the session tool runner each
+ * need to issue requests authenticated by a per-helper credential (a
+ * self-hosted environment key, today) rather than the parent client's own
+ * `X-Api-Key`, *and* tagged with their own `x-stainless-helper` telemetry
+ * value. Each wants to inherit the parent's full configuration — `timeout`,
+ * `maxRetries`, `fetch`, `fetchOptions`, custom `defaultHeaders`,
+ * `defaultQuery` — and override only the auth + telemetry bits.
+ *
+ * {@link copyClientForHelper} is the one shared construction.
+ */
+
+/**
+ * The closed set of `x-stainless-helper` telemetry tags the runner helpers
+ * stamp on outgoing requests. Constrained as a string union so a typo at any
+ * call site is a type error rather than silently mistagged telemetry.
+ */
+export type HelperTag = 'environments-work-poller' | 'environments-worker' | 'session-tool-runner';
+
+interface ClientInternalAccess {
+  _options: { defaultHeaders?: HeadersLike };
+  _authState?: { extraHeaders?: Record<string, string> };
+}
+
+/**
+ * Return a `withOptions()` clone of `client` set up for use *by* one of the
+ * runner helpers: authenticated with `authToken` as Bearer credentials, with
+ * the parent's `X-Api-Key` cleared, and tagged with the helper's
+ * `x-stainless-helper` value on every outgoing request.
+ *
+ * The returned sub-client inherits the parent's full configuration
+ * (`baseURL`, `timeout`, `maxRetries`, `fetch`, `fetchOptions`, custom
+ * `defaultHeaders`, `defaultQuery`). Overrides applied:
+ *
+ * - `authToken: authToken` — the new credential.
+ * - `apiKey: null` — the parent's `X-Api-Key` is cleared. `withOptions`
+ *   inherits the parent's `apiKey` by default; without this, both
+ *   `X-Api-Key` *and* `Authorization: Bearer …` would land on the wire.
+ *   `client.ts` only triggers the env-var fallback when `apiKey === undefined`,
+ *   so explicit `null` is honored.
+ * - `credentials: undefined` — opts the clone out of any inherited
+ *   credentials/config/profile so the explicit bearer is the unambiguous auth.
+ * - `baseURL: client.baseURL` — pins the parent's resolved host (auth override otherwise resets it).
+ * - `defaultHeaders` is rebuilt as `parent._authState.extraHeaders ⊕ parent.defaultHeaders ⊕
+ *   {'x-stainless-helper': helper}`. `withOptions` *replaces* (does not
+ *   merge) `defaultHeaders`, so we merge here so any custom headers the
+ *   caller set on the parent client survive on the sub-client.
+ */
+export function copyClientForHelper<T extends Anthropic>(
+  client: T,
+  { authToken, helper }: { authToken: string; helper: HelperTag },
+): T {
+  if (!authToken) {
+    throw new AnthropicError(
+      `copyClientForHelper: expected a non-empty authToken but received ${JSON.stringify(authToken)}`,
+    );
+  }
+  const internal = client as unknown as ClientInternalAccess;
+  const parentDefaults = internal._options.defaultHeaders;
+  // Carry the parent's credential/profile headers; strip the auth ones (we re-auth below).
+  const parentAuthExtraHeaders = internal._authState?.extraHeaders;
+  const inheritedAuthExtraHeaders: Record<string, string> | undefined =
+    parentAuthExtraHeaders ?
+      Object.fromEntries(
+        Object.entries(parentAuthExtraHeaders).filter(([name]) => {
+          const lower = name.toLowerCase();
+          return lower !== 'authorization' && lower !== 'x-api-key';
+        }),
+      )
+    : undefined;
+  const defaultHeaders: NullableHeaders = buildHeaders([
+    inheritedAuthExtraHeaders,
+    parentDefaults,
+    { 'x-stainless-helper': helper },
+  ]);
+  return client.withOptions({
+    apiKey: null,
+    authToken,
+    baseURL: client.baseURL,
+    credentials: undefined,
+    defaultHeaders,
+  }) as T;
+}

package/src/lib/tools/BetaRunnableTool.ts

@@ -11,8 +11,14 @@ import {
   BetaToolTextEditor20250124,
   BetaToolTextEditor20250429,
   BetaToolTextEditor20250728,
+  BetaToolUnion,
   BetaToolUseBlock,
 } from '../../resources/beta';
+import type {
+  BetaManagedAgentsAgentCustomToolUseEvent,
+  BetaManagedAgentsAgentToolUseEvent,
+} from '../../resources/beta/sessions/events';
+import { ToolError } from './ToolError';
 
 export type Promisable<T> = T | Promise<T>;
 
@@ -33,8 +39,32 @@ export type BetaClientRunnableToolType =
   | BetaToolTextEditor20250429
   | BetaToolTextEditor20250728;
 
+/**
+ * The tool-use that triggered a {@link BetaRunnableTool.run}:
+ *
+ * - from `client.beta.messages.toolRunner`, a Messages `tool_use` content block;
+ * - from `client.beta.sessions.events.toolRunner`, the `agent.tool_use` /
+ *   `agent.custom_tool_use` session event.
+ *
+ * The shapes overlap on the common fields (`id`, `name`, `input`), so code that
+ * only reads those works without narrowing; narrow on the shape (e.g. `'type' in
+ * x`) when you need surface-specific properties.
+ */
+export type BetaToolUse =
+  | BetaToolUseBlock
+  | BetaManagedAgentsAgentToolUseEvent
+  | BetaManagedAgentsAgentCustomToolUseEvent;
+
 export type BetaToolRunContext = {
-  toolUseBlock: BetaToolUseBlock;
+  /** The tool-use that triggered this run. See {@link BetaToolUse}. */
+  toolUse: BetaToolUse;
+  /**
+   * @deprecated Renamed to `toolUse`. Also note that for
+   * `client.beta.sessions.events.toolRunner` this is the `agent.tool_use` /
+   * `agent.custom_tool_use` *event*, not a Messages content block, despite the
+   * name — which is why it was renamed.
+   */
+  toolUseBlock: BetaToolUse;
   signal?: AbortSignal | null | undefined;
 };
 
@@ -46,4 +76,55 @@ export type BetaRunnableTool<Input = any> = BetaClientRunnableToolType & {
     context?: BetaToolRunContext,
   ) => Promisable<string | Array<BetaToolResultContentBlockParam>>;
   parse: (content: unknown) => Input;
+  /**
+   * Optional cleanup hook for tools that hold process-level resources (e.g. a
+   * persistent shell). `SessionToolRunner` (`client.beta.sessions.events.toolRunner`)
+   * calls it once when iteration ends.
+   */
+  close?: () => Promisable<void>;
 };
+
+/**
+ * Resolve the registry key for a tool — the name the model addresses it by.
+ * MCP toolsets are keyed on `mcp_server_name`; every other tool on `name`.
+ * Shared so the tool-name lookup is identical across `toolRunner()` surfaces.
+ */
+export function toolName(tool: BetaToolUnion | BetaRunnableTool): string {
+  return 'name' in tool ? tool.name : tool.mcp_server_name;
+}
+
+/**
+ * Format a thrown value into tool-result content: a {@link ToolError} carries
+ * its own structured content, anything else becomes an `Error: <message>`
+ * string. Shared so every `toolRunner()` surface reports tool failures the
+ * same way to the model.
+ */
+export function toolErrorContent(e: unknown): string | Array<BetaToolResultContentBlockParam> {
+  return e instanceof ToolError ? e.content : `Error: ${e instanceof Error ? e.message : String(e)}`;
+}
+
+/** Outcome of {@link runRunnableTool}: the content to post back and whether it is an error. */
+export interface RunnableToolOutcome {
+  content: string | Array<BetaToolResultContentBlockParam>;
+  isError: boolean;
+}
+
+/**
+ * Run a {@link BetaRunnableTool} end-to-end: parse the raw input, invoke `run`,
+ * and format any thrown value via {@link toolErrorContent}. Shared so the
+ * parse → run → catch → format pipeline is identical across `toolRunner()`
+ * surfaces.
+ */
+export async function runRunnableTool(
+  tool: BetaRunnableTool,
+  rawInput: unknown,
+  context: BetaToolRunContext,
+): Promise<RunnableToolOutcome> {
+  try {
+    const input = tool.parse ? tool.parse(rawInput) : rawInput;
+    const content = await tool.run(input, context);
+    return { content, isError: false };
+  } catch (e) {
+    return { content: toolErrorContent(e), isError: true };
+  }
+}

package/src/lib/tools/BetaToolRunner.ts

@@ -6,26 +6,10 @@ import { BetaMessage, BetaMessageParam, BetaToolUnion, MessageCreateParams } fro
 import { BetaMessageStream } from '../BetaMessageStream';
 import { RequestOptions } from '../../internal/request-options';
 import { buildHeaders } from '../../internal/headers';
+import { promiseWithResolvers } from '../../internal/utils/promise';
 import { CompactionControl, DEFAULT_SUMMARY_PROMPT, DEFAULT_TOKEN_THRESHOLD } from './CompactionControl';
 import { collectStainlessHelpers } from '../stainless-helper-header';
 
-/**
- * Just Promise.withResolvers(), which is not available in all environments.
- */
-function promiseWithResolvers<T>(): {
-  promise: Promise<T>;
-  resolve: (value: T) => void;
-  reject: (reason?: any) => void;
-} {
-  let resolve: (value: T) => void;
-  let reject: (reason?: any) => void;
-  const promise = new Promise<T>((res, rej) => {
-    resolve = res;
-    reject = rej;
-  });
-  return { promise, resolve: resolve!, reject: reject! };
-}
-
 /**
  * A ToolRunner handles the automatic conversation loop between the assistant and tools.
  *
@@ -484,6 +468,7 @@ async function generateToolResponse(
         }
 
         const result = await tool.run(input, {
+          toolUse: toolUse,
           toolUseBlock: toolUse,
           signal: requestOptions?.signal,
         });