@genkit-ai/ai/lib/agent.mjs.map

Genkit AI framework generative AI APIs.

{"version":3,"sources":["../src/agent.ts"],"sourcesContent":["/**\n * Copyright 2026 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport {\n  GenkitError,\n  deepEqual,\n  defineAction,\n  defineBidiAction,\n  getContext,\n  run,\n  z,\n  type Action,\n  type ActionContext,\n  type ActionFnArg,\n  type BidiAction,\n} from '@genkit-ai/core';\nimport { Channel } from '@genkit-ai/core/async';\nimport type { Registry } from '@genkit-ai/core/registry';\nimport {\n  createAgentAPI,\n  type AgentAPI,\n  type AgentTransport,\n  type SnapshotLookup,\n} from './agent-core.mjs';\n\nimport { parseSchema, toJsonSchema } from '@genkit-ai/core/schema';\nimport {\n  setCustomMetadataAttribute,\n  setCustomMetadataAttributes,\n} from '@genkit-ai/core/tracing';\nimport {\n  AgentAbortRequestSchema,\n  AgentAbortResponseSchema,\n  AgentInitSchema,\n  AgentInputSchema,\n  AgentOutputSchema,\n  AgentStreamChunkSchema,\n  GetSnapshotRequestSchema,\n  type AgentInit,\n  type AgentInput,\n  type AgentResult,\n  type AgentStreamChunk,\n} from './agent-types.mjs';\nimport { generateStream } from './generate.mjs';\nimport { diff, type JsonPatch } from './json-patch.mjs';\nimport { MessageData } from './model-types.mjs';\nimport { type ToolRequestPart, type ToolResponsePart } from './parts.mjs';\nimport {\n  definePrompt,\n  type PromptAction,\n  type PromptConfig,\n} from './prompt.mjs';\nimport { InMemorySessionStore } from './session-stores.mjs';\nimport {\n  Session,\n  SessionSnapshot,\n  SessionSnapshotSchema,\n  SessionState,\n  SessionStore,\n  reserveSnapshotId,\n  runWithSession,\n  type AgentFinishReason,\n  type Artifact,\n  type SessionSnapshotInput,\n  type SessionStoreOptions,\n} from './session.mjs';\n\n// Re-export the shared agent/session wire schemas + types from their canonical\n// home (./agent-types.ts) so existing imports from './agent.mjs' (and the\n// package barrel) keep working.\nexport {\n  AgentAbortRequestSchema,\n  AgentAbortResponseSchema,\n  AgentInitSchema,\n  AgentInputSchema,\n  AgentOutputSchema,\n  AgentResultSchema,\n  AgentStreamChunkSchema,\n  GetSnapshotRequestSchema,\n  JsonPatchOperationSchema,\n  JsonPatchSchema,\n  TurnEndSchema,\n  type AgentInit,\n  type AgentInput,\n  type AgentResult,\n  type AgentStreamChunk,\n  type TurnEnd,\n} from './agent-types.mjs';\n\n/**\n * Default interval (ms) at which a detached (background) turn refreshes its\n * pending snapshot's heartbeat. Each beat is a write to the session store.\n */\nconst DEFAULT_HEARTBEAT_INTERVAL_MS = 30_000;\n\n/**\n * Default staleness threshold (ms) after which a `pending` snapshot whose\n * heartbeat has not advanced is reported as `expired` on read. Should be\n * comfortably larger than {@link DEFAULT_HEARTBEAT_INTERVAL_MS} so a single\n * missed beat does not trip expiry.\n */\nconst DEFAULT_HEARTBEAT_TIMEOUT_MS = 60_000;\n\n/**\n * Returns `true` when a snapshot is a `pending` (detached, in-flight) snapshot\n * whose heartbeat is older than `timeoutMs` - i.e. its background worker is\n * presumed dead. A pending snapshot that has not yet written a first heartbeat\n * is not considered expired (the beat may simply not have fired yet).\n */\nfunction isHeartbeatExpired(\n  snapshot: SessionSnapshot,\n  timeoutMs: number = DEFAULT_HEARTBEAT_TIMEOUT_MS\n): boolean {\n  if (snapshot.status !== 'pending' || !snapshot.heartbeatAt) {\n    return false;\n  }\n  const last = Date.parse(snapshot.heartbeatAt);\n  if (Number.isNaN(last)) {\n    return false;\n  }\n  return Date.now() - last > timeoutMs;\n}\n\n/**\n * Result returned by a single turn handler passed to {@link SessionRunner.run}.\n *\n * Returning a `finishReason` lets a custom agent explicitly state why the turn\n * ended (e.g. `interrupted`, `length`). When omitted, no per-turn reason is\n * reported.\n */\nexport interface TurnResult {\n  finishReason?: AgentFinishReason;\n}\n\n/**\n * Per-turn context handed to the handler passed to {@link SessionRunner.run}.\n *\n * The `snapshotId` is *reserved at turn start* (before the handler runs) and is\n * the id the snapshot persisted at turn end will reuse. This lets a handler\n * name external, snapshot-correlated resources - e.g. a git branch / worktree\n * named after the snapshot - up front, then commit them under that id, so a\n * later rollback to the snapshot can restore the external state too.\n */\nexport interface TurnContext {\n  /**\n   * The id the snapshot produced by this turn will be saved under (reserved\n   * ahead of time so it is known before the turn runs).\n   */\n  snapshotId: string;\n  /**\n   * The id of the parent snapshot this turn continues from, or `undefined` on\n   * the first turn of a fresh session.\n   */\n  parentSnapshotId?: string;\n  /** Zero-based index of this turn within the current invocation. */\n  turnIndex: number;\n}\n\n/**\n * Output returned at turn completion.\n */\nexport interface AgentOutput<S = unknown> {\n  sessionId?: string;\n  artifacts?: Artifact[];\n  message?: MessageData;\n  snapshotId?: string;\n  state?: SessionState<S>;\n  finishReason?: AgentFinishReason;\n  /**\n   * Present when `finishReason` is `failed`. Carries the original error\n   * details (RuntimeError shape); `state`/`snapshotId` hold the last-good state.\n   */\n  error?: {\n    status?: string;\n    message: string;\n    details?: any;\n  };\n}\n\n/**\n * Structured error details surfaced on the failure path.\n */\ninterface AgentErrorDetails {\n  status: string;\n  message: string;\n  details?: any;\n}\n\n/**\n * Normalizes a thrown value into the structured error shape used across the\n * agent (in `AgentOutput.error` and `SessionRunner.lastTurnError`).\n */\nfunction toErrorDetails(e: any): AgentErrorDetails {\n  return {\n    status: e?.status || 'INTERNAL',\n    message: e?.message || 'Internal failure',\n    // Only surface explicitly-provided structured details. Never fall back to\n    // the raw thrown value: it is serialized over the wire (AgentOutput.error)\n    // and persisted into snapshots, so leaking it could expose stack traces /\n    // internal state, or break JSON.stringify on circular error objects.\n    details: e?.detail ?? e?.details,\n  };\n}\n\n/**\n * Builds an abort-aware `saveSnapshot` mutator: it skips the write (returns\n * `null`) when the current snapshot was concurrently aborted, otherwise writes\n * `input`. This prevents a \"done\"/\"failed\" write from clobbering an \"aborted\"\n * status set by a concurrent abort.\n */\nfunction abortAwareMutator<S>(input: SessionSnapshotInput<S>) {\n  return (current: SessionSnapshot<S> | undefined) =>\n    current?.status === 'aborted' ? null : input;\n}\n\n/**\n * Asserts that an operation requiring a persistent store is not being invoked\n * on a store-less (client-managed) agent.\n */\nfunction requireStore<S>(\n  store: SessionStore<S> | undefined,\n  operation: string,\n  agentName: string\n): asserts store is SessionStore<S> {\n  if (!store) {\n    throw new GenkitError({\n      status: 'FAILED_PRECONDITION',\n      message: `${operation} requires a persistent store. Provide a 'store' when defining '${agentName}'.`,\n    });\n  }\n}\n\n/**\n * Sets a snapshot's status to `aborted` (unless it already reached a terminal\n * state) and returns its previous status, or `undefined` when the snapshot\n * does not exist.\n */\nasync function abortSnapshotInStore<S>(\n  store: SessionStore<S>,\n  snapshotId: string,\n  options?: SessionStoreOptions\n): Promise<SessionSnapshot['status'] | undefined> {\n  let previousStatus: SessionSnapshot['status'] | undefined;\n  await store.saveSnapshot(\n    snapshotId,\n    (current) => {\n      if (!current) return null;\n      previousStatus = current.status;\n      if (\n        current.status === 'completed' ||\n        current.status === 'failed' ||\n        current.status === 'aborted'\n      ) {\n        return null; // Already terminal - don't override.\n      }\n      return { ...current, status: 'aborted' };\n    },\n    options\n  );\n  return previousStatus;\n}\n\n/**\n * Executor responsible for running turns over input streams and persisting state.\n */\nexport class SessionRunner<State = unknown> {\n  readonly session: Session<State>;\n  readonly inputCh: AsyncIterable<AgentInput>;\n\n  turnIndex: number = 0;\n  public onEndTurn?: (\n    snapshotId?: string,\n    finishReason?: AgentFinishReason\n  ) => void;\n  public onDetach?: (snapshotId: string) => void;\n  public newSnapshotId?: string;\n  /** The finish reason of the most recently completed turn. */\n  public lastTurnFinishReason?: AgentFinishReason;\n  /**\n   * Error details of the most recent failed turn. Set when a turn throws and\n   * the runner resolves gracefully instead of propagating the exception.\n   */\n  public lastTurnError?: AgentErrorDetails;\n  /**\n   * The state the most recently *successful* turn left behind. On a failed\n   * turn this is the state the failed turn started with - the last-good state\n   * returned to the caller (for client-managed agents).\n   */\n  public lastGoodState?: SessionState<State>;\n  /**\n   * The snapshotId of the most recently *successful* (persisted, `done`) turn.\n   * On a failed turn this is the last-good snapshot the caller resumes from;\n   * `undefined` when no turn has succeeded yet (e.g. a first-turn failure).\n   */\n  public lastGoodSnapshotId?: string;\n  private lastSnapshot?: SessionSnapshot<State>;\n\n  private lastSnapshotVersion: number = 0;\n\n  private store?: SessionStore<State>;\n  public isDetached: boolean = false;\n  /**\n   * Aborts in-flight turns. When set and aborted, a turn that rejects out of\n   * `generate` is reported as `aborted` (not `failed`) and its failed snapshot\n   * write is skipped (the abort path already persisted the `aborted` status).\n   */\n  private abortSignal?: AbortSignal;\n\n  /**\n   * True until the first `customPatch` chunk of the current turn has been\n   * emitted. The first patch of every turn is a whole-document replace\n   * (re-basing clients that may not share the server's baseline); reset to\n   * `true` at the start of each turn.\n   */\n  public firstCustomPatchInTurn: boolean = true;\n\n  constructor(\n    session: Session<State>,\n    inputCh: AsyncIterable<AgentInput>,\n    options?: {\n      lastSnapshot?: SessionSnapshot<State>;\n      store?: SessionStore<State>;\n      abortSignal?: AbortSignal;\n      onEndTurn?: (\n        snapshotId?: string,\n        finishReason?: AgentFinishReason\n      ) => void;\n      onDetach?: (snapshotId: string) => void;\n    }\n  ) {\n    this.session = session;\n    this.inputCh = inputCh;\n\n    this.lastSnapshot = options?.lastSnapshot;\n    this.store = options?.store;\n    this.abortSignal = options?.abortSignal;\n    this.onEndTurn = options?.onEndTurn;\n    this.onDetach = options?.onDetach;\n\n    // Seed the last-good state with the initial session state so that a\n    // failure on the very first turn still has a valid state to fall back to\n    // (the seed/loaded state, excluding the failed turn's mutations). The\n    // last-good snapshotId is undefined until a turn successfully persists.\n    this.lastGoodState = this.session.getState();\n    this.lastGoodSnapshotId = options?.lastSnapshot?.snapshotId;\n  }\n\n  // ── Session delegate methods ────────────────────────────────────────\n  // These forward to `this.session` so callers can write `sess.addMessages()`\n  // instead of the verbose `sess.session.addMessages()`.\n\n  /** Returns a deep copy of the current session state. */\n  getState(): SessionState<State> {\n    return this.session.getState();\n  }\n\n  /** Retrieves all messages associated with the session. */\n  getMessages(): MessageData[] {\n    return this.session.getMessages();\n  }\n\n  /** Appends messages to the session. */\n  addMessages(messages: MessageData[]): void {\n    this.session.addMessages(messages);\n  }\n\n  /** Overwrites the session messages. */\n  setMessages(messages: MessageData[]): void {\n    this.session.setMessages(messages);\n  }\n\n  /** Retrieves the custom state of the session. */\n  getCustom(): State | undefined {\n    return this.session.getCustom();\n  }\n\n  /** Updates the custom state using a mutator function. */\n  updateCustom(fn: (custom?: State) => State): void {\n    this.session.updateCustom(fn);\n  }\n\n  /** Retrieves the list of artifacts generated during the session. */\n  getArtifacts(): Artifact[] {\n    return this.session.getArtifacts();\n  }\n\n  /** Adds artifacts to the session, deduplicating by name. */\n  addArtifacts(artifacts: Artifact[]): void {\n    this.session.addArtifacts(artifacts);\n  }\n\n  /** Invokes the end-of-turn callback, absorbing errors from a closed stream. */\n  private notifyEndTurn(\n    snapshotId: string | undefined,\n    finishReason?: AgentFinishReason\n  ): void {\n    try {\n      this.onEndTurn?.(snapshotId, finishReason);\n    } catch {\n      // Stream was closed, absorb exception.\n    }\n  }\n\n  /**\n   * Executes the flow handler against incoming input messages sequentially.\n   *\n   * The handler receives the turn's {@link AgentInput} and a {@link TurnContext}\n   * whose `snapshotId` is *reserved up front* - it is the id the snapshot\n   * persisted at turn end will reuse. This lets a handler set up external,\n   * snapshot-correlated state (e.g. a git branch/worktree named after the\n   * snapshot) before generating, then commit it under that id.\n   *\n   * The handler may return a {@link TurnResult} carrying an explicit\n   * `finishReason` for the just-completed turn. When omitted, no per-turn\n   * reason is reported. Failures always report `failed`.\n   */\n  async run(\n    fn: (input: AgentInput, ctx: TurnContext) => Promise<TurnResult | void>\n  ): Promise<void> {\n    for await (const input of this.inputCh) {\n      if (input.message) {\n        this.session.addMessages([input.message]);\n      }\n\n      // The first customPatch of every turn is a whole-document replace that\n      // re-bases clients which may not share the server's baseline.\n      this.firstCustomPatchInTurn = true;\n\n      const parentSnapshotId = this.lastSnapshot?.snapshotId;\n\n      // Reserve the turn's snapshotId up front (when a store is configured) so\n      // the handler can name snapshot-correlated external resources before the\n      // turn runs. The detach path may have already reserved one; reuse it.\n      // The persisted snapshot at turn end reuses this id (maybeSnapshot\n      // prefers `newSnapshotId`).\n      if (this.store && !this.newSnapshotId) {\n        this.newSnapshotId = reserveSnapshotId();\n      }\n\n      const turnSnapshotId = this.newSnapshotId;\n      this.newSnapshotId = undefined;\n\n      const turnContext: TurnContext = {\n        snapshotId: turnSnapshotId!,\n        parentSnapshotId,\n        turnIndex: this.turnIndex,\n      };\n\n      try {\n        await run(`runTurn-${this.turnIndex + 1}`, input, async () => {\n          const turnResult = await fn(input, turnContext);\n          const finishReason = turnResult?.finishReason;\n          this.lastTurnFinishReason = finishReason;\n          this.lastTurnError = undefined;\n\n          const snapshotId = await this.maybeSnapshot(\n            'completed',\n            undefined,\n            turnSnapshotId,\n            finishReason\n          );\n\n          // Capture the state this successful turn produced. This becomes the\n          // last-good state to fall back to if a later turn fails, and its\n          // snapshotId is the last-good snapshot a failed turn resumes from.\n          this.lastGoodState = this.session.getState();\n          this.lastGoodSnapshotId = snapshotId;\n\n          // Tag the turn span with the snapshotId this turn persisted under, so\n          // a trace can correlate the turn with its snapshot (server-managed\n          // agents only; client-managed turns have no snapshotId).\n          if (snapshotId) {\n            setCustomMetadataAttribute('agent:snapshotId', snapshotId);\n          }\n\n          this.notifyEndTurn(snapshotId, finishReason);\n\n          // The turn span's output is the session state this turn produced -\n          // applies to both client- and server-managed agents.\n          return { state: this.session.getState() };\n        });\n        this.turnIndex++;\n      } catch (e: any) {\n        // An aborted turn rejects out of `generate` and lands here. Treat it as\n        // `aborted` rather than `failed`: the abort path already persisted the\n        // `aborted` status (the abort-aware mutator would skip a `failed` write\n        // anyway), so we record the finish reason and skip the failed snapshot\n        // write entirely instead of reporting a spurious error.\n        if (this.abortSignal?.aborted) {\n          this.lastTurnFinishReason = 'aborted';\n          this.lastTurnError = undefined;\n          this.notifyEndTurn(this.lastSnapshot?.snapshotId, 'aborted');\n          break;\n        }\n\n        this.lastTurnFinishReason = 'failed';\n        this.lastTurnError = toErrorDetails(e);\n        const snapshotId = await this.maybeSnapshot(\n          'failed',\n          this.lastTurnError,\n          turnSnapshotId,\n          'failed'\n        );\n        this.notifyEndTurn(snapshotId, 'failed');\n\n        // Graceful failure: rather than propagating the exception (which would\n        // discard the action's final return - and with it the last-good state\n        // and all prior successful turns), stop processing further inputs and\n        // let the invocation resolve with `finishReason: 'failed'`. The caller\n        // recovers the last-good state from the returned AgentOutput.\n        break;\n      }\n    }\n  }\n\n  /**\n   * Saves a snapshot of the current session state to the persistent store.\n   *\n   * When a store is configured every turn is persisted (snapshotting is no\n   * longer opt-out). Uses the mutator-based `saveSnapshot` to atomically check\n   * that the snapshot has not been concurrently aborted before writing -\n   * preventing a race where a \"done\" write could overwrite a concurrent\n   * \"aborted\" status.\n   */\n  async maybeSnapshot(\n    status?: 'pending' | 'completed' | 'failed',\n    error?: { status?: string; message: string; details?: any },\n    snapshotId?: string,\n    finishReason?: AgentFinishReason\n  ): Promise<string | undefined> {\n    if (\n      !this.store ||\n      (this.isDetached && snapshotId !== this.lastSnapshot?.snapshotId)\n    )\n      return this.lastSnapshot?.snapshotId;\n\n    const currentVersion = this.session.getVersion();\n    if (currentVersion === this.lastSnapshotVersion && !status) {\n      return this.lastSnapshot?.snapshotId;\n    }\n\n    const currentState = this.session.getState();\n\n    const snapshotInput: SessionSnapshotInput<State> = {\n      ...(snapshotId || this.newSnapshotId\n        ? { snapshotId: (snapshotId || this.newSnapshotId)! }\n        : {}),\n      // Stamp the session id onto every snapshot in the chain so callers can\n      // resolve a snapshot's session without reaching into its state.\n      sessionId: this.session.sessionId,\n      createdAt: new Date().toISOString(),\n      updatedAt: new Date().toISOString(),\n      state: currentState as SessionState<State>,\n      parentId: this.lastSnapshot?.snapshotId,\n      // Default to a resumable `completed` status. The only caller that omits a\n      // status is the post-invocation write (which fires when the handler\n      // mutates state after the last turn); persisting it as `completed` keeps\n      // it a valid resume target under the \"only `completed` is resumable\" rule.\n      status: status ?? 'completed',\n      // Stamp an initial heartbeat on a `pending` (detached, in-flight)\n      // snapshot. A background heartbeat loop refreshes it; if it goes stale the\n      // snapshot is reported as `expired` on read (the worker is presumed dead).\n      ...(status === 'pending' && { heartbeatAt: new Date().toISOString() }),\n      ...(finishReason && { finishReason }),\n      error,\n    };\n\n    const effectiveId = snapshotId || this.newSnapshotId;\n\n    // Use the mutator-based saveSnapshot to atomically check the current\n    // status before writing.  If the snapshot was concurrently aborted,\n    // the mutator returns null and the write is skipped.\n    const assignedId = await this.store.saveSnapshot(\n      effectiveId,\n      abortAwareMutator(snapshotInput),\n      { context: getContext() }\n    );\n    if (assignedId === null) {\n      // Snapshot was aborted concurrently; preserve the existing ID\n      // without overwriting.\n      return effectiveId;\n    }\n\n    this.lastSnapshot = { ...snapshotInput, snapshotId: assignedId };\n    this.lastSnapshotVersion = currentVersion;\n\n    return assignedId;\n  }\n}\n\n/**\n * Projects an agent's server-side data onto the view a client should see.\n *\n * Every member is optional; an omitted member passes the corresponding data\n * through unchanged. Use this to redact sensitive fields or reshape data\n * before it leaves the server - covering both data at rest and data in flight:\n *\n * - `state` reshapes/redacts session state at rest. Applied to\n *   `AgentOutput.state` (client-managed agents), to snapshots returned by\n *   `getSnapshotData`, and as the baseline for streamed `customPatch` diffs\n *   (so streamed custom-state deltas stay consistent with the transformed\n *   full state). Note: `state.artifacts` is part of session state, so artifact\n *   redaction at rest happens here too.\n * - `chunk` reshapes/redacts each stream chunk in flight (`modelChunk`,\n *   `artifact`, `customPatch`, `turnEnd`) - e.g. filtering \"internal\" tool\n *   request/response parts out of model chunks, or redacting streamed\n *   artifacts. Return `null`/`undefined` to drop the chunk entirely.\n *\n * When both `state` and `chunk` touch the same data (e.g. artifacts), keeping\n * the two projections consistent is the author's responsibility.\n */\nexport interface ClientTransform<S = unknown> {\n  /**\n   * Reshapes/redacts session state before it is exposed to the client (at\n   * rest: `AgentOutput.state`, snapshots, and the streamed `customPatch`\n   * baseline).\n   */\n  state?: (state: SessionState<S>) => SessionState;\n  /**\n   * Reshapes/redacts each stream chunk before it is sent to the client.\n   * Return `null`/`undefined` to drop the chunk entirely.\n   */\n  chunk?: (chunk: AgentStreamChunk) => AgentStreamChunk | null | undefined;\n}\n\n/**\n * Function handler definition for custom agent actions.\n */\nexport type AgentFn<State> = (\n  sess: SessionRunner<State>,\n  options: {\n    sendChunk: (chunk: AgentStreamChunk) => void;\n    abortSignal?: AbortSignal;\n    context?: ActionContext;\n  }\n) => Promise<AgentResult>;\n\n/**\n * Lookup input for the `getSnapshotData` action / method.\n *\n * Mirrors {@link GetSnapshotOptions}: provide exactly one of `snapshotId`\n * (an exact snapshot) or `sessionId` (the session's latest leaf snapshot).\n */\nexport const GetSnapshotDataInputSchema = z.object({\n  snapshotId: z.string().optional(),\n  sessionId: z.string().optional(),\n});\n\n/**\n * Lookup input for `getSnapshotData`.\n */\nexport interface GetSnapshotDataInput {\n  snapshotId?: string;\n  sessionId?: string;\n  context?: ActionContext;\n}\n\nexport type GetSnapshotDataAction<S = unknown> = Action<\n  typeof GetSnapshotDataInputSchema,\n  z.ZodType<SessionSnapshot<S>>\n>;\n\n/**\n * Represents a configured, registered Agent.\n *\n * An `Agent` exposes two surfaces:\n *\n * 1. The ergonomic, transport-agnostic {@link AgentAPI} (`chat`, `loadChat`,\n *    `getSnapshot`, `abort`) - the same surface returned by `remoteAgent` on\n *    the client, so server- and client-side code share one interface.\n * 2. The lower-level {@link BidiAction} surface (`run`, `streamBidi`, …) for\n *    advanced use and for serving over HTTP.\n */\nexport interface Agent<State = unknown>\n  extends BidiAction<\n      typeof AgentInputSchema,\n      typeof AgentOutputSchema,\n      typeof AgentStreamChunkSchema,\n      typeof AgentInitSchema\n    >,\n    AgentAPI<State> {\n  getSnapshotData(\n    opts: GetSnapshotDataInput\n  ): Promise<SessionSnapshot<State> | undefined>;\n\n  abort(\n    snapshotId: string,\n    options?: SessionStoreOptions\n  ): Promise<SessionSnapshot['status'] | undefined>;\n\n  readonly getSnapshotDataAction: GetSnapshotDataAction<State>;\n  readonly abortAgentAction: Action<\n    typeof AgentAbortRequestSchema,\n    typeof AgentAbortResponseSchema\n  >;\n}\n\n/**\n * Error thrown for agent init *API misuse* that should surface to the caller as\n * a real, thrown error (mapped to an HTTP status by the server handler) rather\n * than being absorbed into a graceful `finishReason: 'failed'` result.\n *\n * Covers calling an agent with an init that does not match its state-management\n * mode (e.g. sending `state` to a server-managed agent, or `snapshotId`/\n * `sessionId` to a client-managed one) and the snapshot/session ownership\n * guard. Other pre-turn failures (missing snapshot, non-resumable snapshot,\n * invalid custom state) remain graceful.\n */\nexport class AgentInitError extends GenkitError {}\n\n/**\n * Asserts that the init strategy matches the agent's state-management mode,\n * throwing an {@link AgentInitError} on a mismatch.\n *\n * Server-managed agents (with a store) resume via a `snapshotId` / `sessionId`;\n * client-managed agents (no store) supply the full `state` blob. This is API\n * misuse, so it propagates as a thrown error rather than a graceful failure.\n */\nfunction assertInitMatchesStateManagement(\n  config: { name: string; store?: SessionStore<unknown> },\n  init: AgentInit | undefined\n): void {\n  if ((init?.snapshotId || init?.sessionId) && !config.store) {\n    throw new AgentInitError({\n      status: 'FAILED_PRECONDITION',\n      message:\n        `Cannot use '${init.snapshotId ? 'snapshotId' : 'sessionId'}' with ` +\n        `agent '${config.name}': this agent has no store configured ` +\n        `(client-managed state). Send 'state' instead.`,\n    });\n  }\n  if (init?.state && config.store) {\n    throw new AgentInitError({\n      status: 'FAILED_PRECONDITION',\n      message:\n        `Cannot send 'state' to agent '${config.name}': this agent uses ` +\n        `a server-managed store. Send 'snapshotId' or 'sessionId' instead.`,\n    });\n  }\n}\n\n/**\n * Resolves the {@link Session} (and originating snapshot, if any) for an agent\n * turn from its {@link AgentInit}.\n *\n * Server-managed agents (with a store) resume via a `snapshotId` (an exact\n * snapshot) or a `sessionId` (the session's latest snapshot); client-managed\n * agents (no store) supply the full `state` blob. Throws a {@link GenkitError}\n * on a missing snapshot, non-resumable snapshot, or invalid custom state - the\n * caller is expected to translate that into a graceful `finishReason: 'failed'`\n * result. The state-management mismatch checks are performed up front by\n * {@link assertInitMatchesStateManagement} and throw {@link AgentInitError}.\n */\nasync function resolveSession<State>(\n  config: { name: string; store?: SessionStore<State> },\n  store: SessionStore<State>,\n  init: AgentInit | undefined,\n  validateCustomState: (custom: unknown) => void\n): Promise<{ session: Session<State>; snapshot?: SessionSnapshot<State> }> {\n  if (init?.snapshotId) {\n    const snapshot = await store.getSnapshot({\n      snapshotId: init.snapshotId,\n      context: getContext(),\n    });\n    if (!snapshot) {\n      throw new GenkitError({\n        status: 'NOT_FOUND',\n        message: `Snapshot ${init.snapshotId} not found`,\n      });\n    }\n    // When both `snapshotId` and `sessionId` are supplied, `snapshotId` selects\n    // the exact snapshot to resume and `sessionId` acts as an ownership guard:\n    // the snapshot must belong to that session. A mismatch is API misuse, so it\n    // propagates as a thrown error (AgentInitError) rather than being absorbed\n    // into a graceful failure.\n    // Prefer the snapshot's top-level `sessionId`; fall back to the id carried\n    // in its state for rows written before snapshot-level ids existed.\n    const snapshotSessionId = snapshot.sessionId ?? snapshot.state?.sessionId;\n    if (init.sessionId && snapshotSessionId !== init.sessionId) {\n      throw new AgentInitError({\n        status: 'INVALID_ARGUMENT',\n        message:\n          `Snapshot ${init.snapshotId} does not belong to session ` +\n          `${init.sessionId} (it belongs to ` +\n          `${snapshotSessionId ?? 'an unknown session'}).`,\n      });\n    }\n\n    // Only `completed` snapshots are resumable. A failed/aborted/pending\n    // snapshot is persisted for inspection but is not a valid resume target.\n    if (snapshot.status !== 'completed') {\n      throw new GenkitError({\n        status: 'INVALID_ARGUMENT',\n        message:\n          `Snapshot ${init.snapshotId} is not resumable (status: ` +\n          `${snapshot.status ?? 'unknown'}). Only 'completed' snapshots can ` +\n          `be resumed.`,\n      });\n    }\n\n    validateCustomState(snapshot.state?.custom);\n    return {\n      snapshot,\n      session: new Session<State>(snapshot.state as SessionState<State>),\n    };\n  }\n\n  if (init?.sessionId) {\n    // Resume the session's latest snapshot. The store returns the latest leaf\n    // regardless of status, but only `completed` snapshots are resumable - so\n    // if the leaf is a non-resumable turn (e.g. a `failed`/`aborted`/`pending`\n    // turn) walk back over its parent chain to the last-good (`completed`)\n    // snapshot. When the session has no resumable snapshot (e.g. a first-turn\n    // failure) seed a fresh session bound to the requested sessionId so\n    // subsequent turns can find it.\n    let snapshot = await store.getSnapshot({\n      sessionId: init.sessionId,\n      context: getContext(),\n    });\n    // Walk back over non-resumable leaves to the last-good (`completed`)\n    // snapshot. Guard against a self-referential or cyclic `parentId` chain\n    // (corrupt history) with a visited set so we fail fast with\n    // `FAILED_PRECONDITION` instead of looping forever on store reads.\n    const visited = new Set<string>();\n    while (snapshot && snapshot.status !== 'completed') {\n      if (visited.has(snapshot.snapshotId)) {\n        throw new GenkitError({\n          status: 'FAILED_PRECONDITION',\n          message:\n            `Session '${init.sessionId}' has a cyclic snapshot parent chain ` +\n            `(snapshot '${snapshot.snapshotId}' was visited twice). Resume by ` +\n            `snapshotId instead.`,\n        });\n      }\n      visited.add(snapshot.snapshotId);\n      snapshot = snapshot.parentId\n        ? await store.getSnapshot({\n            snapshotId: snapshot.parentId,\n            context: getContext(),\n          })\n        : undefined;\n    }\n    if (snapshot) {\n      validateCustomState(snapshot.state?.custom);\n      return {\n        snapshot,\n        session: new Session<State>(snapshot.state as SessionState<State>),\n      };\n    }\n    return {\n      session: new Session<State>({\n        custom: undefined,\n        artifacts: [],\n        messages: [],\n        sessionId: init.sessionId,\n      }),\n    };\n  }\n\n  if (init?.state && !config.store) {\n    validateCustomState(init.state.custom);\n    return {\n      session: new Session<State>(init.state as SessionState<State>),\n    };\n  }\n\n  return {\n    session: new Session<State>({\n      custom: undefined,\n      artifacts: [],\n      messages: [],\n    }),\n  };\n}\n\n/**\n * Pumps the action's raw input stream into the runner's input channel while\n * intercepting `detach: true` directives.\n *\n * Running this proxy concurrently lets a detach directive take effect\n * immediately rather than waiting for the runner to drain a backlog of\n * pre-queued inputs. A detach-only message (no payload) is consumed here and\n * not forwarded, since it has no turn to process.\n */\nfunction pipeInputWithDetach<State>(\n  inputStream: AsyncIterable<AgentInput>,\n  target: Channel<AgentInput>,\n  getRunner: () => SessionRunner<State>,\n  storeEnabled: boolean,\n  rejectDetach: (reason: any) => void\n): void {\n  (async () => {\n    try {\n      for await (const input of inputStream) {\n        if (input.detach) {\n          if (!storeEnabled) {\n            rejectDetach(\n              new GenkitError({\n                status: 'FAILED_PRECONDITION',\n                message:\n                  'Detach is only supported when a session store is provided.',\n              })\n            );\n          } else {\n            const runner = getRunner();\n            // Reserve the in-flight snapshot's id up front so the detached\n            // snapshot and any handler-named external resources share one id.\n            const turnSnapshotId = runner.newSnapshotId || reserveSnapshotId();\n            runner.newSnapshotId = turnSnapshotId;\n            await runner.maybeSnapshot('pending', undefined, turnSnapshotId);\n            runner.isDetached = true;\n\n            if (runner.onDetach) {\n              runner.onDetach(turnSnapshotId);\n            }\n          }\n          // Only forward to the runner if the input carries a payload beyond\n          // the detach directive; a detach-only message has no turn to process.\n          const hasPayload = !!(\n            input.message ||\n            input.resume?.restart?.length ||\n            input.resume?.respond?.length\n          );\n          if (hasPayload) {\n            target.send(input);\n          }\n        } else {\n          target.send(input);\n        }\n      }\n      target.close();\n    } catch (e) {\n      target.error(e);\n    }\n  })();\n}\n\n/**\n * Registers a multi-turn custom agent action capable of maintaining persistent state.\n *\n * When `stateSchema` is provided the custom state is validated at load time\n * (from a snapshot store or from the client-supplied `init.state`) and the\n * JSON Schema representation is included in the action metadata so that\n * tooling (e.g. the Dev UI) can inspect / validate the state shape.\n */\nexport function defineCustomAgent<State = unknown>(\n  registry: Registry,\n  config: {\n    name: string;\n    description?: string;\n    stateSchema?: z.ZodType<State>;\n    store?: SessionStore<State>;\n    clientTransform?: ClientTransform<State>;\n  },\n  fn: AgentFn<State>\n): Agent<State> {\n  // Helper that applies the optional state transform before exposing state to\n  // the client.  When no transform is configured it returns the raw state.\n  const toClientState = (\n    state: SessionState<State>\n  ): SessionState | undefined => {\n    if (config.clientTransform?.state) {\n      return config.clientTransform.state(state);\n    }\n    return state as SessionState;\n  };\n\n  // If a state schema was provided, pre-compute the JSON schema once so it\n  // can be embedded in metadata and reused for validation.\n\n  const stateJsonSchema = config.stateSchema\n    ? toJsonSchema({ schema: config.stateSchema })\n    : undefined;\n\n  /**\n   * Validates the `custom` field of a session state against the configured\n   * `stateSchema`.  No-ops when no schema was provided.\n   */\n  const validateCustomState = (custom: unknown): void => {\n    if (config.stateSchema && custom !== undefined) {\n      parseSchema(custom, { schema: config.stateSchema });\n    }\n  };\n\n  const primaryAction = defineBidiAction(\n    registry,\n    {\n      name: config.name,\n      description: config.description,\n      actionType: 'agent',\n      inputSchema: AgentInputSchema,\n      outputSchema: AgentOutputSchema,\n      streamSchema: AgentStreamChunkSchema,\n      initSchema: AgentInitSchema,\n      metadata: {\n        agent: {\n          stateManagement: config.store ? 'server' : 'client',\n          abortable: !!config.store?.onSnapshotStateChange,\n          ...(stateJsonSchema && { stateSchema: stateJsonSchema }),\n        },\n      },\n    },\n    async function* (\n      arg: ActionFnArg<AgentStreamChunk, AgentInput, AgentInit>\n    ) {\n      const init = arg.init;\n      const store = config.store || new InMemorySessionStore<State>();\n\n      // API-misuse checks (init does not match the agent's state-management\n      // mode) throw out of the generator so the server handler maps them to a\n      // proper HTTP status, rather than being absorbed into a graceful\n      // `finishReason: 'failed'` result below.\n      assertInitMatchesStateManagement(config, init);\n\n      let session!: Session<State>;\n      let snapshot: SessionSnapshot<State> | undefined;\n\n      try {\n        ({ session, snapshot } = await resolveSession<State>(\n          config,\n          store,\n          init,\n          validateCustomState\n        ));\n      } catch (e: any) {\n        // An AgentInitError signals API misuse (e.g. the snapshot/session\n        // ownership guard) that must surface as a thrown error; re-throw it so\n        // the server handler maps it to a proper HTTP status.\n        if (e instanceof AgentInitError) {\n          throw e;\n        }\n        // Other pre-turn / setup failures (missing snapshot, non-resumable\n        // snapshot, invalid client state). Resolve gracefully with\n        // `finishReason: 'failed'` - preserving the original `error.status` -\n        // rather than throwing, so the caller gets a structured, inspectable\n        // result. There is no last-good turn yet; echo back the\n        // client-supplied state when present.\n        return {\n          finishReason: 'failed' as AgentFinishReason,\n          error: toErrorDetails(e),\n          ...(!config.store &&\n            init?.state && { state: init.state as SessionState }),\n        };\n      }\n\n      // Tag the current trace span with the sessionId so that traces\n      // belonging to the same agent conversation can be correlated.\n      setCustomMetadataAttributes({\n        'agent:sessionId': session.sessionId,\n      });\n\n      let detachedSnapshotId: string | undefined;\n      let resolveDetach:\n        | ((value: void | PromiseLike<void>) => void)\n        | undefined;\n      let rejectDetach: ((reason: any) => void) | undefined;\n      const detachPromise = new Promise<void>((resolve, reject) => {\n        resolveDetach = resolve;\n        rejectDetach = reject;\n      });\n\n      const abortController = new AbortController();\n      let unsubscribe: any = undefined;\n      // Background heartbeat timer for the detached snapshot. Started in\n      // `onDetach`, cleared when the flow settles (or on abort).\n      let heartbeatTimer: ReturnType<typeof setInterval> | undefined;\n      const stopHeartbeat = () => {\n        if (heartbeatTimer) {\n          clearInterval(heartbeatTimer);\n          heartbeatTimer = undefined;\n        }\n      };\n\n      let runner!: SessionRunner<State>;\n\n      // Centralized chunk emitter: every stream chunk passes through here so\n      // the optional `clientTransform.chunk` can reshape/redact it (or drop it\n      // by returning a nullish value) before it reaches the client.\n      //\n      // The actual dispatch is failure-isolated (like `notifyEndTurn`): the\n      // artifact/customPatch emitters fire synchronously from inside\n      // `Session.updateCustom`/`addArtifacts` (i.e. from the user's handler).\n      // If the client stream is already closed, `sendChunk` throws; absorbing\n      // it here prevents that from propagating out of the handler and turning\n      // a normal turn into a `failed` one.\n      const emitChunk = (chunk: AgentStreamChunk) => {\n        try {\n          let toSend: AgentStreamChunk | null | undefined = chunk;\n          if (config.clientTransform?.chunk) {\n            toSend = config.clientTransform.chunk(chunk);\n          }\n          if (!toSend) return;\n          arg.sendChunk(toSend);\n        } catch {\n          // Stream was closed (or the transform threw); absorb the exception.\n        }\n      };\n\n      // We construct an asynchronous proxy channel over the inputStream.\n      // This enables immediate interception of `detach: true` directives. Without this proxy,\n      // a backlog of pre-queued inputs would have to be resolved sequentially by the runner first.\n      const runnerInputChannel = new Channel<AgentInput>();\n\n      pipeInputWithDetach(\n        arg.inputStream,\n        runnerInputChannel,\n        () => runner,\n        !!config.store,\n        (reason) => rejectDetach?.(reason)\n      );\n\n      runner = new SessionRunner<State>(session, runnerInputChannel, {\n        store,\n        lastSnapshot: snapshot,\n        abortSignal: abortController.signal,\n\n        onDetach: (snapshotId) => {\n          detachedSnapshotId = snapshotId;\n          if (resolveDetach) {\n            resolveDetach();\n          }\n\n          // Refresh the detached snapshot's heartbeat periodically. The mutator\n          // only touches a still-`pending` snapshot (returns null otherwise) so\n          // it never resurrects a terminal snapshot or clobbers a concurrent\n          // abort. If a read sees this heartbeat go stale, the snapshot is\n          // reported as `expired` (the worker is presumed dead). `unref` so the\n          // timer never keeps the process alive on its own.\n          const ctx = getContext();\n          heartbeatTimer = setInterval(() => {\n            void store\n              .saveSnapshot(\n                snapshotId,\n                (current) =>\n                  current?.status === 'pending'\n                    ? { ...current, heartbeatAt: new Date().toISOString() }\n                    : null,\n                { context: ctx }\n              )\n              .catch(() => {\n                // Best-effort heartbeat; ignore transient store errors.\n              });\n          }, DEFAULT_HEARTBEAT_INTERVAL_MS);\n          heartbeatTimer.unref?.();\n\n          if (store.onSnapshotStateChange) {\n            unsubscribe = store.onSnapshotStateChange(\n              snapshotId,\n              (snap) => {\n                if (snap.status === 'aborted') {\n                  stopHeartbeat();\n                  abortController.abort();\n                  if (unsubscribe) unsubscribe();\n                }\n              },\n              { context: getContext() }\n            );\n          }\n        },\n\n        onEndTurn: (snapshotId, finishReason) => {\n          if (!runner.isDetached) {\n            emitChunk({\n              turnEnd: {\n                ...(config.store && { snapshotId }),\n                ...(finishReason && { finishReason }),\n              },\n            });\n          }\n        },\n      });\n\n      const sendArtifactChunk = (a: Artifact) => {\n        if (!runner.isDetached) {\n          emitChunk({ artifact: a });\n        }\n      };\n\n      session.on('artifactAdded', sendArtifactChunk);\n      session.on('artifactUpdated', sendArtifactChunk);\n\n      // Auto-emit a `customPatch` chunk whenever custom state is mutated.\n      // The diff is computed AFTER the clientStateTransform so streamed deltas\n      // honor redaction and stay consistent with the transformed full state in\n      // snapshots / final output. The first patch of every turn is a\n      // whole-document replace (re-basing clients that may lack the baseline);\n      // subsequent patches are incremental diffs against the last sent value.\n      let lastSentCustom: unknown;\n      const sendCustomPatch = () => {\n        if (runner.isDetached) return;\n        const transformed = toClientState(session.getState())?.custom;\n        let patch: JsonPatch;\n        if (runner.firstCustomPatchInTurn) {\n          patch = [\n            { op: 'replace', path: '', value: structuredClone(transformed) },\n          ];\n          runner.firstCustomPatchInTurn = false;\n        } else {\n          patch = diff(lastSentCustom, transformed);\n        }\n        lastSentCustom = structuredClone(transformed);\n        if (patch.length) {\n          emitChunk({ customPatch: patch });\n        }\n      };\n      session.on('customChanged', sendCustomPatch);\n\n      const sendChunk = (chunk: AgentStreamChunk) => {\n        if (!runner.isDetached) {\n          emitChunk(chunk);\n        }\n      };\n\n      const flowPromise = (async () => {\n        try {\n          const result = await runWithSession(registry, session, () =>\n            fn(runner, {\n              sendChunk,\n              abortSignal: abortController.signal,\n              context: getContext(),\n            })\n          );\n          // After the handler resolves, persist any state it mutated after the\n          // last turn. Omitting a status defaults to a resumable `completed`\n          // write, which the version guard skips when nothing changed.\n          const finalSnapshotId = await runner.maybeSnapshot();\n          return { result, finalSnapshotId };\n        } finally {\n          // The turn has settled (the snapshot reached a terminal status), so\n          // stop refreshing its heartbeat.\n          stopHeartbeat();\n          if (unsubscribe) unsubscribe();\n          session.off('artifactAdded', sendArtifactChunk);\n          session.off('artifactUpdated', sendArtifactChunk);\n          session.off('customChanged', sendCustomPatch);\n        }\n      })();\n\n      // We race the background flow execution against the detach signal.\n      // If detachment is requested, we yield output metadata early, but allow\n      // the flow handler promise to continue its asynchronous completion.\n      const outcome = await Promise.race([\n        flowPromise,\n        detachPromise.then(() => 'detached' as const),\n      ]);\n\n      if (outcome === 'detached') {\n        return {\n          sessionId: session.sessionId,\n          snapshotId: detachedSnapshotId!,\n          finishReason: 'detached' as AgentFinishReason,\n          ...(!config.store && { state: toClientState(session.getState()) }),\n        };\n      }\n\n      const { result, finalSnapshotId } = outcome;\n\n      // A turn failed: resolve gracefully with `finishReason: 'failed'` and the\n      // last-good state (what the failed turn started with), rather than the\n      // live state which may hold the failed turn's partial mutations.\n      if (runner.lastTurnFinishReason === 'failed' && runner.lastTurnError) {\n        const lastGood = (runner.lastGoodState ??\n          session.getState()) as SessionState<State>;\n        const lastGoodMessages = lastGood.messages;\n        return {\n          sessionId: session.sessionId,\n          finishReason: 'failed' as AgentFinishReason,\n          error: runner.lastTurnError,\n\n          ...(result.artifacts?.length && { artifacts: result.artifacts }),\n          ...(lastGoodMessages?.length && {\n            message: lastGoodMessages[lastGoodMessages.length - 1],\n          }),\n          // Server-managed: the last successful turn is already persisted (every\n          // turn is snapshotted), so point at its snapshot. The failed turn's\n          // own snapshot is persisted too but is not resumable - `sessionId`\n          // resume skips it back to this last-good `done` snapshot. Undefined on\n          // a first-turn failure (no successful turn yet; client holds the seed).\n          ...(config.store && { snapshotId: runner.lastGoodSnapshotId }),\n          // Client-managed: return the last-good state directly.\n          ...(!config.store && { state: toClientState(lastGood) }),\n        };\n      }\n\n      const finishReason = result.finishReason ?? runner.lastTurnFinishReason;\n\n      return {\n        sessionId: session.sessionId,\n        ...(result.artifacts?.length && { artifacts: result.artifacts }),\n        ...(result.message && { message: result.message }),\n        ...(finishReason && { finishReason }),\n        ...(config.store && { snapshotId: finalSnapshotId }),\n        ...(!config.store && { state: toClientState(session.getState()) }),\n      };\n    }\n  );\n\n  // Helper that applies the clientTransform.state projection to a snapshot's\n  // state, returning a new snapshot object with the transformed state.\n  const toClientSnapshot = (\n    snapshot: SessionSnapshot<State>\n  ): SessionSnapshot => {\n    if (!config.clientTransform?.state || !snapshot.state) {\n      return snapshot as SessionSnapshot;\n    }\n    return {\n      ...snapshot,\n      state: config.clientTransform.state(snapshot.state),\n    };\n  };\n\n  // Shared snapshot/abort implementations, reused by both the `defineAction`\n  // surfaces (which inject the ambient request context) and the ergonomic\n  // composite methods (which accept caller-supplied options).\n  const resolveSnapshot = async (\n    lookup: GetSnapshotDataInput\n  ): Promise<SessionSnapshot | undefined> => {\n    requireStore(config.store, 'getSnapshotData', config.name);\n    const snapshot = await config.store.getSnapshot(lookup);\n    if (!snapshot) return undefined;\n    // Compute `expired` on read: a `pending` snapshot whose heartbeat has gone\n    // stale is presumed orphaned (its background worker died), so surface it as\n    // `expi