From f976890517f41c880e46151c5d6adea0c14692be Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 13 Jul 2026 13:21:31 +0200 Subject: [PATCH 1/6] docs: design full AG-UI interrupt support --- .../2026-07-13-ag-ui-interrupts-design.md | 1036 +++++++++++++++++ 1 file changed, 1036 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-13-ag-ui-interrupts-design.md diff --git a/docs/superpowers/specs/2026-07-13-ag-ui-interrupts-design.md b/docs/superpowers/specs/2026-07-13-ag-ui-interrupts-design.md new file mode 100644 index 000000000..37e17fc78 --- /dev/null +++ b/docs/superpowers/specs/2026-07-13-ag-ui-interrupts-design.md @@ -0,0 +1,1036 @@ +# AG-UI Interrupts: Full-Spec, Typed, Durable Resolution + +Status: Approved design + +Date: 2026-07-13 + +Implementation base: [TanStack AI PR #785](https://github.com/TanStack/ai/pull/785), commit `61026d52669fc261b7a051ba20ba4bd806ce84fd` + +## Summary + +TanStack AI will adopt AG-UI native interrupts as the only interrupt protocol emitted by new servers. The client API will expose bound interrupt objects whose methods stage typed responses, validate them, and submit a complete response batch only after every simultaneously open interrupt has been answered or cancelled. + +The selected design is staged bound interrupts with auto-submit: + +- A single interrupt auto-submits after `resolveInterrupt` or `cancelInterrupt`. +- A batch waits until every open interrupt has a staged resolution or cancellation, then submits exactly once. +- `resolveInterrupts(callback)` provides a synchronous root transaction for heterogeneous batches. +- Tool approvals infer edited arguments from `inputSchema` and custom approval data from `approvalSchema`. +- Generic AG-UI interrupts retain `unknown` payloads statically and validate them from `responseSchema` at runtime. +- The server validates the entire batch, reports all validation errors together, and executes no tools when any entry is invalid. +- Final response batches use a mandatory atomic compare-and-swap persistence operation. Exact replays are idempotent, conflicting responses are rejected, and the first valid writer wins. + +This is a breaking protocol and API transition relative to the pre-PR approval flow and parts of PR #785's current surface. Compatibility is handled by a pre-1.0 deprecation period, retained legacy readers in new clients, coordinated client and server upgrades, and a manual migration guide. There will be no codemod. + +## Normative basis and requirement labels + +This design separates protocol requirements from TanStack guarantees: + +| Label | Meaning | +| --- | --- | +| `[AG-UI]` | Required for AG-UI protocol conformance. The normative source is the interrupts document at immutable commit [`2ed25b2009ff900315ab686dc9f86147b48fe382`](https://github.com/ag-ui-protocol/ag-ui/blob/2ed25b2009ff900315ab686dc9f86147b48fe382/docs/concepts/interrupts.mdx). The PR-head type baseline is `@ag-ui/core@0.0.57`. | +| `[TanStack]` | A stronger TanStack API, validation, safety, persistence, or durability guarantee. It is not claimed as an AG-UI protocol requirement. | +| `[Compatibility]` | Temporary pre-1.0 behavior for old TanStack clients or servers. It is outside the native AG-UI interrupt protocol. | + +The AG-UI requirements used here are: an interrupted run terminates with a nonempty interrupt outcome; required message and state snapshots precede that terminal event; the continuation uses the same thread and a new run; one resume array covers every open interrupt; pending interrupts block unrelated input; `responseSchema` describes resolved payloads; exact replay is idempotent; expiry rejects stale resumes; denial is a resolved payload; cancellation has no payload; and a resumed tool emits only its result against the original `toolCallId`. + +TanStack additionally guarantees typed tool shorthands, client preflight validation, authoritative server validation, exhaustive batch errors, full-set atomic compare-and-swap persistence, first-writer conflict handling, local draft durability, recovery state, and cross-framework API parity. In particular, AG-UI describes validation against `responseSchema`, but this design's mandatory server-side validation of every entry, all-errors response, and no-write-on-any-invalid rule are explicit TanStack guarantees. + +## Evidence and breaking-change analysis + +The protocol requirements in this design follow the primary [AG-UI interrupts specification](https://docs.ag-ui.com/concepts/interrupts) and its [immutable source revision](https://github.com/ag-ui-protocol/ag-ui/blob/2ed25b2009ff900315ab686dc9f86147b48fe382/docs/concepts/interrupts.mdx). The installed `@ag-ui/core@0.0.57` types define an interrupt descriptor, `resolved` and `cancelled` resume statuses, a `resume` array on run input, and an interrupt outcome on `RUN_FINISHED`. + +Before PR #785, TanStack AI used approval-specific custom events and `addToolApprovalResponse`. The old server emitted `approval-requested` and `tool-input-available` custom events, as shown in the [pre-PR chat implementation](https://github.com/TanStack/ai/blob/5fcaf90dc82bc20b8c7a75faa3c129da04858af5/packages/ai/src/activities/chat/index.ts#L1719-L1772). That flow represented approval state in message parts and resumed through updated conversation history. + +PR #785 changes the lifecycle to native interrupt outcomes and resumable run state. That is a breaking protocol change for an old client connected to a new server: the old client depends on custom approval events that the new server will no longer emit. Emitting both paths is not safe because duplicate prompts, duplicate callbacks, and ordering ambiguity can cause two responses to the same logical decision. No capability negotiation exists to make dual emission reliable. + +PR #785 also has correctness gaps that this design closes: + +- [`addToolApprovalResponse`](https://github.com/TanStack/ai/blob/61026d52669fc261b7a051ba20ba4bd806ce84fd/packages/ai-client/src/chat-client.ts#L1514-L1534) currently maps denial to `cancelled`. AG-UI represents denial as a resolved response such as `{ "approved": false }`; cancellation abandons an interrupt and has no payload. +- The [resume request path](https://github.com/TanStack/ai/blob/61026d52669fc261b7a051ba20ba4bd806ce84fd/packages/ai-client/src/chat-client.ts#L1054-L1067) reuses the interrupted run ID. AG-UI ends the interrupted run and resumes in a new run on the same thread. +- The [interrupt construction path](https://github.com/TanStack/ai/blob/61026d52669fc261b7a051ba20ba4bd806ce84fd/packages/ai/src/activities/chat/index.ts#L1762-L1825) uses noncanonical reasons. Tool approvals must use `tool_call`; a TanStack-specific client-tool interruption must use a namespaced reason. +- The [resumed tool path](https://github.com/TanStack/ai/blob/61026d52669fc261b7a051ba20ba4bd806ce84fd/packages/ai/src/activities/chat/index.ts#L1837-L1873) can replay tool call start, arguments, and end events. A resumed edited tool call emits only the result against the original `toolCallId`. +- The [persistence middleware](https://github.com/TanStack/ai/blob/61026d52669fc261b7a051ba20ba4bd806ce84fd/packages/ai-persistence/src/middleware.ts#L93-L154) validates fail-fast and stores responses sequentially. A response batch needs exhaustive validation followed by one atomic commit. +- The [interrupt store contract](https://github.com/TanStack/ai/blob/61026d52669fc261b7a051ba20ba4bd806ce84fd/packages/ai-persistence/src/types.ts#L47-L73) exposes per-item `resolve` and `cancel`, so it cannot guarantee all-or-none batch persistence or first-writer conflict handling. + +Because PR #785 has not established a stable released API for native interrupt resumption, its raw `resumeInterrupts` surface can be replaced by the safe bound API in this work. An explicitly unsafe raw escape hatch remains available. + +## Goals + +- Fully implement the AG-UI interrupt lifecycle, event ordering, resume batching, cancellation, expiry, and idempotency requirements. +- Make tool approvals ergonomic and type-safe from definition through every framework hook. +- Support edited tool arguments as an optional full replacement of the original arguments. +- Support branch-specific typed custom approval and rejection payloads. +- Support arbitrary AG-UI interrupt reasons and response schemas without pretending runtime JSON Schema creates static TypeScript types. +- Make multiple simultaneous interrupts safe and easy to resolve individually or as a root transaction. +- Validate every entry on the server and return all errors in one response. +- Preserve client response drafts across reloads while keeping final response authority on the server. +- Prevent partial response persistence, duplicate continuation scheduling, and last-writer overwrites. +- Give React, Preact, Solid, Vue, Svelte, and Angular equivalent behavior through the headless client. +- Provide a clear pre-1.0 deprecation and migration path. + +## Non-goals + +- A codemod for approval APIs. +- Static type inference from an interrupt's runtime JSON Schema. +- Protocol negotiation or dual emission of native and legacy approval events. +- Partial submission of a batch while sibling interrupts remain unanswered. +- Merging `editedArgs` into original arguments. Edited arguments are always a complete replacement. +- Deprecating `addToolResult`; it remains the direct client-tool result API. +- Combining client-tool execution and human tool approval into one interrupt type. +- Exactly-once external side effects for arbitrary tools. Tool implementations still need idempotency where crash recovery can re-enter external work. +- Changing general checkpoint, tool-result, or message persistence beyond the interrupt response transaction required here. + +## Considered architectures + +### 1. Staged bound interrupts with auto-submit, selected + +Every exposed interrupt carries methods bound to the owning chat client. Item methods stage a decision. Once all currently open interrupts are covered, the client submits one batch. A singleton is therefore both staged and submitted by one call. + +This design makes the common case concise, supports per-card user interfaces, and preserves the AG-UI all-open-interrupt batch invariant. It also provides a natural place for typed tool overloads, validation status, item errors, replacement, and clearing. + +### 2. Pure response builders plus explicit submit + +Item methods could return response objects and require a separate root submit call. This makes batching explicit but adds boilerplate to the singleton case, makes per-card code manage a second response collection, and does not match the requested binding behavior. + +### 3. Root-only reducer + +The hook could expose only `resolveInterrupts(interrupts => responses)`. This centralizes atomicity but makes independent interrupt components awkward, prevents direct item binding, and forces a heterogeneous response construction problem into every consumer. + +The selected architecture keeps atomic submission centralized internally while exposing both convenient item bindings and an explicit root transaction. + +## Public type model + +The examples below describe the public contract. Implementation helper types may use different internal names, but they must preserve these inference and assignability rules. + +### Common bound interrupt fields + +```ts +type InterruptStatus = + | 'pending' + | 'validating' + | 'staged' + | 'submitting' + | 'error' + +interface InterruptCorrelation { + threadId: string + interruptedRunId: string + generation: number + submissionId?: string + continuationRunId?: string +} + +type ItemInterruptErrorCode = + | 'invalid-payload' + | 'invalid-edited-args' + | 'invalid-tool-output' + | 'invalid-response-schema' + | 'unknown-interrupt' + | 'expired' + | 'stale' + | 'conflict' + | 'legacy-unsupported' + +interface ItemInterruptError extends InterruptCorrelation { + scope: 'item' + interruptId: string + code: ItemInterruptErrorCode + message: string + path?: readonly (string | number)[] + source: 'client' | 'server' + retryable: boolean +} + +type BatchInterruptErrorCode = + | 'incomplete-batch' + | 'item-validation-failed' + | 'unsupported-bulk-operation' + | 'async-resolver' + | 'inactive-transaction' + | 'mixed-provenance' + | 'transport' + | 'server' + | 'protocol' + | 'invalid-response-schema' + | 'expired' + | 'stale' + | 'conflict' + | 'persistence-required' + | 'atomic-commit-unsupported' + | 'recovery-unavailable' + | 'legacy-submit-failed' + +interface BatchInterruptError extends InterruptCorrelation { + scope: 'batch' + code: BatchInterruptErrorCode + message: string + source: 'client' | 'server' | 'transport' + retryable: boolean + interruptIds: readonly string[] +} + +type InterruptSubmissionError = ItemInterruptError | BatchInterruptError + +interface BoundInterruptBase { + id: string + reason: string + message?: string + toolCallId?: string + responseSchema?: JSONSchema + expiresAt?: string + metadata?: Record + status: InterruptStatus + provenance: 'native' | 'legacy' + generation: number + errors: readonly ItemInterruptError[] + stagedResponse?: TStagedResponse + clearResolution(): void + cancelInterrupt(): void +} +``` + +The AG-UI fields retain the installed core package's actual exported types. A typed `JSONSchema` descriptor still does not create a static application payload type. `stagedResponse` is a read-only view of the local draft and never includes bound functions when persisted. + +Every public variant has a TanStack-owned `kind` discriminator. Known tool variants also carry the configured tool's literal `toolName`; neither field is read from untrusted resume metadata. + +```ts +declare const toolApprovalCapability: unique symbol + +interface ToolApprovalCapabilityMarker { + readonly [toolApprovalCapability]: TNeedsApproval +} + +type ApprovalCapabilityOf = + TTool extends ToolApprovalCapabilityMarker + ? TNeedsApproval + : false + +type ToolName = TTools[number]['name'] + +type ToolByName< + TTools extends readonly AnyClientTool[], + TName extends ToolName, +> = Extract + +type ApprovalToolName = { + [TName in ToolName]: ApprovalCapabilityOf< + ToolByName + > extends true + ? TName + : never +}[ToolName] + +type ToolApprovalInterruptsFor< + TTools extends readonly AnyClientTool[], +> = { + [TName in ApprovalToolName]: ToolApprovalInterrupt< + ToolByName + > & { + kind: 'tool-approval' + toolName: TName + } +}[ApprovalToolName] + +type ClientToolExecutionInterruptsFor< + TTools extends readonly AnyClientTool[], +> = { + [TName in ToolName]: ClientToolExecutionInterrupt< + ToolByName + > & { + kind: 'client-tool-execution' + toolName: TName + } +}[ToolName] + +type ChatInterrupt = + | ToolApprovalInterruptsFor + | ClientToolExecutionInterruptsFor + | GenericAGUIInterrupt +``` + +A generic interrupt has `kind: 'generic'` and no `toolName`. A known tool interrupt narrows by `kind` and literal `toolName` to the configured tool's input, output, and approval payload types. `ToolApprovalInterruptsFor` contains only names whose preserved `TNeedsApproval` capability is the literal `true`; tools with literal `false`, absent approval capability, or a widened nonliteral capability cannot enter the approval variant. A native descriptor that cannot be matched safely remains generic even if it happens to contain a tool-like reason. + +At submission, the server looks up the authoritative original call by `(threadId, interruptedRunId, interruptId)` and obtains its stored `toolCallId`, `toolName`, original arguments, schema hashes, and generation. It never selects a tool or arguments from client-supplied `kind`, `toolName`, `toolCallId`, or metadata. Any conflicting client correlation field is ignored for binding and reported as `stale` or `protocol` when it indicates a mismatched descriptor. + +### Tool definition configuration + +`approvalSchema` is legal only when `needsApproval` is statically `true`. + +```ts +const transferFunds = toolDefinition({ + name: 'transferFunds', + inputSchema: transferInputSchema, + outputSchema: transferOutputSchema, + needsApproval: true, + approvalSchema: { + approve: approvalPayloadSchema, + reject: rejectionPayloadSchema, + }, +}) +``` + +Supported forms are: + +```ts +approvalSchema: sharedPayloadSchema + +approvalSchema: { + approve: approvalPayloadSchema, + reject: rejectionPayloadSchema, +} +``` + +The branch object must contain at least one branch. A shared schema applies to both approval and rejection. In a branch object, an omitted branch forbids a custom payload for that decision. If `approvalSchema` is absent, both branches forbid a custom payload. `approvalSchema` by itself never enables approval; `needsApproval: true` is required. + +Each shared or branch schema accepts the full existing `SchemaInput` family: Standard JSON Schema, Standard Schema, or TanStack's raw `JSONSchema`. Standard JSON Schema and raw JSON Schema are serialized into the AG-UI `responseSchema`. A Standard Schema validator that also exports Standard JSON Schema uses that export. A Standard Schema validator without JSON Schema export remains supported: matching TanStack clients and the server use the configured validator, while the wire schema expresses the branch envelope with an unconstrained nested payload. Generic clients therefore keep that nested payload as `unknown`, and server validation remains authoritative. + +Schema inference rules are explicit: + +- A Standard Schema or Standard JSON Schema input contributes its declared input type. +- A raw `JSONSchema` contributes `unknown`; TanStack does not infer a TypeScript type from JSON Schema syntax. +- An absent `inputSchema` forbids `editedArgs` in types and at runtime. The original stored arguments are used unchanged. +- A present raw `inputSchema` permits `editedArgs: unknown`, then validates the replacement at runtime. +- An absent `outputSchema` makes client-tool output `unknown` and performs no output-schema validation. +- A present raw `outputSchema` keeps output `unknown` and validates it at runtime. +- An absent approval branch schema forbids nested `payload`. +- A present raw approval schema requires `payload: unknown`; Standard Schema input that includes `undefined` makes the property optional, while other typed schema inputs require it. + +The branch-map form is recognized only when its own keys are a nonempty subset of `approve` and `reject` and its values are `SchemaInput` values. A Standard Schema object is recognized by its standard marker first, and an object containing JSON Schema keywords is a shared raw schema. Ambiguous or malformed configuration fails tool-definition construction. + +The tool-definition builder carries its `TNeedsApproval` generic through the non-runtime `ToolApprovalCapabilityMarker` and preserves that marker, the input schema, output schema, and approval schema generics through client and server tool transformations. The mapped union reads this type marker rather than the optional runtime `needsApproval` property. Configurations with `needsApproval: false` or without `needsApproval` reject `approvalSchema` at compile time and runtime. + +### Typed tool approval signatures + +A tool approval interrupt has this normative branch-aware public contract: + +```ts +interface ToolApprovalInterrupt< + TTool extends AnyClientTool, +> extends BoundInterruptBase, + ApproveSchema, + RejectSchema + >> { + kind: 'tool-approval' + reason: 'tool_call' + toolName: TTool['name'] + toolCallId: string + + resolveInterrupt( + approved: true, + ...options: OptionsTuple, + ApproveSchema + >> + ): void + + resolveInterrupt( + approved: false, + ...options: OptionsTuple>> + ): void +} + +type EditedArgsField = TInputSchema extends NoSchema + ? { editedArgs?: never } + : { editedArgs?: InferSchemaInput } + +type PayloadField = TPayloadSchema extends NoSchema + ? { payload?: never } + : TPayloadSchema extends JSONSchema + ? { payload: unknown } + : undefined extends InferSchemaInput + ? { payload?: InferSchemaInput } + : { payload: InferSchemaInput } + +type ApprovalFields = + EditedArgsField & PayloadField + +type RejectionFields = PayloadField & { + editedArgs?: never +} + +declare const noSchema: unique symbol +type NoSchema = typeof noSchema + +type RequiredKeys = { + [TKey in keyof T]-?: {} extends Pick ? never : TKey +}[keyof T] + +type OptionsTuple = [RequiredKeys] extends [never] + ? [options?: TFields] + : [options: TFields] + +type ToolApprovalDraft = + | ({ status: 'resolved'; approved: true } & ApprovalFields< + TInputSchema, + TApproveSchema + >) + | ({ status: 'resolved'; approved: false } & RejectionFields< + TRejectSchema + >) + | { + status: 'cancelled' + } +``` + +`OptionsTuple` makes the argument optional only when every property in `TFields` is optional; otherwise it requires exactly one options object. The resolver overloads and `ToolApprovalDraft` deliberately share `ApprovalFields` and `RejectionFields`, so `stagedResponse` cannot represent a payload that the resolver could not create. + +The conditional option tuple has these rules: + +- `editedArgs` is optional and exists only on the approval branch. +- `editedArgs` is inferred from the tool's `inputSchema` and is forbidden when that schema is absent. +- If an approval payload schema requires a value, the options object and its `payload` field are required. +- If an approval payload is optional, the options object is optional unless `editedArgs` is supplied. +- If a branch has no payload schema, that branch's options type has no `payload` field. Raw JSON Schema payloads are required but remain `unknown` statically. +- The rejection branch has no `editedArgs` field. +- Excess-property checks and runtime validation reject custom fields outside the nested `payload` object. + +The intended calls are: + +```ts +interrupt.resolveInterrupt(true) + +interrupt.resolveInterrupt(true, { + editedArgs: updatedInput, + payload: { note: 'Reviewed' }, +}) + +interrupt.resolveInterrupt(false, { + payload: { reason: 'Insufficient evidence' }, +}) +``` + +The first argument always represents approval or denial. On the AG-UI wire, both are `resolved` entries. TanStack's response schema and resume payload use this envelope: + +```ts +type ToolApprovalResumePayload = + | { + approved: true + editedArgs?: unknown + payload?: unknown + } + | { + approved: false + payload?: unknown + } +``` + +`approvalSchema` governs only the nested `payload`. The generated `responseSchema` is a JSON Schema union of the approval and rejection envelopes. The approval branch includes optional `editedArgs` governed by `inputSchema`; the rejection branch forbids `editedArgs`. Branches without an approval schema forbid `payload`. + +`editedArgs`, when present, fully replaces the original tool arguments. Neither client nor server performs a shallow or deep merge. The server validates the replacement against `inputSchema` before any response is committed or tool is executed. + +### Generic AG-UI interrupts + +An interrupt that is not a known tool approval or a known TanStack client-tool execution remains generic: + +```ts +interface GenericAGUIInterrupt + extends BoundInterruptBase< + | { status: 'resolved'; payload: unknown } + | { status: 'cancelled' } + > { + kind: 'generic' + toolName?: never + resolveInterrupt(payload: unknown): void +} +``` + +All standard descriptor fields are statically known. The response remains `unknown` because a JSON Schema received at runtime cannot safely produce a compile-time TypeScript type. The client validates against `responseSchema` as a preflight when a schema is present. The server repeats that validation authoritatively. + +The user documentation will show how to convert the received JSON Schema into a schema supported by the user's chosen validation library, validate the payload, and narrow the result before calling `resolveInterrupt`. It will also show a Standard Schema compatible validation path. Documentation examples must not claim static inference from the wire schema. + +### Client-tool execution interrupts + +A client-owned tool that is awaiting browser execution is separate from human approval: + +```ts +interface ClientToolExecutionInterrupt + extends BoundInterruptBase< + | { status: 'resolved'; payload: InferToolOutput } + | { status: 'cancelled' } + > { + kind: 'client-tool-execution' + reason: 'tanstack:client_tool_execution' + toolName: TTool['name'] + toolCallId: string + resolveInterrupt(output: InferToolOutput): void +} +``` + +The output is inferred from the tool's `outputSchema` and is validated before submission and again on the server when that schema exists. With no output schema, output is `unknown` and no schema validation is attempted. A raw output schema also remains `unknown` statically but is validated at runtime. `addToolResult` remains supported and delegates to this machinery when it targets a pending client-tool interrupt. It is not deprecated. + +### Hook and headless-client surface + +Every framework hook exposes the same logical surface from the headless client: + +```ts +const { + interrupts, + resolveInterrupts, + cancelInterrupts, + retryInterrupts, + interruptErrors, + isResuming, + resumeInterruptsUnsafe, +} = useChat(options) +``` + +The root operations are: + +```ts +resolveInterrupts(decision: boolean): Promise + +resolveInterrupts( + resolve: (interrupt: ChatInterrupt) => undefined, +): Promise + +cancelInterrupts(): Promise + +retryInterrupts(): Promise + +resumeInterruptsUnsafe( + entries: readonly RunAgentResumeItem[], + state?: ChatResumeState, +): Promise +``` + +The boolean overload is intentionally narrow. It is allowed only when every open interrupt is a tool approval and the selected decision branch requires no custom payload. It never edits arguments. The client rejects unsupported boolean bulk operations before staging anything. Heterogeneous batches and payload-bearing decisions use the callback overload. + +Root promises resolve only after the server accepts the atomic response batch. They reject with a `BatchInterruptError` when validation, transport, protocol, expiry, staleness, or conflict prevents acceptance. Bound item methods are synchronous and return `void`; they update local state and may start asynchronous submission internally. + +`resumeInterruptsUnsafe` is the raw protocol escape hatch. It bypasses the type-safe builder API but still passes through server validation, exact-set checks, expiry checks, and atomic persistence. Its name is deliberately explicit. Normal application code should use bound or root resolvers. + +## Client staging and submission state machine + +The client maintains one draft slot per open interrupt. A draft is a resolved response or a cancellation. + +```text +pending -> staged -> submitting -> accepted + ^ | | + | | +-> error + +----------+ +``` + +The state machine follows these rules: + +1. `resolveInterrupt` constructs a candidate synchronously and starts client preflight. Synchronous validators finish in the call; an async Standard Schema validator sets status to `validating` and the submission gate waits for it. +2. `cancelInterrupt` stages `status: 'cancelled'` with no payload and needs no response-schema validation. +3. A valid candidate replaces the prior draft, clears superseded item errors, and sets status to `staged` while no submission is in flight. +4. An invalid replacement does not destroy the last valid draft. It sets status to `error`, records errors for the attempted candidate, preserves the previous `stagedResponse` for display, and blocks submission until the user supplies a valid replacement or clears it. +5. Async validation captures the item's draft generation. A result from an older candidate is discarded and cannot overwrite newer state. +6. `clearResolution()` removes the draft and candidate errors, increments the draft generation, and returns the item to `pending` while no submission is in flight. +7. With one open interrupt, a valid response or cancellation immediately starts submission after preflight completes. +8. With multiple open interrupts, item calls do not submit until every interrupt has a valid draft and no item is validating or in `error`. +9. When the final uncovered interrupt becomes valid, the client freezes the complete canonical batch, its draft generations, and fingerprint, then starts one submission. +10. During `submitting`, item replacement and clearing are rejected to avoid mutating the in-flight fingerprint. +11. A transport or retryable availability failure unfreezes interaction but retains the complete valid batch and its retry fingerprint. Items return to `staged`, root `interruptErrors` receives the failure, and `retryInterrupts()` is enabled. +12. Editing, replacing, cancelling, or clearing any item after that failure increments the batch generation, discards the frozen retry fingerprint, and clears the superseded transport error. `retryInterrupts()` then rejects until a new complete batch is frozen. +13. A server validation failure unfreezes the batch, maps every item error, and sets affected items to `error` without deleting their last candidate. A corrected valid candidate replaces it. +14. Pending interrupts block unrelated new user input as required by AG-UI. +15. After acceptance, the resolved descriptors and drafts leave the active `interrupts` collection and the continuation owns subsequent state. + +The callback overload is a synchronous transaction. Its return type is `undefined`, not `void`, because TypeScript permits an async function where a callback merely returns `void`: + +```ts +await resolveInterrupts((interrupt) => { + if (interrupt.kind === 'tool-approval') { + if (interrupt.toolName === 'transferFunds') { + interrupt.resolveInterrupt(true, { + payload: { note: 'Reviewed' }, + }) + return + } + + interrupt.cancelInterrupt() + return + } + + if (interrupt.kind === 'client-tool-execution') { + interrupt.cancelInterrupt() + return + } + + interrupt.resolveInterrupt(buildGenericResponse(interrupt)) +}) +``` + +The client invokes the callback exactly once for each interrupt in a stable snapshot. It passes transaction-scoped bound wrappers carrying a unique transaction token and draft-generation number. Item auto-submission is suppressed during the callback. On normal `undefined` return, every snapshot item must be covered; the complete batch is submitted once and the token is sealed. + +If the callback throws, returns any non-`undefined` value, returns a promise or thenable, or leaves an interrupt uncovered, all staging changes made by that callback are rolled back, its token and generation are invalidated, and no network request starts. Every transaction-scoped item method checks that token before touching state. This prevents a bypassed async callback from returning a promise, awaiting, and then mutating live drafts after rollback. Such a late call is ignored, records a root `inactive-transaction` error, and never submits. The original non-transaction bound objects remain usable after the failed root operation. + +`cancelInterrupts()` stages cancellation for every open interrupt and submits one batch. A cancellation counts as coverage but is never rewritten as an approval denial. + +`retryInterrupts()` resubmits the same frozen complete batch only after a retryable transport or server availability failure. Expired, stale, and conflicting batches refresh authoritative state and cannot be retried blindly. + +## Validation and error propagation + +### JSON Schema contract + +`[TanStack]` All serialized `responseSchema`, raw input, raw output, and raw approval schemas use JSON Schema Draft 2020-12. An absent `$schema` is interpreted as Draft 2020-12. A different declared dialect is rejected rather than silently reinterpreted. + +Browser and server use Ajv 8's 2020 entry point plus `ajv-formats` in full mode with the same options: `allErrors: true`, strict schema checking, format validation enabled, type coercion disabled, defaults disabled, and data mutation disabled. Unknown formats are schema compilation errors. Only document-local references, including `#` and `#/$defs/...`, are supported. Network, file, and cross-document `$ref` resolution is disabled in both environments. Unresolved references and invalid schemas produce `invalid-response-schema`; a TanStack server rejects an invalid schema before emitting an interrupt descriptor. + +Ajv `instancePath` values are decoded using RFC 6901 escaping. Normalized public paths are arrays of decoded string segments; validators that natively report numeric array indexes may use numbers. For `required` and `additionalProperties` errors, the missing or extra property is appended to the instance path. Multiple Ajv errors remain multiple errors and retain deterministic Ajv evaluation order within an item; items are ordered by interrupt ID in the aggregate response. + +Standard Schema validators run through their own validation surface and their issue paths are normalized to the same public path form. Client preflight may await an async Standard Schema validator behind the synchronous staging API. Server validation may also await it before commit. The serialized envelope is additionally checked with the common Ajv configuration. Validation never coerces, fills defaults, strips properties, or mutates `editedArgs` or payloads. + +If a client receives an invalid generic schema from a non-TanStack source, resolved submission is disabled for that item and it receives `invalid-response-schema`; cancellation remains available because it has no payload. `resumeInterruptsUnsafe` cannot bypass authoritative schema compilation. + +Validation happens in two phases: + +- Client preflight gives immediate feedback and avoids known-invalid requests. +- `[TanStack]` Server validation is mandatory and authoritative because clients, schemas, state, and expiry can be stale or bypassed. Exhaustive validation and no-write-on-any-invalid are stronger TanStack guarantees, not additional claims about the AG-UI base protocol. + +The server validates every entry in the proposed batch before persisting or executing anything. It collects all failures rather than stopping at the first failure. Checks include: + +- The response covers exactly the current open interrupt ID set. +- Every ID belongs to the interrupted run and has not expired. +- Each entry uses an allowed `resolved` or `cancelled` status. +- A cancelled entry has no payload. +- A resolved generic payload matches `responseSchema` when present. +- A tool approval envelope selects a valid branch. +- Approved `editedArgs`, when present, match `inputSchema` as a full replacement. +- Rejection never includes `editedArgs`. +- The nested custom payload matches the selected approval schema branch. +- A client-tool output matches `outputSchema`. + +If any validation fails, the server performs no interrupt response write, schedules no continuation tools, and returns one AG-UI `RUN_ERROR`. It does not emit a second error event or a compensating `RUN_FINISHED`. Because `RUN_ERROR` is the protocol's error event and its schema permits extension fields, TanStack attaches a serialized `InterruptSubmissionError[]` at the top-level namespaced field `tanstack:interruptErrors`. Generic AG-UI consumers still receive the standard error message and code. + +`scope: 'item'` errors populate only the matching bound item's `errors`; `scope: 'batch'` errors populate only root `interruptErrors`. An unknown interrupt ID, missing item, or invalid exact-set check is a batch error unless the authoritative state can correlate the problem to one still-current item. Every error carries thread, interrupted-run, and generation correlation; submission and continuation IDs are included when known. The promise rejection is a `BatchInterruptError` with `item-validation-failed` when only item errors caused failure, and its `interruptIds` lists every affected item. + +Retryability is explicit, not inferred only from code. Transport and transient server availability errors are retryable when the frozen fingerprint remains intact. Client validation, invalid schema, incomplete batch, unsupported bulk, async resolver, inactive transaction, mixed provenance, persistence capability, expiry, staleness, and conflict errors are not retryable without a state or draft change. A server error may mark itself retryable only when it proves no commit occurred. + +`interruptErrors` is cleared or replaced by the next relevant operation. A transport failure keeps the complete valid draft batch staged and adds a retryable root error. Stale, expired, and conflict responses replace local descriptors and drafts with authoritative server state before exposing the non-retryable error. + +### Authoritative recovery state + +`[TanStack]` Stale, expired, replayed, and conflicting submissions use one versioned recovery DTO: + +```ts +interface InterruptRecoveryStateV1 { + schemaVersion: 1 + state: 'pending' | 'committed' | 'expired' | 'missing' | 'legacy-committed' + threadId: string + interruptedRunId: string + generation: number + pendingInterrupts: readonly Interrupt[] + committed?: { + fingerprint: string + resolutions: readonly RunAgentResumeItem[] + continuationRunId?: string + committedAt: string + } +} +``` + +For every newly committed batch, `state: 'committed'` requires `committed.continuationRunId`; that is the winning continuation ID. `resolutions` is included only for authorized callers and is redacted according to the same policy as interrupt payload persistence. `legacy-committed` exists only for migrated pre-batch records whose historical continuation ID cannot be reconstructed. `pendingInterrupts` is nonempty only in `pending`; expired, missing, and committed states return an empty set. + +The primary recovery wire path is the top-level `tanstack:interruptRecovery` field on the same `RUN_ERROR` that reports stale, expired, or conflict. The server obtains this DTO in the same persistence read or transaction that detects the condition, so its generation and winning commit are consistent with the error. + +If the stream ends before that event or a transport cannot carry extension fields, the client calls an explicit recovery fetch contract: + +```ts +type InterruptStateFetcher = ( + input: { + threadId: string + interruptedRunId: string + knownGeneration: number + }, + options: { signal: AbortSignal }, +) => Promise +``` + +Connection adapters expose the equivalent `loadInterruptState` method. Fetcher-mode clients receive `interruptStateFetcher` beside `fetcher`. Server integrations expose the matching authenticated `getInterruptRecoveryState` handler over their chosen transport; the HTTP helper returns the DTO as JSON. The standard fetch and connection bridges wire these contracts automatically. This is a TanStack recovery extension, not an AG-UI run request. + +The client accepts only a matching thread and interrupted-run ID and a generation at least as new as its current generation. `pending` replaces descriptors, rebinds methods from the configured tool registry, and reconciles drafts. `committed` clears drafts and attaches to or reloads the winning continuation. `expired` clears drafts and reports expiry. `missing` and `legacy-committed` clear unsafe retry state and report a non-retryable protocol or migration error. If neither the event extension nor a configured recovery fetcher is available, state is not guessed: drafts remain quarantined, submission is disabled, and root receives `recovery-unavailable`. + +## Client draft durability + +Draft decisions are client state, not final server responses. `ChatPersistenceController` owns serialization, hydration, generation invalidation, and cleanup. The configured `persistence.server` `ChatStorageAdapter` remains an opaque get/set/remove transport keyed by thread ID; it never creates bound objects or decides reconciliation. + +The stored DTO is versioned and contains raw JSON only: + +```ts +interface ChatResumeSnapshotV2 { + schemaVersion: 2 + resumeState: ChatResumeState + interruptState?: { + provenance: 'native' | 'legacy' + phase: 'pending' | 'accepted' + interruptedRunId: string + generation: number + descriptorsHash: string + descriptors: readonly Interrupt[] + drafts: readonly { + interruptId: string + responseSchemaHash: string + response: RunAgentResumeItem + savedAt: string + }[] + continuationRunId?: string + } +} +``` + +The effective draft identity is `(threadId, interruptedRunId, generation, interruptId, responseSchemaHash)`. Functions, validators, tool implementations, `kind`, and hydrated error objects are never serialized. On hydration, the controller validates the DTO shape and version, obtains authoritative recovery state when native interrupts are present, then rebuilds bound variants from descriptors plus the configured tool registry. A draft survives only when the authoritative run, generation, exact interrupt ID, canonical response schema hash, and expiry match. Everything else is discarded. + +After acceptance, the client clears active in-memory descriptors and drafts immediately, writes an `accepted` tombstone containing the continuation ID and no drafts, then queues adapter removal. If removal fails, the tombstone prevents an old pending UI from reappearing; a later authoritative recovery confirms the commit and cleanup retries. Thread reset and disposal also invalidate pending writes by generation. + +The unversioned PR #785 snapshot shape is treated as version 1. Migration copies `resumeState` and raw `pendingInterrupts` into a version 2 native descriptor set with no drafts, generation `0`, and a mandatory recovery refresh before resolution. Invalid or uncorrelatable version 1 snapshots are removed and produce a root `protocol` error. Legacy CUSTOM-event state continues to hydrate from message history through the compatibility backend, not by pretending it is a native resume snapshot. + +`pendingInterrupts` is a deprecated getter alias over the same hydrated `interrupts` array. It does not read a second snapshot, does not persist functions, and does not return raw descriptor DTOs. + +This persistence allows a user to answer part of a multi-interrupt batch, reload, and continue. It does not make a partial response visible to the server. Only the final complete batch is committed server-side. + +## Atomic server persistence and concurrency + +Every persistence adapter must implement an atomic batch compare-and-swap operation. Sequential per-item `resolve` and `cancel` calls are not an acceptable fallback. + +The store contract adds this normative operation: + +```ts +commitInterruptResolutions({ + threadId, + interruptedRunId, + continuationRunId, + expectedGeneration, + expectedInterruptIds, + resolutions, + fingerprint, +}): Promise< + | { status: 'committed'; continuationRunId: string } + | { status: 'replayed'; continuationRunId: string } + | { status: 'conflict'; authoritativeState: InterruptBatchState } +> +``` + +The commit invariant is: + +```text +All currently pending interrupt responses are persisted, or none are. +``` + +The server computes a stable canonical fingerprint over the sorted complete resolution set. Property ordering and interrupt input order do not affect it. The transaction compares the exact expected pending ID set, records the response batch and continuation run ID together, and transitions every item atomically. + +Concurrency semantics are: + +- The first valid complete batch commits and becomes authoritative. +- An exact replay with the same fingerprint is accepted idempotently as `replayed` and returns the first continuation run ID. +- A later different batch for the same interrupted run returns `conflict` and the authoritative state. +- A replay attaches to or replays the existing continuation. It never schedules the tools a second time. +- A missing, extra, expired, already-transitioned, or otherwise stale expected ID set cannot commit. + +Database adapters use a transaction with a uniqueness or version predicate. The memory adapter uses a per-interrupted-run critical section and versioned batch record. Custom persistence adapters must implement the atomic operation; servers reject safe interrupt resumption when the configured adapter cannot provide it. + +The atomic response record prevents duplicate continuation scheduling at the resume gateway. Existing run locks and checkpoints remain responsible for continuation recovery. Arbitrary external tools must still be idempotent when exactly-once side effects matter across process crashes. + +This operation applies only to final client interrupt responses. Client drafts, general messages, checkpoints, and tool-result persistence retain their existing ownership and contracts unless separately changed. + +### Persistence availability and adapter work + +`[TanStack]` Safe native interrupt resumption requires an `InterruptStore` with atomic batch capability. Without configured persistence, the server must not emit an interrupt outcome that it cannot later resolve safely. At the point an operation would interrupt, it emits `RUN_ERROR` with `persistence-required` instead. Development applications may configure the existing memory persistence explicitly; production and multi-process applications use a durable adapter. There is no hidden process-global fallback. + +Built-in work is mandatory: + +- Memory: update `packages/ai-persistence/src/memory.ts` with a per-interrupted-run critical section, generation counter, immutable batch record, canonical fingerprint, replay lookup, and recovery DTO projection. It is single-process and loses state on restart, which documentation must state. +- Drizzle: update schema and stores in `packages/ai-persistence-drizzle`, add a numbered migration and packaged migration asset, and perform interrupt row transitions plus batch insert in one database transaction. +- Prisma: update the packaged and development Prisma models in `packages/ai-persistence-prisma`, its stores and model CLI, and document the required user migration before deploying the new package. +- Cloudflare: add the equivalent D1 migration and packaged asset in `packages/ai-persistence-cloudflare`, use a D1 transaction for the CAS, and leave R2 artifact storage unchanged. + +Durable schemas add an `interrupt_batches` record keyed by interrupted run ID with thread ID, generation, expected ID set, fingerprint, serialized resolutions, winning continuation run ID, and commit time. Interrupt rows gain a generation column. Indexes cover `(run_id, status, generation)`, `(thread_id, status)`, and unique continuation run ID. The CAS transaction checks generation and the exact pending set, transitions rows, and inserts the batch record together. + +Migration behavior is deterministic. Existing pending interrupt rows backfill to generation `1` and can enter the new CAS. Existing resolved or cancelled rows remain historical and are marked as legacy committed; because no reliable winning continuation ID can be inferred, recovery returns `legacy-committed` and never replays them. New writes always have a batch record. Migrations must be idempotent under the repository migration tooling and safe on nonempty databases. + +The shared persistence testkit gains interrupt batch conformance covering commit, rollback, replay, conflict, recovery, generation, and concurrent writers. Memory, Drizzle, Prisma, and Cloudflare must all run it. A custom adapter lacking the capability fails server construction when interrupt middleware is configured, or fails the first dynamically discovered interrupt with `atomic-commit-unsupported` before emitting an interrupt outcome. It never falls back to sequential writes. + +## AG-UI wire lifecycle + +### Interrupted run + +The server emits any required state before terminating the run: + +1. Emit `MESSAGES_SNAPSHOT` when continuation needs updated message state. +2. Emit `STATE_SNAPSHOT` when continuation needs updated agent state. +3. Emit one `RUN_FINISHED` whose outcome is `type: 'interrupt'` and whose interrupt list is nonempty. + +The interrupted run is terminal. Its interrupt descriptors contain globally unique IDs, canonical or namespaced reasons, optional `responseSchema`, optional expiry, and enough metadata to bind known tools without weakening generic fallback behavior. + +Tool approvals use reason `tool_call`. General requests use the canonical `input_required` or `confirmation` reason when their semantics match it. TanStack client-tool execution uses `tanstack:client_tool_execution`. Other TanStack-specific reasons use the `tanstack:` namespace. Unknown reasons remain generic. + +### Resume request + +The client starts a new continuation run on the same thread. It sends one `resume` array containing every open interrupt exactly once. No new user message is required. + +- Approval is `status: 'resolved'` with `{ approved: true, ... }`. +- Denial is `status: 'resolved'` with `{ approved: false, ... }`. +- Cancellation is `status: 'cancelled'` with no payload. +- An approved edit is a complete `editedArgs` replacement in the resolved payload. + +The client checks `expiresAt` before sending. The server checks expiry authoritatively. An expired or stale resume ends with `RUN_ERROR`, refreshes authoritative state, and executes nothing. + +### Continuation run + +The continuation has a new run ID and the same thread ID. After the batch commits, a resumed server tool emits only `TOOL_CALL_RESULT` for the original `toolCallId`. It does not replay `TOOL_CALL_START`, `TOOL_CALL_ARGS`, or `TOOL_CALL_END`. Subsequent model output and tool calls proceed normally in the new run. + +Downstream behavior is defined per entry after the complete batch commits: + +- Approved tool approval: invoke the authoritative tool once with `editedArgs` when supplied, otherwise with the stored original arguments. Persist its actual tool-result history part and emit one `TOOL_CALL_RESULT` with the original `toolCallId`. +- Denied tool approval: do not invoke the tool. Persist a terminal tool-result history part whose value is `{ status: 'denied', payload?: unknown }` and emit one result-only `TOOL_CALL_RESULT` whose string content is the canonical JSON serialization of that value. This gives the continued model a resolved tool outcome without confusing denial with cancellation. +- Cancelled tool approval: do not invoke the tool and emit no `TOOL_CALL_RESULT`. Persist the cancellation in the interrupt batch audit record and mark the original call not-executed when reconstructing provider history. This follows the AG-UI parallel-interrupt example, which emits results for approved calls and treats the cancelled call as not-executed. +- Client-tool execution: treat the resolved output as the actual result, persist one tool-result history part, and emit one result-only `TOOL_CALL_RESULT`. Cancellation is not-executed and emits no result. +- Generic interrupt: pass the resolved payload or cancellation state to the continuation handler and emit no tool event unless the application itself creates a later tool call. + +A mixed batch may therefore produce actual approved results, synthetic denied results, and no event for cancelled entries. Denied synthetic results are available immediately after commit; approved tool results are emitted as executions finish, so cross-tool result order is not guaranteed. Correlation is solely by the original `toolCallId`, and every result-producing entry emits exactly once. The resumed model starts only after required terminal history records for the batch are available; cancelled calls are excluded from its actionable tool-call set. + +Snapshots and persisted resume state must make a refresh capable of reconstructing the pending interrupt set before response, and the authoritative resolution or continuation after response. + +## Compatibility, deprecation, and migration + +`[Compatibility]` New servers emit only native AG-UI interrupts. New clients temporarily retain legacy custom-event readers and a real legacy submission backend so they can consume an older server during migration. Old clients cannot consume the new native-only server approval flow, so the recommended deployment is a coordinated server and client upgrade. + +Every hydrated item has `provenance: 'native' | 'legacy'`. Native items submit an AG-UI `resume` batch. Legacy `approval-requested` and `tool-input-available` events normalize into bound objects with `provenance: 'legacy'`, but they submit through the old history protocol: + +1. Stage all current legacy approval decisions with the same singleton and all-items-covered gate used by native items. +2. On completion, clone the current message history and update every matching approval part in one in-memory transaction. +3. Send that complete updated history in one legacy follow-up request, using no native `resume` array. +4. If history construction fails, change no live message. If transport fails, restore the staged batch and expose `legacy-submit-failed`. + +This legacy backend supports the behavior the old protocol can represent: boolean approve or deny and existing client-tool results. Edited arguments, custom approval payloads, generic interrupts, native cancellation, expiry, atomic server CAS, and native conflict recovery are rejected with `legacy-unsupported`; they are not silently dropped. Multiple legacy approvals received together are updated in one history transaction and one request. Native and legacy provenance cannot coexist in one batch; if malformed input creates such a set, submission stops with `mixed-provenance`. + +The following are deprecated immediately and scheduled for removal at 1.0: + +- `addToolApprovalResponse` +- `pendingInterrupts` +- Legacy approval custom-event types +- Legacy approval custom-event readers + +`pendingInterrupts` remains a deprecated alias of the new bound `interrupts` collection during the transition. It does not expose a second state store. `addToolApprovalResponse` selects the native bound backend or legacy history backend from provenance and emits a deprecation notice in development. `addToolResult` similarly preserves the legacy history path for an old client-tool event and the typed native path for a native client-tool interrupt. Legacy readers never cause new servers to dual-emit. + +The raw PR #785 `resumeInterrupts` name is replaced by `resumeInterruptsUnsafe`. The normal API is `interrupts`, item resolvers, and root resolvers. + +There is no codemod. The migration guide will cover: + +1. Coordinated client and server package upgrades. +2. Replacing `pendingInterrupts` reads with `interrupts`. +3. Replacing `addToolApprovalResponse` with bound `resolveInterrupt` calls. +4. The distinction between denial and cancellation. +5. Adding `approvalSchema` only beside `needsApproval: true`. +6. Moving custom response data under nested `payload`. +7. Treating `editedArgs` as an optional full replacement. +8. Resolving simultaneous interrupts individually or with the root callback. +9. Validating generic `unknown` payloads from `responseSchema`. +10. Upgrading custom persistence adapters to atomic batch commit. +11. Using `resumeInterruptsUnsafe` only for raw protocol integration. + +The implementation will include a changeset that calls out the breaking protocol/API transition under the repository's pre-1.0 release convention. + +## Documentation Impact + +Implementation adds two user-facing pages: + +- `docs/chat/interrupts.md`: the native AG-UI lifecycle and complete client API, with server and client examples. +- `docs/migration/interrupts.md`: the manual pre-1.0 migration from custom approval events and PR #785 raw resume APIs. There is no codemod. + +Implementation updates these existing pages: + +- `docs/tools/tool-approval.md`: `needsApproval`, shared and branch-specific `approvalSchema`, nested payloads, full-replacement edits, denial, and cancellation. +- `docs/tools/client-tools.md`: separate client-tool execution interrupts and retained `addToolResult`. +- `docs/chat/persistence.md`: version 2 client resume snapshots and draft hydration. +- `docs/persistence/custom-stores.md`: mandatory atomic interrupt batch capability and conformance kit. +- `docs/persistence/migrations.md`: Drizzle, Prisma, and Cloudflare schema migration steps. +- `docs/persistence/delivery-durability.md`: replay, first-writer conflicts, process-memory limitations, and tool idempotency. +- `docs/architecture/approval-flow-processing.md`: native event flow, new-run continuation, and legacy compatibility boundary. +- `docs/reference/functions/toolDefinition.md`: the `approvalSchema` type surface and `SchemaInput` behavior. + +`docs/config.json` adds navigation entries and `addedAt: "2026-07-13"` for both new pages. It refreshes `updatedAt` to `2026-07-13` for every content-updated page listed above, leaving `addedAt` intact. The page examples show both server and client, avoid type-assertion casts, and demonstrate user-library conversion and validation of generic JSON Schema payloads. They distinguish `[AG-UI]` protocol rules from TanStack guarantees in prose rather than presenting CAS or exhaustive server validation as base-protocol requirements. + +This internal file, `docs/superpowers/specs/2026-07-13-ag-ui-interrupts-design.md`, stays ignored and is not added to `docs/config.json`. + +## Framework parity + +Interrupt state, staging, validation, persistence, batching, and submission live in `@tanstack/ai-client`. Framework packages only adapt reactivity and preserve generics. + +React, Preact, Solid, Vue, Svelte, and Angular must expose equivalent: + +- Bound `interrupts` +- Deprecated `pendingInterrupts` alias +- Item resolver, clearer, and canceller behavior +- Root resolve, cancel, and retry functions +- `interruptErrors` +- `isResuming` +- `resumeInterruptsUnsafe` +- Tool and approval-schema inference + +Framework-specific tests must prove that updates to staged status, errors, and replacement are observable in the framework's native reactive model. Behavior must not be reimplemented independently in wrappers. + +## Implementation sequence + +1. Extend core tool-definition types to preserve approval schema generics and generate branch-aware response schemas. +2. Add the atomic interrupt batch persistence contract and implement it in built-in adapters. +3. Correct server interrupt reasons, run lifecycle, exhaustive validation, batch commit, continuation IDs, and result-only resumed tool events. +4. Build the headless client's bound interrupt normalization, type mapping, staging transaction, validation, error mapping, draft persistence, and retry behavior. +5. Add root operations, the unsafe escape hatch, and the `addToolResult` delegation path. +6. Thread the headless surface and generics through all framework packages. +7. Add legacy client readers and deprecation annotations without server dual emission. +8. Add unit, integration, persistence, type, framework, and end-to-end coverage. +9. Publish the dedicated interrupts documentation, manual migration guide, docs metadata updates, and changeset. + +Protocol and persistence foundations precede client convenience APIs so tests can exercise one authoritative wire contract. Compatibility shims are added after the new internal model exists, ensuring they delegate to one implementation. + +## Test Strategy + +Tests are layered so type construction and state-machine failures are found before persistence races and end-to-end transport runs. Existing suites are extended where they already own the behavior; new focused files are added only where mixing concerns would obscure failures. + +Suite ownership is exact: + +- Core protocol and execution: `packages/ai/tests/interrupts.test.ts` and new `packages/ai/tests/interrupts-types.test-d.ts`. +- Headless API and state machine: new `packages/ai-client/tests/chat-client-interrupts.test.ts`, plus `packages/ai-client/tests/chat-client-resume.test.ts` and `packages/ai-client/tests/chat-persistence-controller.test.ts`. +- Persistence middleware and memory: `packages/ai-persistence/tests/interrupts.test.ts`, `packages/ai-persistence/tests/memory.conformance.test.ts`, and interrupt coverage added to `packages/ai-persistence/src/testkit/conformance.ts`. +- Durable adapters: `packages/ai-persistence-drizzle/tests/drizzle.conformance.test.ts`, `packages/ai-persistence-drizzle/tests/migrations.test.ts`, `packages/ai-persistence-prisma/tests/prisma.conformance.test.ts`, `packages/ai-persistence-prisma/tests/models.test.ts`, `packages/ai-persistence-cloudflare/tests/runtime.conformance.test.ts`, and `packages/ai-persistence-cloudflare/tests/migrations.test.ts`. +- Framework behavior and types: `packages/ai-react/tests/use-chat.test.ts`, `packages/ai-react/tests/use-chat-types.test.ts`, the equivalent `use-chat` files in `ai-preact`, `ai-solid`, and `ai-vue`, `packages/ai-svelte/tests/use-chat.test.ts`, `packages/ai-svelte/tests/create-chat-types.test.ts`, `packages/ai-angular/tests/inject-chat.test.ts`, and `packages/ai-angular/tests/inject-chat-types.test.ts`. +- Browser-to-server behavior: new `testing/e2e/tests/interrupts.spec.ts` with its route fixture under `testing/e2e/src/routes/`. +- Documentation and links: the repository `test:docs` target after the pages and `docs/config.json` entries in Documentation Impact are changed. + +### Type tests + +- `approvalSchema` is accepted only with `needsApproval: true`. +- A tool with preserved `TNeedsApproval: false` or no approval marker produces `never` when extracted by `{ kind: 'tool-approval'; toolName: thatName }`, while a literal-true tool appears in the mapped approval union. +- Shared schema types both branches. +- Branch schemas type approval and rejection independently. +- An omitted branch forbids custom payload. +- Required and optional payload schemas produce the correct options arity. +- Approved resolution accepts optional full-replacement `editedArgs` inferred from `inputSchema`. +- Rejection rejects `editedArgs`. +- Custom fields are accepted only under nested `payload`. +- Client-tool resolution infers `outputSchema`. +- Generic interrupt resolution accepts `unknown` without unsafe inferred narrowing. +- Generics survive definition, client tool, server tool, headless client, and every framework hook. + +### Client unit tests + +- Singleton resolve and cancel auto-submit. +- Multi-item resolution waits for all items. +- Staged responses can be replaced and cleared before submission. +- Root callback visits each item once, suppresses intermediate submission, and submits once. +- Root callback rollback covers throws, non-`undefined` returns, promises, and incomplete batches. +- A deliberately cast async callback returns a promise, awaits a deferred gate, then calls an item method; token invalidation prevents the delayed post-await call from changing drafts or submitting. +- Boolean bulk resolution accepts only the allowed homogeneous case. +- Invalid replacement preserves the prior valid draft, sets item error state, and blocks submission until corrected or cleared. +- Editing or clearing after transport failure invalidates the frozen retry fingerprint. +- Transport failure retains a retryable frozen batch. +- Stale, expired, and conflict responses refresh authoritative state. +- Item and root errors map correctly and clear when superseded. +- Drafts survive reload and invalid drafts are discarded during reconciliation. +- Pending interrupts block unrelated new input. +- `addToolResult` and deprecated compatibility entry points delegate without duplicate submission. + +### Server and protocol tests + +- Snapshots precede interrupt `RUN_FINISHED`. +- Interrupted and continuation run IDs differ while thread ID remains stable. +- Resume covers the exact open set. +- Approval, denial, and cancellation produce distinct correct wire entries. +- Canonical and namespaced reasons are used. +- Edited arguments fully replace original arguments. +- Resumed tools emit result only against the original `toolCallId`. +- Generic payloads, tool envelopes, edited arguments, and client-tool outputs validate against their schemas. +- All invalid items are returned in one `RUN_ERROR` and no tool executes. +- Expired and stale resumes fail without writes or execution. +- Stale, expired, and conflict failures carry a correlated recovery DTO; fallback recovery fetch applies only newer matching generations. +- Invalid Draft 2020-12 schemas, remote references, unknown formats, and normalized paths behave identically in browser and server. + +### Persistence and concurrency tests + +- All responses commit or none do. +- A failure in any validation leaves every interrupt pending. +- Exact replay returns the original continuation run ID and does not schedule twice. +- Two different concurrent batches produce one commit and one conflict. +- Property and input ordering do not change the canonical fingerprint. +- Missing, extra, stale, and expired ID sets cannot commit. +- Built-in memory and durable adapters meet identical semantics. +- A restart after commit reconnects to the authoritative continuation. +- Custom adapters without atomic batch support fail loudly. +- No configured persistence fails before an unresumable interrupt outcome can be emitted. +- Nonempty-database migrations backfill pending generation state and mark historical commits as legacy. + +### Compatibility and framework tests + +- New clients translate legacy approval custom events into the new bound model. +- New servers do not emit legacy approval events. +- Deprecated aliases share state with the new APIs. +- Every framework observes the same staging, error, replacement, resolve, cancel, and retry behavior. +- Framework type tests preserve tool-specific approval payload inference. + +### End-to-end scenarios + +- One approval with and without edited arguments. +- Three simultaneous approvals resolved one card at a time. +- A heterogeneous batch resolved through the root callback. +- Approve, deny, and cancel decisions in one batch. +- In a mixed batch, approved and denied entries each emit one result-only event, cancelled entries emit none, no tool start/args/end event is replayed, and history contains the specified actual, denied, and not-executed audit states. +- Generic runtime-schema validation failure followed by correction. +- Multiple invalid entries returned together and corrected in one follow-up. +- Reload after one of three local drafts, then complete the batch. +- Network failure, explicit retry, and exact replay. +- Two browser tabs submit different complete batches and observe first-writer conflict behavior. +- Client-tool execution remains separate and type-safe. + +Validation runs narrow package and type tests first, then the repository's mandatory pre-PR gate: + +```text +pnpm test:pr +pnpm --filter @tanstack/ai-e2e test:e2e +``` + +The end-to-end suite is mandatory because this changes observable behavior. + +## Risks and mitigations + +### Type complexity leaks into user errors + +Branch-aware conditional types can produce unreadable diagnostics. Keep public helper aliases named, document the two schema forms, and add compile-failure tests for the most common mistakes. + +### Runtime and static schema behavior diverge + +Generate the response JSON Schema from the same normalized approval-schema representation used by TypeScript helper types. Test every shared, branch-specific, absent, required, and optional combination at both type and runtime levels. + +### Partial persistence or duplicate continuations + +Sequential writes and best-effort fallback are forbidden. Require atomic adapter support, fingerprint the complete batch, store the continuation run ID in the same transaction, and test races with deterministic barriers. + +### Duplicate external tool side effects after crashes + +The response gateway prevents duplicate scheduling for exact replays, while run locks and checkpoints govern recovery. Documentation must still require idempotency keys for tools whose external side effects cannot tolerate retry. + +### Stale client drafts overwrite authoritative responses + +Draft keys include the interrupted run and interrupt ID, and reload reconciliation checks the authoritative set, schema identity, and expiry. CAS rejects any different response after the first commit. + +### Compatibility shims create two sources of truth + +Legacy readers and deprecated methods normalize into the same bound model. New servers emit only native interrupts. There is one client state store and one submission pipeline. + +### Structured validation errors exceed the base protocol + +Use standard `RUN_ERROR` fields for generic clients and a namespaced TanStack extension for the exhaustive error array. Do not introduce a competing error event. + +### Framework behavior drifts + +Keep all behavioral state in the headless client and limit wrappers to reactive projection. Run shared contract tests against every framework adapter. + +## Acceptance criteria + +The work is complete when all of the following are true: + +- Native AG-UI interrupt runs, resumes, snapshots, expiry, cancellation, and result-only tool continuation conform to the primary specification. +- The public bound and root APIs implement the exact staging and submission semantics in this design. +- Tool approvals are branch-aware and type-safe from `inputSchema` and `approvalSchema` across all frameworks. +- Generic interrupts remain runtime-validated with statically `unknown` payloads. +- Every server validation error is returned together, with no write or execution on any invalid batch. +- Final response persistence is atomic, exact replay is idempotent, conflicting replay is rejected, and the first valid writer wins. +- Client drafts survive safe reload and never become partial server responses. +- Legacy client readers and deprecated APIs delegate to the new model, while new servers emit no legacy approval events. +- `addToolResult` remains supported and client-tool execution remains a separate typed interrupt. +- The migration guide, feature documentation, changeset, type tests, protocol tests, framework tests, persistence tests, and end-to-end scenarios are complete. +- The mandatory local quality gates pass before commit and review. From d735e8502a9ff666c47028376274382f2f3b0b59 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 13 Jul 2026 16:43:38 +0200 Subject: [PATCH 2/6] docs: plan full AG-UI interrupt support --- .../plans/2026-07-13-ag-ui-interrupts.md | 5205 +++++++++++++++++ 1 file changed, 5205 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-13-ag-ui-interrupts.md diff --git a/docs/superpowers/plans/2026-07-13-ag-ui-interrupts.md b/docs/superpowers/plans/2026-07-13-ag-ui-interrupts.md new file mode 100644 index 000000000..0ba63a140 --- /dev/null +++ b/docs/superpowers/plans/2026-07-13-ag-ui-interrupts.md @@ -0,0 +1,5205 @@ +# AG-UI Interrupts Full-Spec Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Implement native AG-UI interrupts with typed bound resolution APIs, exhaustive validation, durable atomic batching, recovery, compatibility shims, and equivalent behavior in every supported framework. + +**Architecture:** Core owns the wire DTOs, Draft 2020-12 validation, tool approval type machinery, and a persistence capability interface so `@tanstack/ai` never imports a persistence package. Persistence implementations atomically open and compare-and-swap complete interrupt batches; the headless client owns hydration, staging, validation, retry, and draft durability, while framework packages expose thin reactive projections. Continuations use the same thread, a fresh `runId`, and the interrupted run as `parentRunId` and the CAS correlation key. + +**Tech Stack:** TypeScript 5.9, pnpm 11.9.0, Nx, Vitest 4, Ajv 8 Draft 2020-12, ajv-formats, Standard Schema, Drizzle/SQLite, Prisma, Cloudflare D1, React, Preact, Solid, Vue, Svelte, Angular, Playwright, TanStack Start + +--- + +## Scope, file structure, and dependency order + +The user already approved this feature-wide scope even though it changes more than five files. Keep each task narrowly scoped where feasible, do not add a codemod, and do not expand into general checkpoint, message, or tool-result persistence beyond what the interrupt transaction requires. + +Implementation must proceed in dependency order: + +1. Core schema validation and tool-definition types establish one runtime/static contract. +2. Core wire DTOs and the core-owned interrupt persistence capability remove package-cycle pressure. +3. The persistence contract, memory reference implementation, and shared conformance suite define atomic semantics. +4. Drizzle, Prisma, and Cloudflare D1 implement the same atomic contract before server resumption depends on it. +5. Server middleware and the chat engine enforce authoritative validation, event ordering, and continuation semantics. +6. The headless client implements bound interrupts and draft durability against the settled wire contract. +7. Compatibility readers and framework wrappers delegate to that single headless state machine. +8. Browser E2E, documentation, release metadata, review, and quality gates close the work. + +The planned file ownership is: + +| Unit | Files | Responsibility | +| --- | --- | --- | +| Core validation | `packages/ai/src/activities/chat/tools/json-schema-validator.ts`, `packages/ai/src/activities/chat/tools/approval-schema.ts` | Compile Draft 2020-12 schemas, normalize issues, distinguish approval schema forms, generate response envelopes and schema identities. | +| Core public contract | `packages/ai/src/interrupts.ts`, `packages/ai/src/types.ts`, `packages/ai/src/activities/chat/tools/tool-definition.ts`, `packages/ai/src/activities/chat/middleware/types.ts` | Public interrupt errors/recovery DTOs, capability seam, conditional generics, continuation state. | +| Persistence core | `packages/ai-persistence/src/types.ts`, `packages/ai-persistence/src/interrupts.ts`, `packages/ai-persistence/src/memory.ts`, `packages/ai-persistence/src/testkit/conformance.ts` | Exact-set CAS contract, canonical fingerprint, memory critical section, reusable adapter tests. | +| Durable stores | `packages/ai-persistence-drizzle/**`, `packages/ai-persistence-prisma/**`, `packages/ai-persistence-cloudflare/**` | Transactional batches, migrations, packaged assets, API and migration tests. R2 remains unchanged. | +| Server lifecycle | `packages/ai-persistence/src/middleware.ts`, `packages/ai/src/activities/chat/index.ts` | Rebind current schemas, validate every entry, commit once, emit one structured error, snapshot before terminal, result-only continuation. | +| Headless client | `packages/ai-client/src/interrupt-manager.ts`, `packages/ai-client/src/chat-client.ts`, `packages/ai-client/src/chat-persistence-controller.ts`, `packages/ai-client/src/connection-adapters.ts`, `packages/ai-client/src/types.ts` | Bound union, staging transaction, recovery, raw V2 drafts, explicit recovery adapters, compatibility delegation. | +| Framework adapters | `packages/ai-{react,preact,solid,vue,svelte,angular}/src/**` | Reactive projection and generic preservation only. | +| Browser/docs/release | `testing/e2e/**`, `docs/**`, `.changeset/ag-ui-interrupts.md` | Approved scenarios, feature/migration guides, timestamps, breaking changeset. | + +## Spec coverage matrix + +| Approved requirement | Owning task(s) | +| --- | --- | +| Protocol wire types and structured item/root errors | 3, 8, 9 | +| Canonical `tool_call` and `tanstack:client_tool_execution` reasons | 9 | +| `MESSAGES_SNAPSHOT` / `STATE_SNAPSHOT` before interrupt terminal | 9, 16 | +| New continuation `runId`, same thread, interrupted run as `parentRunId` | 8, 9, 12, 16 | +| Result-only resumed tool events and mixed audit history | 9, 16 | +| Conditional `approvalSchema`, Standard Schema inference, `.client()` / `.server()` preservation | 2, 10, 15 | +| Tool approval, client-tool, and generic bound unions | 10 | +| Singleton auto-submit; multi-item wait; replacement; clear; callback rollback/late-call protection | 11 | +| Root boolean/callback/cancel/retry, statuses, errors, staged response, unsafe raw resume | 11, 12 | +| Draft 2020-12 Ajv validation and deterministic paths | 1, 2, 8, 10 | +| Versioned client draft persistence and recovery reconciliation | 13 | +| Exhaustive authoritative server validation, one error event, no write on invalid | 8 | +| Atomic CAS memory/Drizzle/Prisma/D1, exact replay, conflict, first writer, recovery | 4, 5, 6, 7 | +| Core-owned persistence capability without package cycles | 3, 8 | +| Validator rebind by registered tool name/schema hash; drift is stale | 8 | +| Explicit recovery handler/fetcher, never a guessed URL | 8, 13 | +| Generic continuation exposed to application middleware | 3, 8, 9 | +| Deprecations, legacy reader, no native dual writer, retained `addToolResult` | 12 | +| Sandbox legacy-event collision guard | 12 | +| React/Preact/Solid/Vue/Svelte/Angular parity | 14, 15 | +| Documentation timestamps, migration guide, breaking changeset, no codemod | 17 | +| All approved unit/type/framework/E2E scenarios | 1-17 | + +### Task 0: Prepare the isolated Windows execution environment + +**Prerequisites:** None. Run every later command from `F:\projects\tanstack\ai-interrupts-full-spec-design` in native Windows PowerShell. + +**Files:** +- Read: `AGENTS.md` +- Read: `CLAUDE.md` +- Read: `docs/superpowers/specs/2026-07-13-ag-ui-interrupts-design.md` +- Read: `docs/superpowers/plans/2026-07-13-ag-ui-interrupts.md` + +- [ ] **Step 1: Verify branch and working tree ownership** + +Run: + +```powershell +git status --short --branch +git branch --show-current +``` + +Expected: branch `codex/interrupts-full-spec-design`; only intentional plan/spec history is present before implementation. Preserve unrelated user changes if any appear. + +- [ ] **Step 2: Install with the repository-enforced package manager** + +Run: + +```powershell +corepack pnpm --version +pnpm install +``` + +Expected: pnpm reports `11.9.0` and install completes successfully. PowerShell environment assignments are process-local, so every later command block that runs Vitest, Nx, or Playwright sets `$env:CI='true'` itself. Run `pnpm install` again after any merge from `main` before resuming work. + +- [ ] **Step 3: Establish execution-time quality peers** + +Use the `test-hygiene` skill for a peer review before modifying test helpers or fixtures. Use the `debugging-discipline` skill whenever a red test fails for a reason other than the named missing symbol/behavior, and retain the investigation ledger for Task 18. These are execution requirements, not optional final polish. + +- [ ] **Step 4: Enforce subprocess hygiene** + +Use one-shot commands only; never start watch mode, a development server, or a WSL shell. Await each pnpm/Nx/Vitest/Playwright process tree before starting the next memory-heavy command. Do not kill all Node processes and do not run `wsl --shutdown` while any command is active. + +### Task 1: Add one Draft 2020-12 validator for browser and server + +**Prerequisites:** Task 0. + +**Files:** +- Create: `packages/ai/src/activities/chat/tools/json-schema-validator.ts` +- Modify: `packages/ai/src/index.ts:43-52` +- Modify: `packages/ai/src/client.ts:229-234,300-301` +- Modify: `packages/ai/package.json:82-94` +- Modify: `pnpm-lock.yaml` +- Test: `packages/ai/tests/interrupts.test.ts` + +- [ ] **Step 1: Write focused failing validator tests** + +Append these tests; they prove Ajv issue order, local references, repeated object reuse, cycle rejection, canonical JSON input, format handling, mutation safety, and RFC 6901 paths: + +```ts +import { + JsonSchemaCompilationError, + compileJsonSchema202012, +} from '../src/activities/chat/tools/json-schema-validator' + +it('normalizes every Draft 2020-12 issue without mutating input', () => { + const validate = compileJsonSchema202012({ + $schema: 'https://json-schema.org/draft/2020-12/schema', + type: 'object', + properties: { + 'a/b': { type: 'string', minLength: 3 }, + email: { type: 'string', format: 'email' }, + }, + required: ['a/b', 'email'], + additionalProperties: false, + }) + const input = { 'a/b': 'x', email: 'bad', extra: true } + + expect(validate(input)).toEqual([ + expect.objectContaining({ keyword: 'additionalProperties', path: ['extra'] }), + expect.objectContaining({ path: ['a/b'] }), + expect.objectContaining({ path: ['email'] }), + ]) + expect(input).toEqual({ 'a/b': 'x', email: 'bad', extra: true }) +}) + +it('accepts #, local $defs, and repeated acyclic object identities', () => { + const sharedProperty = { type: 'string', minLength: 2 } as const + const sharedObject = { + type: 'object', + properties: { label: sharedProperty }, + required: ['label'], + additionalProperties: false, + } as const + const root = { + type: 'object', + $defs: { entry: sharedObject }, + properties: { + selfCopy: { $ref: '#' }, + first: { $ref: '#/$defs/entry' }, + second: sharedObject, + third: sharedObject, + }, + additionalProperties: false, + } as const + const validate = compileJsonSchema202012(root) + + expect(validate({ + first: { label: 'ok' }, + second: { label: 'yes' }, + third: { label: 'no' }, + })).toEqual([]) +}) + +it('rejects unresolved local references, true cycles, Date, and Map', () => { + expect(() => compileJsonSchema202012({ $ref: '#/$defs/missing' })).toThrow( + JsonSchemaCompilationError, + ) + + const cyclic: Record = { type: 'object' } + cyclic['properties'] = { self: cyclic } + expect(() => compileJsonSchema202012(cyclic)).toThrow(/cycles/) + expect(() => compileJsonSchema202012(new Date())).toThrow(/plain JSON/) + expect(() => compileJsonSchema202012(new Map())).toThrow(/plain JSON/) +}) + +it('accepts canonical JSON primitives and rejects non-JSON values', () => { + const validate = compileJsonSchema202012({ + type: 'array', + items: { type: ['string', 'number', 'boolean', 'null'] }, + }) + expect(validate(['text', 1, true, null])).toEqual([]) + expect(() => compileJsonSchema202012({ const: undefined })).toThrow( + /JSON-compatible/, + ) +}) + +it('rejects other dialects, remote references, and unknown formats', () => { + expect(() => + compileJsonSchema202012({ + $schema: 'http://json-schema.org/draft-07/schema#', + type: 'string', + }), + ).toThrow(JsonSchemaCompilationError) + expect(() => + compileJsonSchema202012({ $ref: 'https://example.com/schema.json' }), + ).toThrow(/document-local/) + expect(() => + compileJsonSchema202012({ type: 'string', format: 'private-id' }), + ).toThrow(JsonSchemaCompilationError) +}) +``` + +- [ ] **Step 2: Run the red test** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai test:lib -- --run tests/interrupts.test.ts +``` + +Expected: FAIL because `json-schema-validator.ts` and `compileJsonSchema202012` do not exist. + +- [ ] **Step 3: Add dependencies and the minimal shared compiler** + +Add runtime dependencies `"ajv": "^8.20.0"` and `"ajv-formats": "^3.0.1"`, run `pnpm install`, and implement this public shape: + +```ts +import Ajv2020 from 'ajv/dist/2020.js' +import addFormats from 'ajv-formats' +import type { ErrorObject } from 'ajv' +import type { JSONSchema } from '../../../types' + +export interface JsonSchemaValidationIssue { + keyword: string + message: string + path: readonly (string | number)[] +} + +export class JsonSchemaCompilationError extends Error { + constructor(message: string, options?: ErrorOptions) { + super(message, options) + this.name = 'JsonSchemaCompilationError' + } +} + +function decodePointer(pointer: string): Array { + if (pointer === '') return [] + return pointer + .slice(1) + .split('/') + .map((part) => part.replaceAll('~1', '/').replaceAll('~0', '~')) +} + +function issuePath(error: ErrorObject): Array { + const path = decodePointer(error.instancePath) + if (error.keyword === 'required') { + path.push(String(error.params['missingProperty'])) + } + if (error.keyword === 'additionalProperties') { + path.push(String(error.params['additionalProperty'])) + } + return path +} + +const draft202012 = 'https://json-schema.org/draft/2020-12/schema' + +function assertSupportedSchemaTree(schema: unknown): asserts schema is JSONSchema { + const active = new WeakSet() + const visit = (value: unknown): void => { + if (value === null || ['string', 'number', 'boolean'].includes(typeof value)) { + return + } + if (typeof value !== 'object') { + throw new JsonSchemaCompilationError('Schema values must be JSON-compatible.') + } + if ( + !Array.isArray(value) && + Object.getPrototypeOf(value) !== Object.prototype && + Object.getPrototypeOf(value) !== null + ) { + throw new JsonSchemaCompilationError('Schema nodes must be plain JSON objects.') + } + if (active.has(value)) { + throw new JsonSchemaCompilationError('Schema values must not contain cycles.') + } + active.add(value) + if (Array.isArray(value)) { + value.forEach(visit) + active.delete(value) + return + } + const record = value as Record + if ( + typeof record['$schema'] === 'string' && + record['$schema'] !== draft202012 + ) { + throw new JsonSchemaCompilationError('Only Draft 2020-12 is supported.') + } + if ( + typeof record['$ref'] === 'string' && + !record['$ref'].startsWith('#') + ) { + throw new JsonSchemaCompilationError('Only document-local $ref values are supported.') + } + Object.values(record).forEach(visit) + active.delete(value) + } + visit(schema) +} + +export function compileJsonSchema202012( + schema: unknown, +): (value: unknown) => readonly JsonSchemaValidationIssue[] { + assertSupportedSchemaTree(schema) + const ajv = new Ajv2020({ + allErrors: true, + strict: true, + validateFormats: true, + coerceTypes: false, + useDefaults: false, + removeAdditional: false, + }) + addFormats(ajv, { mode: 'full' }) + let validate + try { + validate = ajv.compile(schema) + } catch (cause) { + throw new JsonSchemaCompilationError('Invalid Draft 2020-12 schema.', { + cause, + }) + } + return (value) => { + if (validate(value)) return [] + return (validate.errors ?? []).map((error) => ({ + keyword: error.keyword, + message: error.message ?? 'Schema validation failed.', + path: issuePath(error), + })) + } +} +``` + +The active recursion set is removed on unwind, so repeated acyclic object identities are legal while true cycles fail. Ajv owns local-reference resolution and issue evaluation order; do not sort `validate.errors`. Export the compiler and issue/error types from both core barrels; do not alter the provider-facing draft-07 conversion in `schema-converter.ts`. + +- [ ] **Step 4: Run the green test and type check** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai test:lib -- --run tests/interrupts.test.ts +pnpm --filter @tanstack/ai test:types +``` + +Expected: both commands PASS; the test reports all three normalized paths in deterministic Ajv order. + +- [ ] **Step 5: Refactor without changing behavior** + +Extract pointer decoding and recursive `$ref` traversal into private pure functions, keep one Ajv option object, rerun both commands, and confirm no validator path mutates input or performs network/file resolution. + +### Task 2: Make `approvalSchema` conditional, branch-aware, and generic-preserving + +**Prerequisites:** Task 1. + +**Files:** +- Create: `packages/ai/src/activities/chat/tools/approval-schema.ts` +- Create: `packages/ai/src/interrupt-serialization.ts` +- Modify: `packages/ai/src/activities/chat/tools/tool-definition.ts:12-227` +- Modify: `packages/ai/src/types.ts:133-153,601-735` +- Modify: `packages/ai/src/index.ts:43-52` +- Modify: `packages/ai/src/client.ts:229-234,300-301` +- Modify: `packages/ai/package.json:82-94` +- Modify: `pnpm-lock.yaml` +- Create: `packages/ai/tests/interrupts-types.test-d.ts` +- Modify: `packages/ai/tests/interrupts.test.ts` + +- [ ] **Step 1: Write failing compile-time and runtime cases** + +Create type tests covering shared/branch schemas, omitted branches, input/output inference, and both transformations: + +```ts +import { expectTypeOf } from 'vitest' +import { z } from 'zod' +import { + toolDefinition, + type ApprovalCapabilityOf, + type ApprovalSchemaOf, + type InferToolInput, + type InferToolOutput, + type InputSchemaOf, + type NoSchema, +} from '../src' + +const transfer = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema: z.object({ cents: z.number() }), + outputSchema: z.object({ receipt: z.string() }), + approvalSchema: { + approve: z.object({ note: z.string() }), + reject: z.object({ reason: z.string() }), + }, +}) + +expectTypeOf>().toEqualTypeOf<{ + cents: number +}>() +expectTypeOf>>().toEqualTypeOf<{ + receipt: string +}>() +expectTypeOf>>().toEqualTypeOf() +expectTypeOf>>().toEqualTypeOf() + +toolDefinition({ + name: 'invalid', + description: 'Cannot declare an approval payload', + // @ts-expect-error approvalSchema requires needsApproval: true + approvalSchema: z.object({ note: z.string() }), +}) + +const noInputDefinition = toolDefinition({ + name: 'noInput', + description: 'Approval without editable input', + needsApproval: true, +}) +const noInputClient = noInputDefinition.client() +const noInputServer = noInputDefinition.server(async () => ({ ok: true })) +expectTypeOf(noInputClient.inputSchema).toEqualTypeOf() +expectTypeOf(noInputServer.inputSchema).toEqualTypeOf() +expectTypeOf>().toEqualTypeOf() +expectTypeOf>().toEqualTypeOf() +expectTypeOf>().toEqualTypeOf() +``` + +Add runtime assertions: + +```ts +it('builds branch-aware approval envelopes and rejects malformed maps', () => { + const normalized = normalizeApprovalSchema({ + approve: { type: 'object', required: ['note'] }, + }) + expect(normalized.responseSchema).toMatchObject({ oneOf: expect.any(Array) }) + expect(normalized.branches.reject).toBeNull() + expect(() => normalizeApprovalSchema({})).toThrow(/approve or reject/) + expect(() => + normalizeApprovalSchema({ approve: { unsupported: Symbol('bad') } }), + ).toThrow(/SchemaInput/) +}) +``` + +- [ ] **Step 2: Run the red targets** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai test:types +pnpm --filter @tanstack/ai test:lib -- --run tests/interrupts.test.ts +``` + +Expected: type check FAILS because `approvalSchema`, `ApprovalCapabilityOf`, and `ApprovalSchemaOf` are missing; runtime test FAILS because `normalizeApprovalSchema` is missing. + +- [ ] **Step 3: Add the type marker and discriminated config** + +Replace the config interface with a discriminated intersection and add an approval generic to definition/client/server forms: + +```ts +export declare const toolApprovalCapability: unique symbol + +export interface ToolApprovalCapabilityMarker< + TNeedsApproval extends boolean, + TApprovalSchema, +> { + readonly [toolApprovalCapability]: { + needsApproval: TNeedsApproval + approvalSchema: TApprovalSchema + } +} + +export type ApprovalSchemaConfig = + | SchemaInput + | { approve: SchemaInput; reject?: SchemaInput } + | { approve?: SchemaInput; reject: SchemaInput } + +type ApprovalConfig< + TNeedsApproval extends boolean, + TApprovalSchema extends ApprovalSchemaConfig | undefined, +> = TNeedsApproval extends true + ? { needsApproval: true; approvalSchema?: TApprovalSchema } + : { needsApproval?: false; approvalSchema?: never } + +export type ToolDefinitionConfig< + TInput extends SchemaInput | undefined = undefined, + TOutput extends SchemaInput | undefined = undefined, + TName extends string = string, + TNeedsApproval extends boolean = false, + TApprovalSchema extends ApprovalSchemaConfig | undefined = undefined, +> = { + name: TName + description: string + inputSchema?: TInput + outputSchema?: TOutput + lazy?: boolean + metadata?: Record +} & ApprovalConfig + +export type ApprovalCapabilityOf = + TTool extends ToolApprovalCapabilityMarker + ? TNeeds + : false + +export type ApprovalSchemaOf = + TTool extends ToolApprovalCapabilityMarker + ? TSchema + : undefined + +export declare const noSchema: unique symbol +export type NoSchema = typeof noSchema + +export type InputSchemaOf = + TTool extends { inputSchema: infer TInput } + ? TInput extends undefined ? NoSchema : TInput + : NoSchema + +export type OutputSchemaOf = + TTool extends { outputSchema: infer TOutput } + ? TOutput extends undefined ? NoSchema : TOutput + : NoSchema +``` + +Apply `TInput extends SchemaInput | undefined = undefined` and `TOutput extends SchemaInput | undefined = undefined` verbatim to `ToolDefinitionConfig`, `ToolDefinition`, `ToolDefinitionInstance`, `ClientTool`, `ServerTool`, `toolDefinition()`, `.client()`, and `.server()`. Every returned intersection carries the exact `TInput`, `TOutput`, `TNeedsApproval`, and `TApprovalSchema` parameters. The runtime `inputSchema` and `outputSchema` fields remain `undefined` when omitted, while `InputSchemaOf`/`OutputSchemaOf` map only that literal absence to `NoSchema`. `InferToolInput` and `InferToolOutput` return `unknown` for a present raw JSON Schema but never widen an omitted schema into a present schema. The compile failures above prove `editedArgs` cannot appear for an omitted input schema before any client-manager work begins. + +- [ ] **Step 4: Normalize schemas once and generate the wire envelope** + +Implement `normalizeApprovalSchema` with this returned contract: + +```ts +export interface NormalizedApprovalSchema { + branches: { + approve: NormalizedSchemaInput | null + reject: NormalizedSchemaInput | null + } + responseSchema: JSONSchema + responseSchemaHash: string + approvalSchemaHash: string +} + +export function normalizeApprovalSchema( + approvalSchema: ApprovalSchemaConfig | undefined, + inputSchema?: SchemaInput, +): NormalizedApprovalSchema +``` + +Add `@noble/hashes` as a runtime dependency and implement `interrupt-serialization.ts` in core. This is the only canonicalization, digest, clone, and deep-freeze implementation used later by persistence and the client: + +```ts +import { sha256 } from '@noble/hashes/sha2.js' +import { bytesToHex, utf8ToBytes } from '@noble/hashes/utils.js' + +function canonical(value: unknown, active: WeakSet): string { + if (value === null) return 'null' + if (typeof value === 'string' || typeof value === 'boolean') { + return JSON.stringify(value) + } + if (typeof value === 'number' && Number.isFinite(value)) { + return JSON.stringify(value) + } + if (typeof value !== 'object') { + throw new TypeError('Interrupt values must be JSON-compatible.') + } + if (active.has(value)) throw new TypeError('Interrupt values must not cycle.') + if ( + !Array.isArray(value) && + Object.getPrototypeOf(value) !== Object.prototype && + Object.getPrototypeOf(value) !== null + ) { + throw new TypeError('Interrupt values must use plain JSON objects.') + } + + active.add(value) + let encoded: string + if (Array.isArray(value)) { + encoded = `[${value.map((item) => canonical(item, active)).join(',')}]` + } else { + const record = value as Record + encoded = `{${Object.keys(record) + .sort() + .map((key) => `${JSON.stringify(key)}:${canonical(record[key], active)}`) + .join(',')}}` + } + active.delete(value) + return encoded +} + +export function canonicalInterruptJson(value: unknown): string { + return canonical(value, new WeakSet()) +} + +export function digestInterruptJson(canonicalJson: string): string { + return `sha256:${bytesToHex(sha256(utf8ToBytes(canonicalJson)))}` +} + +function freezeTree(value: unknown): void { + if (value === null || typeof value !== 'object' || Object.isFrozen(value)) return + Object.values(value).forEach(freezeTree) + Object.freeze(value) +} + +export function cloneAndDeepFreezeJson(value: T): T { + const clone: T = JSON.parse(canonicalInterruptJson(value)) + freezeTree(clone) + return clone +} +``` + +Implement the approval normalizer with no undefined helper. `schemaToWire` checks Standard Schema first, retains its validator, uses its Standard JSON Schema export when present, and otherwise uses `{}` only for the nested wire payload while keeping the validator. Raw JSON Schema is recognized before a branch map by the keyword set: + +```ts +const jsonSchemaKeywords = new Set([ + '$schema', '$id', '$ref', '$defs', 'type', 'properties', 'required', + 'additionalProperties', 'items', 'oneOf', 'anyOf', 'allOf', 'enum', 'const', + 'format', 'minimum', 'maximum', 'minLength', 'maxLength', 'pattern', +]) + +function isPlainRecord(value: unknown): value is Record { + return value !== null && typeof value === 'object' && !Array.isArray(value) +} + +function isRawJsonSchema(value: unknown): value is JSONSchema { + return isPlainRecord(value) && + Object.keys(value).some((key) => jsonSchemaKeywords.has(key)) +} + +function isApprovalBranchMap( + value: unknown, +): value is { approve?: SchemaInput; reject?: SchemaInput } { + if (!isPlainRecord(value)) return false + const keys = Object.keys(value) + return ( + keys.length > 0 && + keys.every((key) => key === 'approve' || key === 'reject') && + keys.every((key) => isSchemaInput(value[key])) + ) +} + +function schemaToWire(schema: SchemaInput): NormalizedSchemaInput { + if (isStandardSchema(schema)) { + const jsonSchema = isStandardJSONSchema(schema) + ? schema['~standard'].jsonSchema.input({ target: 'draft-2020-12' }) + : undefined + return { source: schema, validator: schema, jsonSchema } + } + if (isRawJsonSchema(schema)) { + compileJsonSchema202012(schema) + return { source: schema, jsonSchema: schema } + } + throw new TypeError('Expected a supported SchemaInput.') +} + +function decisionEnvelope(input: { + approved: boolean + payload: NormalizedSchemaInput | null + inputSchema: NormalizedSchemaInput | null +}): JSONSchema { + const properties: Record = { + approved: { const: input.approved }, + } + const required = ['approved'] + if (input.approved && input.inputSchema) { + properties['editedArgs'] = input.inputSchema.jsonSchema ?? {} + } + if (input.payload) { + properties['payload'] = input.payload.jsonSchema ?? {} + required.push('payload') + } + return { + type: 'object', + properties, + required, + additionalProperties: false, + } +} + +export function normalizeApprovalSchema( + approvalSchema: ApprovalSchemaConfig | undefined, + inputSchema?: SchemaInput, +): NormalizedApprovalSchema { + const normalizedInput = inputSchema ? schemaToWire(inputSchema) : null + let approve: NormalizedSchemaInput | null = null + let reject: NormalizedSchemaInput | null = null + if (approvalSchema !== undefined) { + if (isStandardSchema(approvalSchema) || isRawJsonSchema(approvalSchema)) { + approve = schemaToWire(approvalSchema) + reject = approve + } else if (isApprovalBranchMap(approvalSchema)) { + approve = approvalSchema.approve + ? schemaToWire(approvalSchema.approve) + : null + reject = approvalSchema.reject + ? schemaToWire(approvalSchema.reject) + : null + } else { + throw new TypeError('approvalSchema must be a SchemaInput or nonempty approve/reject map.') + } + } + const responseSchema: JSONSchema = { + oneOf: [ + decisionEnvelope({ approved: true, payload: approve, inputSchema: normalizedInput }), + decisionEnvelope({ approved: false, payload: reject, inputSchema: null }), + ], + } + const responseCanonical = canonicalInterruptJson(responseSchema) + const approvalCanonical = canonicalInterruptJson({ + approve: approve?.jsonSchema ?? null, + reject: reject?.jsonSchema ?? null, + }) + return { + branches: { approve, reject }, + responseSchema, + responseSchemaHash: digestInterruptJson(responseCanonical), + approvalSchemaHash: digestInterruptJson(approvalCanonical), + } +} + +export function hashSchemaInput(schema: SchemaInput | undefined): string { + if (schema === undefined) return digestInterruptJson('undefined') + const normalized = schemaToWire(schema) + const identity = normalized.jsonSchema ?? { standardValidator: 'unserialized' } + return digestInterruptJson(canonicalInterruptJson(identity)) +} +``` + +Add the imported `isSchemaInput` guard to `schema-converter.ts` using the same Standard-first/raw-keyword checks, export it, and use it above. Add `NormalizedSchemaInput` with `source`, optional `validator`, and optional `jsonSchema`. Add runtime tests for every shared/branch/omitted form, malformed empty/unknown-key maps, Standard Schema with and without JSON export, raw schema compilation, response/approval hash stability, and rejection of function/Symbol/cyclic values. Never derive identity from function source text. + +- [ ] **Step 5: Run green type/runtime checks** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai test:types +pnpm --filter @tanstack/ai test:lib -- --run tests/interrupts.test.ts +``` + +Expected: both PASS; `.client()` and `.server()` preserve `true`, input/output, and approval-schema generics; malformed branch objects throw synchronously. + +- [ ] **Step 6: Focused refactor and recheck** + +Make the type-level `ApproveSchema`/`RejectSchema` selectors read the same normalized forms used at runtime, name conditional helpers so diagnostics remain legible, and rerun both commands. Confirm no provider schema-converter code changed. + +### Task 3: Define core wire DTOs, recovery DTO, and persistence capability + +**Prerequisites:** Task 2. + +**Files:** +- Create: `packages/ai/src/interrupts.ts` +- Modify: `packages/ai/src/types.ts:1081-1124,1371-1400,1480-1550` +- Modify: `packages/ai/src/activities/chat/middleware/types.ts:88-115,208-235` +- Modify: `packages/ai/src/index.ts:111-172` +- Modify: `packages/ai/src/client.ts:229-234,300-301` +- Modify: `packages/ai/tests/interrupts.test.ts` +- Modify: `packages/ai/tests/interrupts-types.test-d.ts` + +- [ ] **Step 1: Write failing DTO and capability tests** + +Add compile/runtime assertions for item/root errors, recovery correlation, generic continuation, and the core-owned capability: + +```ts +import { + InterruptPersistenceCapability, + canonicalizeInterruptResolutions, + type InterruptRecoveryStateV1, + type InterruptSubmissionError, +} from '../src/interrupts' + +it('correlates interrupt errors and recovery to the interrupted run', () => { + const recovery = { + schemaVersion: 1, + state: 'committed', + threadId: 'thread-1', + interruptedRunId: 'run-old', + generation: 2, + pendingInterrupts: [], + committed: { + fingerprint: 'v1:abc', + resolutions: [], + continuationRunId: 'run-new', + committedAt: '2026-07-13T10:00:00.000Z', + }, + } satisfies InterruptRecoveryStateV1 + const error: InterruptSubmissionError = { + scope: 'batch', + code: 'conflict', + message: 'Another response batch won.', + source: 'server', + retryable: false, + interruptIds: ['interrupt-1'], + threadId: 'thread-1', + interruptedRunId: 'run-old', + generation: 2, + } + + expect(recovery.committed.continuationRunId).toBe('run-new') + expect(error.interruptedRunId).toBe('run-old') + expect(InterruptPersistenceCapability.capabilityName).toBe( + 'interrupt-persistence', + ) +}) + +it('canonicalizes and deeply freezes response batches in core', () => { + const left = canonicalizeInterruptResolutions([ + { interruptId: 'b', status: 'cancelled' }, + { interruptId: 'a', status: 'resolved', payload: { z: 1, a: 2 } }, + ]) + const right = canonicalizeInterruptResolutions([ + { interruptId: 'a', status: 'resolved', payload: { a: 2, z: 1 } }, + { interruptId: 'b', status: 'cancelled' }, + ]) + expect(left.canonicalResolutions).toBe(right.canonicalResolutions) + expect(left.fingerprint).toMatch(/^sha256:[0-9a-f]{64}$/) + expect(left.fingerprint).toBe(right.fingerprint) + expect(Object.isFrozen(left.resolutions)).toBe(true) + expect(Object.isFrozen(left.resolutions[0])).toBe(true) +}) +``` + +- [ ] **Step 2: Run the red targets** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai test:lib -- --run tests/interrupts.test.ts +pnpm --filter @tanstack/ai test:types +``` + +Expected: FAIL because `interrupts.ts`, `InterruptRecoveryStateV1`, and `InterruptPersistenceCapability` do not exist. + +- [ ] **Step 3: Implement the core-owned seam and public DTOs** + +Define the complete error code unions from the approved spec, `InterruptCorrelation`, item/root errors, and recovery DTO. The atomic gateway contract must include opening descriptors atomically as well as committing responses: + +```ts +export interface OpenInterruptBatchInput { + threadId: string + interruptedRunId: string + descriptors: readonly Interrupt[] + bindings: readonly UnopenedInterruptBinding[] +} + +export interface CommitInterruptResolutionsInput { + threadId: string + interruptedRunId: string + continuationRunId: string + expectedGeneration: number + expectedInterruptIds: readonly string[] + resolutions: readonly RunAgentResumeItem[] + fingerprint: string + canonicalResolutions: string +} + +export interface InterruptCorrelation { + threadId: string + interruptedRunId: string + generation: number + submissionId?: string + continuationRunId?: string +} + +export type ItemInterruptErrorCode = + | 'invalid-payload' + | 'invalid-edited-args' + | 'invalid-tool-output' + | 'invalid-response-schema' + | 'unknown-interrupt' + | 'expired' + | 'stale' + | 'conflict' + | 'legacy-unsupported' + +export type BatchInterruptErrorCode = + | 'incomplete-batch' + | 'item-validation-failed' + | 'unsupported-bulk-operation' + | 'async-resolver' + | 'inactive-transaction' + | 'mixed-provenance' + | 'transport' + | 'server' + | 'protocol' + | 'invalid-response-schema' + | 'expired' + | 'stale' + | 'conflict' + | 'persistence-required' + | 'atomic-commit-unsupported' + | 'recovery-unavailable' + | 'legacy-submit-failed' + +export interface ItemInterruptError extends InterruptCorrelation { + scope: 'item' + interruptId: string + code: ItemInterruptErrorCode + message: string + path?: readonly (string | number)[] + source: 'client' | 'server' + retryable: boolean +} + +export interface BatchInterruptError extends InterruptCorrelation { + scope: 'batch' + code: BatchInterruptErrorCode + message: string + source: 'client' | 'server' | 'transport' + retryable: boolean + interruptIds: readonly string[] +} + +export type InterruptSubmissionError = ItemInterruptError | BatchInterruptError + +export type InterruptBinding = + | { + kind: 'tool-approval' + interruptId: string + interruptedRunId: string + generation: number + toolName: string + toolCallId: string + originalArgs: unknown + inputSchemaHash: string + approvalSchemaHash: string + responseSchemaHash: string + expiresAt?: string + } + | { + kind: 'client-tool-execution' + interruptId: string + interruptedRunId: string + generation: number + toolName: string + toolCallId: string + outputSchemaHash: string + responseSchemaHash: string + expiresAt?: string + } + | { + kind: 'generic' + interruptId: string + interruptedRunId: string + generation: number + responseSchemaHash: string + expiresAt?: string + } + +export type UnopenedInterruptBinding = InterruptBinding extends infer TBinding + ? TBinding extends InterruptBinding + ? Omit + : never + : never + +export type InterruptCommitResult = + | { status: 'committed'; continuationRunId: string } + | { status: 'replayed'; continuationRunId: string } + | { status: 'conflict'; authoritativeState: InterruptRecoveryStateV1 } + +export interface InterruptRecoveryQuery { + threadId: string + interruptedRunId: string + knownGeneration: number +} + +export interface InterruptRecoveryStateV1 extends InterruptCorrelation { + schemaVersion: 1 + state: 'pending' | 'committed' | 'expired' | 'missing' | 'legacy-committed' + pendingInterrupts: readonly Interrupt[] + committed?: { + fingerprint: string + resolutions?: readonly RunAgentResumeItem[] + continuationRunId?: string + committedAt: string + } +} + +export interface InterruptPersistenceGateway { + openInterruptBatch( + input: OpenInterruptBatchInput, + ): Promise<{ generation: number; descriptors: readonly Interrupt[] }> + commitInterruptResolutions( + input: CommitInterruptResolutionsInput, + ): Promise + getInterruptRecoveryState( + input: InterruptRecoveryQuery, + ): Promise +} + +export const InterruptPersistenceCapability = + createCapability()('interrupt-persistence') + +export const [getInterruptPersistence, provideInterruptPersistence] = + InterruptPersistenceCapability +``` + +Add this core-owned response canonicalizer; persistence stores compare both the SHA-256 digest and canonical bytes on replay: + +```ts +export function canonicalizeInterruptResolutions( + resolutions: readonly RunAgentResumeItem[], +): { + resolutions: readonly RunAgentResumeItem[] + canonicalResolutions: string + fingerprint: string +} { + const sorted = [...resolutions].sort((left, right) => + left.interruptId.localeCompare(right.interruptId), + ) + const frozen = cloneAndDeepFreezeJson(sorted) + const canonicalResolutions = canonicalInterruptJson(frozen) + return Object.freeze({ + resolutions: frozen, + canonicalResolutions, + fingerprint: digestInterruptJson(canonicalResolutions), + }) +} +``` + +Extend `RunErrorEvent` with optional top-level fields: + +```ts +export interface RunErrorEvent extends AGUIRunErrorEvent { + model?: string + 'tanstack:interruptErrors'?: readonly InterruptSubmissionError[] + 'tanstack:interruptRecovery'?: InterruptRecoveryStateV1 + error?: { message: string; code?: string | undefined } | undefined +} +``` + +Extend middleware context/config state without importing persistence: + +```ts +export interface ChatMiddlewareContext { + runId: string + parentRunId?: string + threadId: string +} + +export type ChatResumeGenericResolution = + | { interruptId: string; status: 'resolved'; payload: unknown } + | { interruptId: string; status: 'cancelled'; payload?: never } + +export interface ChatResumeToolState { + approvals?: ReadonlyMap + clientToolResults?: ReadonlyMap + genericInterrupts?: ReadonlyMap + deniedToolResults?: ReadonlyMap + cancelledToolCallIds?: ReadonlySet +} +``` + +Add `parentRunId` when `TextEngine` constructs `middlewareCtx`. Deprecate `ApprovalRequestedEvent` and `ToolInputAvailableEvent` with removal-at-1.0 JSDoc, but keep them in `KnownCustomEvent` for compatibility. Export every new public DTO/capability from root and client barrels. + +- [ ] **Step 4: Run green tests** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai test:lib -- --run tests/interrupts.test.ts +pnpm --filter @tanstack/ai test:types +``` + +Expected: both PASS; no dependency from `packages/ai` to any persistence package appears. + +- [ ] **Step 5: Refactor and verify the seam** + +Group wire-only types separately from capability methods in `interrupts.ts`, keep `parentRunId` optional for non-resume requests, and rerun the focused commands. Run `pnpm --filter @tanstack/ai test:lib -- --run tests/custom-events-integration.test.ts` and expect PASS to prove deprecated types remain readable. + +### Task 4: Replace sequential persistence with the atomic reference contract + +**Prerequisites:** Task 3. + +**Files:** +- Modify: `packages/ai-persistence/src/types.ts:36-55,139,255-287` +- Modify: `packages/ai-persistence/src/interrupts.ts` +- Modify: `packages/ai-persistence/src/capabilities.ts` +- Modify: `packages/ai-persistence/src/memory.ts:65-126,355` +- Modify: `packages/ai-persistence/src/testkit/conformance.ts:114-184` +- Modify: `packages/ai-persistence/src/index.ts` +- Modify: `packages/ai-persistence/tests/interrupts.test.ts` +- Modify: `packages/ai-persistence/tests/memory.conformance.test.ts` +- Modify: `packages/ai-persistence/tests/with-persistence.test.ts` +- Modify: `packages/ai-persistence/tests/persistence-validation.test.ts` +- Modify: `packages/ai-persistence/tests/persistence-types.test-d.ts` + +- [ ] **Step 1: Rewrite the shared conformance tests to fail on sequential stores** + +Replace the interrupt conformance block with complete-batch cases. Use stable IDs per test so adapters can run the same suite: + +```ts +it('commits all responses, replays exactly, and rejects a different writer', async () => { + const store = persistence.stores.interrupts + expect(store, 'interrupt conformance requires stores.interrupts').toBeDefined() + if (!store) throw new Error('interrupt store missing') + const opened = await store.openInterruptBatch({ + threadId: 'thread-cas', + interruptedRunId: 'run-interrupted', + descriptors: [ + { id: 'int-a', reason: 'confirmation' }, + { id: 'int-b', reason: 'input_required' }, + ], + bindings: [ + { interruptId: 'int-a', kind: 'generic', responseSchemaHash: 'v1:a' }, + { interruptId: 'int-b', kind: 'generic', responseSchemaHash: 'v1:b' }, + ], + }) + const resolutions = [ + { interruptId: 'int-b', status: 'cancelled' }, + { interruptId: 'int-a', status: 'resolved', payload: { ok: true } }, + ] as const + const canonical = canonicalizeInterruptResolutions(resolutions) + const input = { + threadId: 'thread-cas', + interruptedRunId: 'run-interrupted', + continuationRunId: 'run-continuation-a', + expectedGeneration: opened.generation, + expectedInterruptIds: ['int-b', 'int-a'], + resolutions: canonical.resolutions, + fingerprint: canonical.fingerprint, + canonicalResolutions: canonical.canonicalResolutions, + } + + await expect(store.commitInterruptResolutions(input)).resolves.toEqual({ + status: 'committed', + continuationRunId: 'run-continuation-a', + }) + await expect( + store.commitInterruptResolutions({ + ...input, + continuationRunId: 'ignored-on-replay', + }), + ).resolves.toEqual({ + status: 'replayed', + continuationRunId: 'run-continuation-a', + }) + await expect( + store.commitInterruptResolutions({ + ...input, + continuationRunId: 'run-continuation-b', + resolutions: [ + { interruptId: 'int-a', status: 'cancelled' }, + { interruptId: 'int-b', status: 'cancelled' }, + ], + fingerprint: canonicalizeInterruptResolutions([ + { interruptId: 'int-a', status: 'cancelled' }, + { interruptId: 'int-b', status: 'cancelled' }, + ]).fingerprint, + canonicalResolutions: canonicalizeInterruptResolutions([ + { interruptId: 'int-a', status: 'cancelled' }, + { interruptId: 'int-b', status: 'cancelled' }, + ]).canonicalResolutions, + }), + ).resolves.toMatchObject({ status: 'conflict' }) +}) +``` + +Define the testkit contract and fixtures in `src/testkit/conformance.ts`; the store is required and no case may return early: + +```ts +export interface InterruptConformanceHarness { + getStore(): InterruptStore | undefined + advanceBy(milliseconds: number): void + inspect(interruptedRunId: string): Promise<{ + statuses: readonly string[] + batchCount: number + }> + failTransitionOnce(interruptId: string): void + reopen?(): Promise +} + +export async function openTwoInterrupts(store: InterruptStore) { + return store.openInterruptBatch({ + threadId: 'thread-cas', + interruptedRunId: 'run-interrupted', + descriptors: [ + { id: 'int-a', reason: 'confirmation' }, + { + id: 'int-b', + reason: 'input_required', + expiresAt: '2026-07-13T10:01:00.000Z', + }, + ], + bindings: [ + { interruptId: 'int-a', kind: 'generic', responseSchemaHash: 'sha256:a' }, + { + interruptId: 'int-b', + kind: 'generic', + responseSchemaHash: 'sha256:b', + expiresAt: '2026-07-13T10:01:00.000Z', + }, + ], + }) +} + +export function validTwoItemCommit( + generation: number, + continuationRunId = 'run-continuation-a', +): CommitInterruptResolutionsInput { + const candidate = canonicalizeInterruptResolutions([ + { interruptId: 'int-b', status: 'cancelled' }, + { interruptId: 'int-a', status: 'resolved', payload: { ok: true } }, + ]) + return { + threadId: 'thread-cas', + interruptedRunId: 'run-interrupted', + continuationRunId, + expectedGeneration: generation, + expectedInterruptIds: ['int-b', 'int-a'], + resolutions: candidate.resolutions, + fingerprint: candidate.fingerprint, + canonicalResolutions: candidate.canonicalResolutions, + } +} + +export function runInterruptStoreConformance( + createHarness: () => Promise, +): void { + it('requires the atomic interrupt capability', async () => { + const harness = await createHarness() + expect( + harness.getStore(), + 'interrupt conformance requires stores.interrupts', + ).toBeDefined() + if (!harness.getStore()) throw new Error('interrupt store missing') + }) + + it('keeps canonical order stable and rejects malformed exact sets', async () => { + const harness = await createHarness() + const store = harness.getStore() + if (!store) throw new Error('interrupt store missing') + const opened = await openTwoInterrupts(store) + const ordered = validTwoItemCommit(opened.generation) + const reordered = canonicalizeInterruptResolutions([ + { interruptId: 'int-a', status: 'resolved', payload: { ok: true } }, + { interruptId: 'int-b', status: 'cancelled' }, + ]) + expect(reordered.fingerprint).toBe(ordered.fingerprint) + expect(reordered.canonicalResolutions).toBe(ordered.canonicalResolutions) + for (const expectedInterruptIds of [ + ['int-a'], + ['int-a', 'int-b', 'int-c'], + ['int-a', 'int-a'], + ]) { + await expect(store.commitInterruptResolutions({ + ...ordered, + expectedInterruptIds, + })).resolves.toMatchObject({ status: 'conflict' }) + } + await expect(store.commitInterruptResolutions({ + ...ordered, + expectedGeneration: opened.generation + 1, + })).resolves.toMatchObject({ status: 'conflict' }) + }) + + it('rolls back every row when a transition fails', async () => { + const harness = await createHarness() + const store = harness.getStore() + if (!store) throw new Error('interrupt store missing') + const opened = await openTwoInterrupts(store) + harness.failTransitionOnce('int-b') + await expect( + store.commitInterruptResolutions(validTwoItemCommit(opened.generation)), + ).rejects.toThrow() + await expect(harness.inspect('run-interrupted')).resolves.toEqual({ + statuses: ['pending', 'pending'], + batchCount: 0, + }) + }) + + it('projects pending, committed replay, restart, and one concurrent winner', async () => { + const harness = await createHarness() + const store = harness.getStore() + if (!store) throw new Error('interrupt store missing') + const opened = await openTwoInterrupts(store) + await expect(store.getInterruptRecoveryState({ + threadId: 'thread-cas', + interruptedRunId: 'run-interrupted', + knownGeneration: opened.generation, + })).resolves.toMatchObject({ + state: 'pending', + generation: opened.generation, + }) + + const first = validTwoItemCommit(opened.generation, 'run-continuation-a') + const secondCandidate = canonicalizeInterruptResolutions([ + { interruptId: 'int-a', status: 'cancelled' }, + { interruptId: 'int-b', status: 'cancelled' }, + ]) + const second = { + ...first, + continuationRunId: 'run-continuation-b', + resolutions: secondCandidate.resolutions, + fingerprint: secondCandidate.fingerprint, + canonicalResolutions: secondCandidate.canonicalResolutions, + } + const results = await Promise.all([ + store.commitInterruptResolutions(first), + store.commitInterruptResolutions(second), + ]) + expect(results.map((result) => result.status).sort()).toEqual([ + 'committed', + 'conflict', + ]) + const winningRunId = results.find( + (result) => result.status === 'committed', + )?.continuationRunId + expect(winningRunId).toBeDefined() + const replayInput = winningRunId === 'run-continuation-a' ? first : second + await expect(store.commitInterruptResolutions({ + ...replayInput, + continuationRunId: 'ignored-replay-run', + })).resolves.toEqual({ + status: 'replayed', + continuationRunId: winningRunId, + }) + + const recoveredStore = harness.reopen ? await harness.reopen() : store + await expect(recoveredStore.getInterruptRecoveryState({ + threadId: 'thread-cas', + interruptedRunId: 'run-interrupted', + knownGeneration: opened.generation, + })).resolves.toMatchObject({ + state: 'committed', + committed: { continuationRunId: winningRunId }, + }) + }) + + it('marks an uncommitted expired batch expired', async () => { + const harness = await createHarness() + const store = harness.getStore() + if (!store) throw new Error('interrupt store missing') + const opened = await openTwoInterrupts(store) + harness.advanceBy(60_001) + await expect( + store.commitInterruptResolutions(validTwoItemCommit(opened.generation)), + ).resolves.toMatchObject({ + status: 'conflict', + authoritativeState: { state: 'expired' }, + }) + }) +} +``` + +Memory omits `reopen` to make process-local loss explicit. Drizzle, Prisma, and Cloudflare supply it and must recover the committed winner. Every harness supplies `failTransitionOnce`; it may wrap the store at the gateway boundary so production adapters do not expose a test-only option. No conformance case returns early. + +- [ ] **Step 2: Run the red persistence targets** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-persistence test:lib -- --run tests/interrupts.test.ts tests/memory.conformance.test.ts tests/with-persistence.test.ts tests/persistence-validation.test.ts +pnpm --filter @tanstack/ai-persistence test:types +``` + +Expected: FAIL because `openInterruptBatch`, `commitInterruptResolutions`, recovery, generation, and `fingerprintInterruptResolutions` are missing. + +- [ ] **Step 3: Define the store records and capability alias** + +Make `InterruptStore` extend the core gateway and retain read methods only; remove sequential `resolve`/`cancel` as a middleware path: + +```ts +export interface InterruptRecord { + interruptId: string + runId: string + threadId: string + generation: number + status: 'pending' | 'resolved' | 'cancelled' + requestedAt: number + resolvedAt?: number + payload: unknown + binding: InterruptBinding + response?: unknown +} + +export interface InterruptBatchRecord { + threadId: string + interruptedRunId: string + generation: number + expectedInterruptIds: readonly string[] + fingerprint: string + canonicalResolutions: string + resolutions: readonly RunAgentResumeItem[] + continuationRunId: string + committedAt: number +} + +export interface InterruptStore extends InterruptPersistenceGateway { + get(interruptId: string): Promise + list(threadId: string): Promise + listPending(threadId: string): Promise + listByRun(runId: string): Promise + listPendingByRun(runId: string): Promise +} +``` + +In `capabilities.ts`, stop constructing a second interrupt token. Re-export the core `InterruptPersistenceCapability`, `getInterruptPersistence`, and `provideInterruptPersistence`. Keep `InterruptsCapability`, `getInterrupts`, and `provideInterrupts` as deprecated aliases referencing those exact same values until 1.0. + +- [ ] **Step 4: Implement exact-set helpers and consume the core identity** + +Delete persistence-owned canonicalization and FNV. Import `canonicalizeInterruptResolutions` and `cloneAndDeepFreezeJson` from `@tanstack/ai`; persistence owns only exact-set checks: + +```ts +export class InterruptStoreCorruptionError extends Error { + constructor(message: string, options?: ErrorOptions) { + super(message, options) + this.name = 'InterruptStoreCorruptionError' + } +} + +export function hasExactInterruptIds( + expected: readonly string[], + actual: readonly string[], +): boolean { + const left = [...new Set(expected)].sort() + const right = [...new Set(actual)].sort() + return left.length === expected.length && + right.length === actual.length && + left.length === right.length && + left.every((id, index) => id === right[index]) +} + +export function projectInterruptRecovery(input: { + query: InterruptRecoveryQuery + rows: readonly InterruptRecord[] + batch: InterruptBatchRecord | null + now: number + includeResolutions: boolean +}): InterruptRecoveryStateV1 { + const correlation = { + schemaVersion: 1 as const, + threadId: input.query.threadId, + interruptedRunId: input.query.interruptedRunId, + generation: + input.batch?.generation ?? + input.rows[0]?.generation ?? + input.query.knownGeneration, + pendingInterrupts: [] as readonly Interrupt[], + } + if (input.batch) { + return cloneAndDeepFreezeJson({ + ...correlation, + state: 'committed', + committed: { + fingerprint: input.batch.fingerprint, + ...(input.includeResolutions && { + resolutions: input.batch.resolutions, + }), + continuationRunId: input.batch.continuationRunId, + committedAt: new Date(input.batch.committedAt).toISOString(), + }, + }) + } + if (input.rows.length === 0) { + return cloneAndDeepFreezeJson({ ...correlation, state: 'missing' }) + } + const pending = input.rows.filter((row) => row.status === 'pending') + if (pending.length > 0) { + const expired = pending.some( + (row) => + row.binding.expiresAt !== undefined && + Date.parse(row.binding.expiresAt) <= input.now, + ) + return cloneAndDeepFreezeJson({ + ...correlation, + state: expired ? 'expired' : 'pending', + pendingInterrupts: expired + ? [] + : pending.map((row) => row.payload as Interrupt), + }) + } + return cloneAndDeepFreezeJson({ + ...correlation, + state: 'legacy-committed', + }) +} +``` + +Before any store operation, recompute `canonicalizeInterruptResolutions(input.resolutions)` and require both its digest and canonical bytes to equal `input.fingerprint` and `input.canonicalResolutions`. Replay requires stored digest equality **and** stored canonical-byte equality, so a digest alone never proves identity. Memory and every durable adapter call `projectInterruptRecovery` from inside their run lock/transaction with rows and winner read from that same snapshot; its five branches are exercised by the shared conformance suite and the migrated-legacy tests. + +- [ ] **Step 5: Implement the memory critical section and immutable winner** + +Use maps keyed by `interruptedRunId` plus a per-key promise chain. `openInterruptBatch` writes every descriptor or none, stamps generation `1` and schema identity into returned descriptors, and returns an existing identical pending batch idempotently. `commitInterruptResolutions` executes inside the same critical section: + +```ts +return this.withInterruptLock(input.interruptedRunId, async () => { + const candidate = canonicalizeInterruptResolutions(input.resolutions) + if ( + candidate.fingerprint !== input.fingerprint || + candidate.canonicalResolutions !== input.canonicalResolutions + ) { + throw new TypeError('Interrupt batch identity does not match its resolutions.') + } + const winner = this.batches.get(input.interruptedRunId) + if (winner) { + return winner.fingerprint === input.fingerprint && + winner.canonicalResolutions === input.canonicalResolutions + ? { status: 'replayed', continuationRunId: winner.continuationRunId } + : { + status: 'conflict', + authoritativeState: this.recoveryFromWinner(winner), + } + } + const pending = this.pendingForRun(input.interruptedRunId) + if ( + !hasExactInterruptIds( + input.expectedInterruptIds, + pending.map((record) => record.interruptId), + ) || + pending.some((record) => record.generation !== input.expectedGeneration) || + pending.some((record) => + record.binding.expiresAt !== undefined && + Date.parse(record.binding.expiresAt) <= this.clock(), + ) + ) { + return { + status: 'conflict', + authoritativeState: this.recoveryForRun(input), + } + } + const committedAt = this.clock() + const nextInterrupts = new Map(this.interrupts) + for (const resolution of candidate.resolutions) { + const current = nextInterrupts.get(resolution.interruptId) + if ( + !current || + current.runId !== input.interruptedRunId || + current.status !== 'pending' || + current.generation !== input.expectedGeneration + ) { + throw new Error(`Pending interrupt changed: ${resolution.interruptId}`) + } + nextInterrupts.set( + resolution.interruptId, + cloneAndDeepFreezeJson({ + ...current, + status: resolution.status === 'cancelled' ? 'cancelled' : 'resolved', + response: resolution, + resolvedAt: committedAt, + }), + ) + } + const immutableWinner = cloneAndDeepFreezeJson({ + ...input, + resolutions: candidate.resolutions, + canonicalResolutions: candidate.canonicalResolutions, + fingerprint: candidate.fingerprint, + committedAt, + }) + const nextBatches = new Map(this.batches) + nextBatches.set(input.interruptedRunId, immutableWinner) + + // Publish only after every row and winner has been validated and built. + this.interrupts = nextInterrupts + this.batches = nextBatches + return { + status: 'committed', + continuationRunId: input.continuationRunId, + } +}) +``` + +`openInterruptBatch` runs under the same per-run lock: verify descriptor/binding exact IDs and correlations, build a cloned next interrupt map and frozen descriptor snapshot, and publish the map only after every record is built. `getInterruptRecoveryState` reads one locked snapshot and returns `pending`, `committed`, `expired`, `missing`, or `legacy-committed` without exposing mutable store objects. + +- [ ] **Step 6: Reject non-atomic custom stores and export the new helpers** + +Update `defineAIPersistence` validation so an interrupt store lacking all three gateway methods throws `atomic-commit-unsupported`; do not adapt `resolve`/`cancel` sequentially. Re-export new types/helpers and update `createInterruptController` to expose `openInterruptBatch`, `commitInterruptResolutions`, `getInterruptRecoveryState`, and pending reads only. + +- [ ] **Step 7: Run green persistence checks and refactor** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-persistence test:lib -- --run tests/interrupts.test.ts tests/memory.conformance.test.ts tests/with-persistence.test.ts tests/persistence-validation.test.ts +pnpm --filter @tanstack/ai-persistence test:types +``` + +Expected: PASS; the concurrency test has one winner, replay returns its continuation ID, and fault injection leaves every row pending. Refactor duplicated recovery construction into pure helpers and rerun both commands. + +### Task 5: Implement transactional CAS and migrations for Drizzle SQLite + +**Prerequisites:** Task 4. + +**Files:** +- Modify: `packages/ai-persistence-drizzle/src/schema.ts:65-78,115-125` +- Modify: `packages/ai-persistence-drizzle/src/stores.ts:1-38,196-286` +- Modify: `packages/ai-persistence-drizzle/src/index.ts` +- Modify: `packages/ai-persistence-drizzle/src/sqlite.ts` +- Modify: `packages/ai-persistence-drizzle/src/migrations.ts` +- Create: `packages/ai-persistence-drizzle/drizzle/0001_tanstack_ai_interrupt_batches.sql` +- Create: `packages/ai-persistence-drizzle/drizzle/meta/0001_snapshot.json` +- Modify: `packages/ai-persistence-drizzle/drizzle/meta/_journal.json` +- Create: `packages/ai-persistence-drizzle/src/assets/0001_tanstack_ai_interrupt_batches.sql` +- Modify: `packages/ai-persistence-drizzle/tests/drizzle.conformance.test.ts` +- Modify: `packages/ai-persistence-drizzle/tests/migrations.test.ts` +- Modify: `packages/ai-persistence-drizzle/tests/migration-cli.test.ts` +- Modify: `packages/ai-persistence-drizzle/tests/api-types.test-d.ts` + +- [ ] **Step 1: Add failing adapter and nonempty-migration tests** + +Extend the shared conformance invocation and assert historical rows migrate deterministically: + +```ts +it('backfills pending generation and identifies historical commits', () => { + const db = new DatabaseSync(':memory:') + applySqliteMigrations(db, [sqliteMigrations[0]]) + const binding = { + kind: 'generic', + interruptId: 'pending', + interruptedRunId: 'run-p', + generation: 1, + responseSchemaHash: 'sha256:pending', + } + const descriptor = { + id: 'pending', + reason: 'confirmation', + responseSchema: { type: 'object', required: ['approved'] }, + } + const insert = db.prepare(`INSERT INTO interrupts ( + interrupt_id, run_id, thread_id, status, requested_at, + resolved_at, payload_json, response_json + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)`) + insert.run('pending', 'run-p', 'thread', 'pending', 1, null, JSON.stringify(descriptor), null) + insert.run('done', 'run-d', 'thread', 'resolved', 1, 2, '{}', '{}') + applySqliteMigrations(db, [sqliteMigrations[1]]) + + expect(db.prepare('SELECT generation FROM interrupts WHERE interrupt_id = ?').get('pending')).toEqual({ generation: 1 }) + expect(db.prepare('SELECT generation FROM interrupts WHERE interrupt_id = ?').get('done')).toEqual({ generation: 0 }) + expect(db.prepare('SELECT COUNT(*) AS count FROM interrupt_batches').get()).toEqual({ count: 0 }) + const migrated = db.prepare( + 'SELECT status, generation, binding_json FROM interrupts WHERE interrupt_id = ?', + ).get('pending') + expect(migrated).toMatchObject({ status: 'pending', generation: 1 }) + expect(JSON.parse(String(migrated.binding_json))).toEqual({ + ...binding, + responseSchemaHash: expect.stringMatching(/^legacy-json:/), + }) + expect( + db.prepare( + 'SELECT payload_json FROM interrupts WHERE interrupt_id = ? AND status = ?', + ).get('pending', 'pending'), + ).toEqual({ payload_json: JSON.stringify(descriptor) }) +}) + +it('upgrades all legacy pending descriptors atomically', () => { + const db = new DatabaseSync(':memory:') + applySqliteMigrations(db, [sqliteMigrations[0]]) + db.prepare(`INSERT INTO interrupts ( + interrupt_id, run_id, thread_id, status, requested_at, payload_json + ) VALUES (?, ?, ?, ?, ?, ?)`) + .run( + 'legacy-tool', + 'run-legacy', + 'thread', + 'pending', + 1, + JSON.stringify({ + id: 'legacy-tool', + reason: 'approval_required', + toolCallId: 'call-legacy', + responseSchema: { + type: 'object', + properties: { approved: { type: 'boolean' } }, + required: ['approved'], + }, + metadata: { kind: 'approval', toolName: 'transfer', input: { cents: 5 } }, + }), + ) + expect(() => applySqliteMigrations(db, [sqliteMigrations[1]])).not.toThrow() + const rows = db.prepare( + 'SELECT status, generation, binding_json FROM interrupts WHERE run_id = ?', + ).all('run-legacy') + expect(rows).toHaveLength(1) + expect(rows[0]).toMatchObject({ status: 'pending', generation: 1 }) + expect(JSON.parse(String(rows[0]?.binding_json))).toMatchObject({ + kind: 'generic', + interruptId: 'legacy-tool', + interruptedRunId: 'run-legacy', + generation: 1, + }) +}) +``` + +- [ ] **Step 2: Run the red Drizzle targets** + +Run: + +```powershell +pnpm --filter @tanstack/ai-persistence-drizzle test:lib -- --run tests/drizzle.conformance.test.ts tests/migrations.test.ts tests/migration-cli.test.ts +pnpm --filter @tanstack/ai-persistence-drizzle test:types +``` + +Expected: FAIL because the schema lacks generation/binding/batch columns and the store lacks atomic gateway methods. + +- [ ] **Step 3: Add schema and an explicit transaction executor** + +Add `generation`, `bindingJson`, and `responseSchemaHash` to `interrupts`; add an `interruptBatches` table keyed by interrupted run with a unique continuation ID. Do not claim atomicity from the schema-agnostic `DrizzleDb` alone: + +```ts +export interface DrizzleTransactionExecutor { + transaction(work: (tx: DrizzleDb) => Promise): Promise +} + +export function createInterruptStore( + db: DrizzleDb, + executor: DrizzleTransactionExecutor, +): InterruptStore + +export function drizzlePersistence( + db: DrizzleDb, + options: { interruptTransactions: DrizzleTransactionExecutor }, +) +``` + +The Node factory must serialize asynchronous transaction callbacks per connection: + +```ts +export function createSerializedSqliteExecutor( + database: DatabaseSync, + db: DrizzleDb, +): DrizzleTransactionExecutor { + let tail: Promise = Promise.resolve() + return { + transaction(work: (tx: DrizzleDb) => Promise): Promise { + const run = tail.then(async () => { + database.exec('BEGIN IMMEDIATE') + try { + const value = await work(db) + database.exec('COMMIT') + return value + } catch (error) { + database.exec('ROLLBACK') + throw error + } + }) + tail = run.then(() => undefined, () => undefined) + return run + }, + } +} +``` + +`sqlitePersistence` constructs exactly one executor per `DatabaseSync`. `drizzlePersistence` requires `interruptTransactions` whenever interrupts are enabled and throws `atomic-commit-unsupported` during construction if it is absent. + +- [ ] **Step 4: Implement the transaction and exact replay** + +Implement `openInterruptBatch`, `commitInterruptResolutions`, and `getInterruptRecoveryState` through this executor. The commit body is: + +```ts +return executor.transaction(async (tx) => { + const existing = await tx.select().from(interruptBatches) + .where(eq(interruptBatches.interruptedRunId, input.interruptedRunId)).get() + if (existing) { + if ( + existing.fingerprint === input.fingerprint && + existing.resolutionsJson === input.canonicalResolutions + ) { + return { status: 'replayed', continuationRunId: existing.continuationRunId } + } + return { status: 'conflict', authoritativeState: recoveryFromBatch(existing) } + } + const pending = await tx.select().from(interrupts).where(and( + eq(interrupts.runId, input.interruptedRunId), + eq(interrupts.status, 'pending'), + )).all() + const now = clock() + if ( + !hasExactInterruptIds(input.expectedInterruptIds, pending.map((row) => row.interruptId)) || + pending.some((row) => row.generation !== input.expectedGeneration) || + pending.some((row) => isExpiredBindingJson(row.bindingJson, now)) + ) { + return { status: 'conflict', authoritativeState: recoveryFromRows(input, pending, now) } + } + await tx.insert(interruptBatches).values({ + interruptedRunId: input.interruptedRunId, + threadId: input.threadId, + generation: input.expectedGeneration, + expectedInterruptIdsJson: JSON.stringify([...input.expectedInterruptIds].sort()), + fingerprint: input.fingerprint, + resolutionsJson: input.canonicalResolutions, + continuationRunId: input.continuationRunId, + committedAt: now, + }).run() + for (const resolution of input.resolutions) { + const result = await tx.update(interrupts).set({ + status: resolution.status === 'cancelled' ? 'cancelled' : 'resolved', + responseJson: resolution, + resolvedAt: now, + }).where(and( + eq(interrupts.interruptId, resolution.interruptId), + eq(interrupts.runId, input.interruptedRunId), + eq(interrupts.status, 'pending'), + eq(interrupts.generation, input.expectedGeneration), + )).run() + if (result.changes !== 1) throw new Error('interrupt CAS row changed') + } + return { status: 'committed', continuationRunId: input.continuationRunId } +}) +``` + +`openInterruptBatch` and recovery use the same executor and the shared projection: + +```ts +async function openInterruptBatch( + input: OpenInterruptBatchInput, +): Promise<{ generation: number; descriptors: readonly Interrupt[] }> { + if (!hasExactInterruptIds( + input.descriptors.map((descriptor) => descriptor.id), + input.bindings.map((binding) => binding.interruptId), + )) { + throw new TypeError('Interrupt descriptors and bindings must have exact IDs.') + } + return executor.transaction(async (tx) => { + const existing = await tx.select().from(interrupts).where( + eq(interrupts.runId, input.interruptedRunId), + ).all() + if (existing.length > 0) { + const existingIds = existing.map((row) => row.interruptId) + if (!hasExactInterruptIds(existingIds, input.descriptors.map((item) => item.id))) { + throw new Error('Interrupt batch already opened with a different set.') + } + return { + generation: existing[0]?.generation ?? 1, + descriptors: cloneAndDeepFreezeJson( + existing.map((row) => JSON.parse(row.payloadJson) as Interrupt), + ), + } + } + + const generation = 1 + const bindings = new Map(input.bindings.map((binding) => [ + binding.interruptId, + cloneAndDeepFreezeJson({ + ...binding, + interruptedRunId: input.interruptedRunId, + generation, + } satisfies InterruptBinding), + ])) + const descriptors = input.descriptors.map((descriptor) => { + const binding = bindings.get(descriptor.id) + if (!binding) throw new Error(`Missing binding for ${descriptor.id}`) + return cloneAndDeepFreezeJson({ + ...descriptor, + metadata: { ...descriptor.metadata, 'tanstack:binding': binding }, + }) + }) + await tx.insert(interrupts).values(descriptors.map((descriptor) => { + const binding = bindings.get(descriptor.id) + if (!binding) throw new Error(`Missing binding for ${descriptor.id}`) + return { + interruptId: descriptor.id, + runId: input.interruptedRunId, + threadId: input.threadId, + generation, + status: 'pending', + requestedAt: clock(), + payloadJson: JSON.stringify(descriptor), + bindingJson: JSON.stringify(binding), + responseSchemaHash: binding.responseSchemaHash, + } + })).run() + return { generation, descriptors: cloneAndDeepFreezeJson(descriptors) } + }) +} + +async function getInterruptRecoveryState( + query: InterruptRecoveryQuery, +): Promise { + return executor.transaction(async (tx) => { + const batchRow = await tx.select().from(interruptBatches).where( + eq(interruptBatches.interruptedRunId, query.interruptedRunId), + ).get() + const rows = await tx.select().from(interrupts).where( + eq(interrupts.runId, query.interruptedRunId), + ).orderBy(interrupts.interruptId).all() + return projectInterruptRecovery({ + query, + batch: batchRow ? decodeInterruptBatchRow(batchRow) : null, + rows: rows.map(decodeInterruptRow), + now: clock(), + includeResolutions: true, + }) + }) +} +``` + +Define the decoders beside the store; these names are used by the preceding methods: + +```ts +function parseJsonRecord(value: string, label: string): Record { + try { + const parsed: unknown = JSON.parse(value) + if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) { + throw new TypeError(`${label} must contain a JSON object.`) + } + return parsed as Record + } catch (cause) { + throw new InterruptStoreCorruptionError(`Invalid ${label}.`, { cause }) + } +} + +function decodeInterruptRow( + row: typeof interrupts.$inferSelect, +): InterruptRecord { + const payload = parseJsonRecord(row.payloadJson, 'interrupt payload') + const binding = parseJsonRecord(row.bindingJson, 'interrupt binding') + if ( + binding['interruptId'] !== row.interruptId || + binding['interruptedRunId'] !== row.runId || + binding['generation'] !== row.generation + ) { + throw new InterruptStoreCorruptionError('Interrupt binding correlation mismatch.') + } + return cloneAndDeepFreezeJson({ + interruptId: row.interruptId, + runId: row.runId, + threadId: row.threadId, + generation: row.generation, + status: row.status, + requestedAt: row.requestedAt, + ...(row.resolvedAt !== null && { resolvedAt: row.resolvedAt }), + payload, + binding, + ...(row.responseJson !== null && { + response: JSON.parse(row.responseJson) as unknown, + }), + }) as InterruptRecord +} + +function decodeInterruptBatchRow( + row: typeof interruptBatches.$inferSelect, +): InterruptBatchRecord { + const expectedInterruptIds: unknown = JSON.parse(row.expectedInterruptIdsJson) + const resolutions: unknown = JSON.parse(row.resolutionsJson) + if (!Array.isArray(expectedInterruptIds) || !Array.isArray(resolutions)) { + throw new InterruptStoreCorruptionError('Invalid interrupt batch arrays.') + } + const resolutionIds = resolutions.map((entry) => { + if (!entry || typeof entry !== 'object' || !('interruptId' in entry)) { + throw new InterruptStoreCorruptionError('Invalid interrupt batch resolution.') + } + return String(entry.interruptId) + }) + if (!hasExactInterruptIds(expectedInterruptIds.map(String), resolutionIds)) { + throw new InterruptStoreCorruptionError('Interrupt batch ID set mismatch.') + } + return cloneAndDeepFreezeJson({ + threadId: row.threadId, + interruptedRunId: row.interruptedRunId, + generation: row.generation, + expectedInterruptIds: expectedInterruptIds.map(String), + fingerprint: row.fingerprint, + canonicalResolutions: row.resolutionsJson, + resolutions, + continuationRunId: row.continuationRunId, + committedAt: row.committedAt, + }) as InterruptBatchRecord +} +``` + +Direct tests assert `pending`, `committed`, `expired`, `missing`, and resolved-without-batch `legacy-committed`; expiry is evaluated inside the executor transaction. + +- [ ] **Step 5: Generate and package the numbered migration** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-persistence-drizzle db:generate -- --name tanstack_ai_interrupt_batches +``` + +Expected: Drizzle creates the `0001` SQL/meta entries. Ensure the SQL represented by both the generated file and packaged asset contains: + +```sql +ALTER TABLE `interrupts` ADD `generation` integer NOT NULL DEFAULT 1; +ALTER TABLE `interrupts` ADD `binding_json` text; +ALTER TABLE `interrupts` ADD `response_schema_hash` text NOT NULL DEFAULT ''; +UPDATE `interrupts` +SET `binding_json` = json_object( + 'kind', 'generic', + 'interruptId', `interrupt_id`, + 'interruptedRunId', `run_id`, + 'generation', 1, + 'responseSchemaHash', + 'legacy-json:' || COALESCE(json(json_extract(`payload_json`, '$.responseSchema')), 'null'), + 'expiresAt', json_extract(`payload_json`, '$.expiresAt') + ), + `response_schema_hash` = + 'legacy-json:' || COALESCE(json(json_extract(`payload_json`, '$.responseSchema')), 'null') +WHERE `status` = 'pending'; +UPDATE `interrupts` SET `generation` = 0 WHERE `status` <> 'pending'; +CREATE TABLE `interrupt_batches` ( + `interrupted_run_id` text PRIMARY KEY NOT NULL, + `thread_id` text NOT NULL, + `generation` integer NOT NULL, + `expected_interrupt_ids_json` text NOT NULL, + `fingerprint` text NOT NULL, + `resolutions_json` text NOT NULL, + `continuation_run_id` text NOT NULL UNIQUE, + `committed_at` integer NOT NULL +); +CREATE INDEX `interrupts_run_status_generation_idx` ON `interrupts` (`run_id`,`status`,`generation`); +CREATE INDEX `interrupts_thread_status_idx` ON `interrupts` (`thread_id`,`status`); +``` + +In `migrations.test.ts`, read the generated SQL, packaged asset, `_journal.json`, and snapshot. Assert: journal entry tag is `0001_tanstack_ai_interrupt_batches`; snapshot contains `interrupt_batches` and both indexes; packaged SQL bytes equal generated SQL bytes; and `sqliteMigrations[1]` has the same ID, filename, and SQL. In `migration-cli.test.ts`, assert the CLI writes `0000` then `0001` with those exact bytes. Do not rename only the SQL; regenerate journal and snapshot together. + +- [ ] **Step 6: Run green adapter checks and refactor** + +Run: + +```powershell +pnpm --filter @tanstack/ai-persistence-drizzle test:lib -- --run tests/drizzle.conformance.test.ts tests/migrations.test.ts tests/migration-cli.test.ts +pnpm --filter @tanstack/ai-persistence-drizzle test:types +``` + +Expected: PASS; migration reruns are idempotent, CLI copies both files, and the conformance race has one winner. Extract row-to-recovery mapping, rerun both commands, and confirm no transaction accepts a partial expected set. + +### Task 6: Implement Prisma CAS and keep all schema copies identical + +**Prerequisites:** Task 4. May run after Task 5 once shared types settle. + +**Files:** +- Modify: `packages/ai-persistence-prisma/src/stores.ts:200-300` +- Modify: `packages/ai-persistence-prisma/prisma/schema.prisma` +- Modify: `packages/ai-persistence-prisma/prisma/tanstack-ai.prisma` +- Modify: `packages/ai-persistence-prisma/src/assets/tanstack-ai.prisma` +- Modify: `packages/ai-persistence-prisma/tests/prisma.conformance.test.ts` +- Modify: `packages/ai-persistence-prisma/tests/models.test.ts` +- Modify: `packages/ai-persistence-prisma/tests/models-cli.test.ts` +- Modify: `packages/ai-persistence-prisma/tests/api-types.test-d.ts` + +- [ ] **Step 1: Add failing schema-copy and transaction tests** + +Add assertions that all model fragments contain the batch model and are byte-identical where they are intended to be copied: + +```ts +expect(prismaModels).toContain('model InterruptBatch') +expect(prismaModels).toContain('generation Int @default(1)') +expect(prismaModels).toContain('continuationRunId String @unique') +expect(copiedModels).toBe(prismaModels) +``` + +Run shared conformance with two Prisma clients submitting different complete batches concurrently. + +- [ ] **Step 2: Run the red Prisma targets** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-persistence-prisma db:generate +pnpm --filter @tanstack/ai-persistence-prisma test:lib -- --run tests/prisma.conformance.test.ts tests/models.test.ts tests/models-cli.test.ts +pnpm --filter @tanstack/ai-persistence-prisma test:types +``` + +Expected: tests FAIL because `InterruptBatch`, generation, and gateway methods are absent. + +- [ ] **Step 3: Update the three Prisma schema sources** + +Add the same fields to `Interrupt` in all copies and this provider-neutral model: + +```prisma +model InterruptBatch { + interruptedRunId String @id @map("interrupted_run_id") + threadId String @map("thread_id") + generation Int + expectedInterruptIdsJson String @map("expected_interrupt_ids_json") + fingerprint String + resolutionsJson String @map("resolutions_json") + continuationRunId String @unique @map("continuation_run_id") + committedAt BigInt @map("committed_at") + + @@map("interrupt_batches") +} +``` + +Add `generation Int @default(1)`, `bindingJson String?`, `responseSchemaHash String @default("")`, plus run/status/generation and thread/status indexes to `Interrupt`. + +- [ ] **Step 4: Implement `$transaction` CAS** + +Use an interactive transaction with the exact-set decision inside it: + +```ts +return prisma.$transaction(async (tx) => { + const winner = await tx.interruptBatch.findUnique({ + where: { interruptedRunId: input.interruptedRunId }, + }) + if (winner) { + if ( + winner.fingerprint === input.fingerprint && + winner.resolutionsJson === input.canonicalResolutions + ) { + return { status: 'replayed', continuationRunId: winner.continuationRunId } + } + return { status: 'conflict', authoritativeState: recoveryFromPrismaBatch(winner) } + } + + const pending = await tx.interrupt.findMany({ + where: { + runId: input.interruptedRunId, + status: 'pending', + generation: input.expectedGeneration, + }, + orderBy: { interruptId: 'asc' }, + }) + const now = clock() + if ( + !hasExactInterruptIds(input.expectedInterruptIds, pending.map((row) => row.interruptId)) || + pending.some((row) => bindingExpired(row.bindingJson, now)) + ) { + return { status: 'conflict', authoritativeState: recoveryFromPrismaRows(input, pending, now) } + } + await tx.interruptBatch.create({ + data: { + interruptedRunId: input.interruptedRunId, + threadId: input.threadId, + generation: input.expectedGeneration, + expectedInterruptIdsJson: JSON.stringify([...input.expectedInterruptIds].sort()), + fingerprint: input.fingerprint, + resolutionsJson: input.canonicalResolutions, + continuationRunId: input.continuationRunId, + committedAt: BigInt(now), + }, + }) + for (const resolution of input.resolutions) { + const updated = await tx.interrupt.updateMany({ + where: { + interruptId: resolution.interruptId, + runId: input.interruptedRunId, + status: 'pending', + generation: input.expectedGeneration, + }, + data: { + status: resolution.status === 'cancelled' ? 'cancelled' : 'resolved', + responseJson: JSON.stringify(resolution), + resolvedAt: BigInt(now), + }, + }) + if (updated.count !== 1) throw new Error('interrupt CAS row changed') + } + return { + status: 'committed', + continuationRunId: input.continuationRunId, + } +}) +``` + +Add these two gateway methods beside the commit method: + +```ts +async function openInterruptBatch( + input: OpenInterruptBatchInput, +): Promise<{ generation: number; descriptors: readonly Interrupt[] }> { + if (!hasExactInterruptIds( + input.descriptors.map((descriptor) => descriptor.id), + input.bindings.map((binding) => binding.interruptId), + )) { + throw new TypeError('Interrupt descriptors and bindings must have exact IDs.') + } + return prisma.$transaction(async (tx) => { + const existing = await tx.interrupt.findMany({ + where: { runId: input.interruptedRunId }, + orderBy: { interruptId: 'asc' }, + }) + if (existing.length > 0) { + if (!hasExactInterruptIds( + existing.map((row) => row.interruptId), + input.descriptors.map((descriptor) => descriptor.id), + )) { + throw new Error('Interrupt batch already opened with a different set.') + } + return { + generation: existing[0]?.generation ?? 1, + descriptors: cloneAndDeepFreezeJson( + existing.map((row) => JSON.parse(row.payloadJson) as Interrupt), + ), + } + } + + const generation = 1 + const unopened = new Map( + input.bindings.map((binding) => [binding.interruptId, binding]), + ) + const descriptors = input.descriptors.map((descriptor) => { + const binding = unopened.get(descriptor.id) + if (!binding) throw new Error(`Missing binding for ${descriptor.id}`) + const authoritativeBinding = { + ...binding, + interruptedRunId: input.interruptedRunId, + generation, + } satisfies InterruptBinding + return { + descriptor: cloneAndDeepFreezeJson({ + ...descriptor, + metadata: { + ...descriptor.metadata, + 'tanstack:binding': authoritativeBinding, + }, + }), + binding: authoritativeBinding, + } + }) + await tx.interrupt.createMany({ + data: descriptors.map(({ descriptor, binding }) => ({ + interruptId: descriptor.id, + runId: input.interruptedRunId, + threadId: input.threadId, + generation, + status: 'pending', + requestedAt: BigInt(clock()), + payloadJson: JSON.stringify(descriptor), + bindingJson: JSON.stringify(binding), + responseSchemaHash: binding.responseSchemaHash, + })), + }) + return { + generation, + descriptors: cloneAndDeepFreezeJson( + descriptors.map(({ descriptor }) => descriptor), + ), + } + }) +} + +async function getInterruptRecoveryState( + query: InterruptRecoveryQuery, +): Promise { + return prisma.$transaction(async (tx) => { + const batch = await tx.interruptBatch.findUnique({ + where: { interruptedRunId: query.interruptedRunId }, + }) + const rawRows = await tx.interrupt.findMany({ + where: { runId: query.interruptedRunId }, + orderBy: { interruptId: 'asc' }, + }) + const rows = [] + for (const row of rawRows) { + if (row.status === 'pending' && row.bindingJson === null) { + const descriptor: Interrupt = JSON.parse(row.payloadJson) + const responseSchemaHash = `legacy-json:${canonicalInterruptJson( + descriptor.responseSchema ?? null, + )}` + const binding = cloneAndDeepFreezeJson({ + kind: 'generic', + interruptId: row.interruptId, + interruptedRunId: row.runId, + generation: row.generation, + responseSchemaHash, + ...(descriptor.expiresAt !== undefined && { + expiresAt: descriptor.expiresAt, + }), + } satisfies InterruptBinding) + await tx.interrupt.update({ + where: { interruptId: row.interruptId }, + data: { + bindingJson: JSON.stringify(binding), + responseSchemaHash, + }, + }) + rows.push(decodePrismaInterruptRow({ ...row, bindingJson: JSON.stringify(binding) })) + } else { + rows.push(decodePrismaInterruptRow(row)) + } + } + return projectInterruptRecovery({ + query, + batch: batch ? decodePrismaBatchRow(batch) : null, + rows, + now: clock(), + includeResolutions: true, + }) + }) +} +``` + +Define the Prisma decoders used above: + +```ts +type PrismaInterruptRow = { + interruptId: string + runId: string + threadId: string + generation: number + status: string + requestedAt: bigint + resolvedAt: bigint | null + payloadJson: string + bindingJson: string | null + responseJson: string | null +} + +function decodePrismaInterruptRow(row: PrismaInterruptRow): InterruptRecord { + if (row.bindingJson === null) { + throw new InterruptStoreCorruptionError('Pending interrupt has no binding.') + } + const payload: unknown = JSON.parse(row.payloadJson) + const binding: unknown = JSON.parse(row.bindingJson) + if ( + !binding || + typeof binding !== 'object' || + !('interruptId' in binding) || + binding.interruptId !== row.interruptId || + !('interruptedRunId' in binding) || + binding.interruptedRunId !== row.runId || + !('generation' in binding) || + binding.generation !== row.generation + ) { + throw new InterruptStoreCorruptionError('Interrupt binding correlation mismatch.') + } + if ( + row.status !== 'pending' && + row.status !== 'resolved' && + row.status !== 'cancelled' + ) { + throw new InterruptStoreCorruptionError('Invalid interrupt status.') + } + return cloneAndDeepFreezeJson({ + interruptId: row.interruptId, + runId: row.runId, + threadId: row.threadId, + generation: row.generation, + status: row.status, + requestedAt: Number(row.requestedAt), + ...(row.resolvedAt !== null && { resolvedAt: Number(row.resolvedAt) }), + payload, + binding, + ...(row.responseJson !== null && { + response: JSON.parse(row.responseJson) as unknown, + }), + }) as InterruptRecord +} + +type PrismaInterruptBatchRow = { + interruptedRunId: string + threadId: string + generation: number + expectedInterruptIdsJson: string + fingerprint: string + resolutionsJson: string + continuationRunId: string + committedAt: bigint +} + +function decodePrismaBatchRow( + row: PrismaInterruptBatchRow, +): InterruptBatchRecord { + const expected: unknown = JSON.parse(row.expectedInterruptIdsJson) + const resolutions: unknown = JSON.parse(row.resolutionsJson) + if (!Array.isArray(expected) || !Array.isArray(resolutions)) { + throw new InterruptStoreCorruptionError('Invalid interrupt batch arrays.') + } + const resolutionIds = resolutions.map((entry) => { + if (!entry || typeof entry !== 'object' || !('interruptId' in entry)) { + throw new InterruptStoreCorruptionError('Invalid interrupt resolution.') + } + return String(entry.interruptId) + }) + if (!hasExactInterruptIds(expected.map(String), resolutionIds)) { + throw new InterruptStoreCorruptionError('Interrupt batch ID set mismatch.') + } + return cloneAndDeepFreezeJson({ + interruptedRunId: row.interruptedRunId, + threadId: row.threadId, + generation: row.generation, + expectedInterruptIds: expected.map(String), + fingerprint: row.fingerprint, + canonicalResolutions: row.resolutionsJson, + resolutions, + continuationRunId: row.continuationRunId, + committedAt: Number(row.committedAt), + }) as InterruptBatchRecord +} +``` + +The in-transaction legacy branch is the Prisma data migration path: it reconstructs a safe generic authoritative binding from the stored descriptor and persists it before returning `pending`; it never guesses a current tool binding. + +Catch a `P2002` unique race outside the commit transaction, reload the authoritative batch, and compare both fingerprint and canonical `resolutionsJson`; replay always returns the stored continuation ID. Invoke the shared conformance kit with a controllable clock, a faulting gateway wrapper, and a `reopen` that constructs a new Prisma client. Add a direct test that inserts a pending row with `bindingJson: null`, calls recovery, asserts the row was atomically upgraded to `kind: 'generic'`, and can then commit; add expiry-inside-transaction and resolved-without-batch `legacy-committed` tests. All three Prisma schema copies use the field-level `String @unique @map("continuation_run_id")` form shown above; no test or copy uses `@@unique([continuationRunId])`. + +- [ ] **Step 5: Generate the client and run green checks** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-persistence-prisma db:generate +pnpm --filter @tanstack/ai-persistence-prisma test:lib -- --run tests/prisma.conformance.test.ts tests/models.test.ts tests/models-cli.test.ts +pnpm --filter @tanstack/ai-persistence-prisma test:types +``` + +Expected: PASS; copied model output contains the new model, and concurrent writers produce one commit and one conflict. + +- [ ] **Step 6: Refactor and recheck** + +Centralize JSON serialization/deserialization for IDs, resolutions, and bindings, rerun the same commands, and verify docs in Task 17 instruct applications to generate and deploy their own Prisma migration before using the new package. + +### Task 7: Implement an atomic Cloudflare D1 conditional batch + +**Prerequisites:** Tasks 4 and 5 for shared schema/migration shape. + +**Files:** +- Modify: `packages/ai-persistence-cloudflare/src/d1.ts` +- Modify: `packages/ai-persistence-cloudflare/src/migrations.ts` +- Create: `packages/ai-persistence-cloudflare/migrations/0001_tanstack_ai_interrupt_batches.sql` +- Create: `packages/ai-persistence-cloudflare/src/assets/0001_tanstack_ai_interrupt_batches.sql` +- Modify: `packages/ai-persistence-cloudflare/tests/runtime.conformance.test.ts` +- Modify: `packages/ai-persistence-cloudflare/tests/migrations.test.ts` +- Modify: `packages/ai-persistence-cloudflare/tests/migration-cli.test.ts` +- Modify: `packages/ai-persistence-cloudflare/tests/api-types.test-d.ts` + +- [ ] **Step 1: Add failing D1 race, rollback, and migration tests** + +Invoke `runInterruptStoreConformance` with `createD1Stores(database).interrupts`, an injected clock, and a reopen callback that reuses the Miniflare database through a new store instance. Add this D1-only rollback assertion: + +```ts +const store = createD1Stores(database, { + interruptTestHook: { failStatementForInterruptId: 'int-b' }, +}).interrupts +if (!store) throw new Error('interrupt store missing') +await openTwoInterrupts(store) +await expect(store.commitInterruptResolutions(validTwoItemCommit())).rejects.toThrow( + /forced D1 statement failure/, +) +expect(await database.prepare('SELECT COUNT(*) AS count FROM interrupt_batches').first()).toEqual({ count: 0 }) +expect(await database.prepare("SELECT COUNT(*) AS count FROM interrupts WHERE status = 'pending'").first()).toEqual({ count: 2 }) +``` + +The fixture functions `openTwoInterrupts` and `validTwoItemCommit` are exported by the shared conformance testkit created in Task 4; they contain the literal two descriptors/bindings and core-canonicalized commit input. Migration tests assert `0001` is bundled and copied after `0000`. + +- [ ] **Step 2: Run the red Cloudflare targets** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-persistence-cloudflare test:lib -- --run tests/runtime.conformance.test.ts tests/migrations.test.ts tests/migration-cli.test.ts +pnpm --filter @tanstack/ai-persistence-cloudflare test:types +``` + +Expected: FAIL because `createD1Stores` delegates interrupts to a Drizzle store without an atomic executor and migration `0001` is missing. + +- [ ] **Step 3: Add and register the D1 migration** + +Both D1 migration locations contain these statements, including the same authoritative generic binding reconstruction for PR-head pending descriptors: + +```sql +ALTER TABLE `interrupts` ADD `generation` integer NOT NULL DEFAULT 1; +ALTER TABLE `interrupts` ADD `binding_json` text; +ALTER TABLE `interrupts` ADD `response_schema_hash` text NOT NULL DEFAULT ''; +UPDATE `interrupts` +SET `binding_json` = json_object( + 'kind', 'generic', + 'interruptId', `interrupt_id`, + 'interruptedRunId', `run_id`, + 'generation', 1, + 'responseSchemaHash', + 'legacy-json:' || COALESCE(json(json_extract(`payload_json`, '$.responseSchema')), 'null'), + 'expiresAt', json_extract(`payload_json`, '$.expiresAt') + ), + `response_schema_hash` = + 'legacy-json:' || COALESCE(json(json_extract(`payload_json`, '$.responseSchema')), 'null') +WHERE `status` = 'pending'; +UPDATE `interrupts` SET `generation` = 0 WHERE `status` <> 'pending'; +CREATE TABLE `interrupt_batches` ( + `interrupted_run_id` text PRIMARY KEY NOT NULL, + `thread_id` text NOT NULL, + `generation` integer NOT NULL, + `expected_interrupt_ids_json` text NOT NULL, + `fingerprint` text NOT NULL, + `resolutions_json` text NOT NULL, + `continuation_run_id` text NOT NULL UNIQUE, + `committed_at` integer NOT NULL +); +CREATE INDEX `interrupts_run_status_generation_idx` + ON `interrupts` (`run_id`,`status`,`generation`); +CREATE INDEX `interrupts_thread_status_idx` + ON `interrupts` (`thread_id`,`status`); +``` + +The D1 migration test first inserts PR-head `approval_required` and `client_tool_input` descriptors with no reserved binding, applies `0001`, and asserts both remain `pending` with generation `1`, a `kind: 'generic'` binding correlated to the stored row/run, and a `legacy-json:` response-schema identity. Apply the migration in one D1 migration transaction and assert the table, columns, indexes, and both upgraded rows appear together. Import the raw asset and append: + +```ts +{ + id: '0001_tanstack_ai_interrupt_batches', + filename: '0001_tanstack_ai_interrupt_batches.sql', + sql: interruptBatchesMigrationSql, +} +``` + +Keep `src/r2.ts` and R2 migrations unchanged. + +- [ ] **Step 4: Override the D1 interrupt store with one conditional batch** + +Implement `createD1InterruptStore(d1, clock)` and return it from `createD1Stores`. The open and recovery methods are: + +```ts +async function openInterruptBatch( + input: OpenInterruptBatchInput, +): Promise<{ generation: number; descriptors: readonly Interrupt[] }> { + const descriptorIds = input.descriptors.map((descriptor) => descriptor.id) + const bindingIds = input.bindings.map((binding) => binding.interruptId) + if (!hasExactInterruptIds(descriptorIds, bindingIds)) { + throw new TypeError('Interrupt descriptors and bindings must have exact IDs.') + } + const existing = await d1.prepare( + 'SELECT * FROM interrupts WHERE run_id = ? ORDER BY interrupt_id', + ).bind(input.interruptedRunId).all() + if (existing.results.length > 0) { + if (!hasExactInterruptIds( + existing.results.map((row) => row.interrupt_id), + descriptorIds, + )) { + throw new Error('Interrupt batch already opened with a different set.') + } + return { + generation: existing.results[0]?.generation ?? 1, + descriptors: cloneAndDeepFreezeJson( + existing.results.map((row) => JSON.parse(row.payload_json) as Interrupt), + ), + } + } + + const generation = 1 + const bindings = new Map(input.bindings.map((binding) => [ + binding.interruptId, + cloneAndDeepFreezeJson({ + ...binding, + interruptedRunId: input.interruptedRunId, + generation, + } satisfies InterruptBinding), + ])) + const descriptors = input.descriptors.map((descriptor) => { + const binding = bindings.get(descriptor.id) + if (!binding) throw new Error(`Missing binding for ${descriptor.id}`) + return cloneAndDeepFreezeJson({ + ...descriptor, + metadata: { ...descriptor.metadata, 'tanstack:binding': binding }, + }) + }) + const results = await d1.batch(descriptors.map((descriptor) => { + const binding = bindings.get(descriptor.id) + if (!binding) throw new Error(`Missing binding for ${descriptor.id}`) + return d1.prepare(`INSERT INTO interrupts ( + interrupt_id, run_id, thread_id, generation, status, requested_at, + payload_json, binding_json, response_schema_hash + ) VALUES (?, ?, ?, ?, 'pending', ?, ?, ?, ?)`).bind( + descriptor.id, + input.interruptedRunId, + input.threadId, + generation, + clock(), + JSON.stringify(descriptor), + JSON.stringify(binding), + binding.responseSchemaHash, + ) + })) + if (results.some((result) => result.meta.changes !== 1)) { + throw new Error('D1 did not open every interrupt row.') + } + return { generation, descriptors: cloneAndDeepFreezeJson(descriptors) } +} + +async function getInterruptRecoveryState( + query: InterruptRecoveryQuery, +): Promise { + const [batchResult, rowsResult] = await d1.batch([ + d1.prepare( + 'SELECT * FROM interrupt_batches WHERE interrupted_run_id = ?', + ).bind(query.interruptedRunId), + d1.prepare( + 'SELECT * FROM interrupts WHERE run_id = ? ORDER BY interrupt_id', + ).bind(query.interruptedRunId), + ]) + const batchRow = batchResult.results[0] as D1InterruptBatchRow | undefined + const rows = rowsResult.results as D1InterruptRow[] + return projectInterruptRecovery({ + query, + batch: batchRow ? decodeD1BatchRow(batchRow) : null, + rows: rows.map(decodeD1InterruptRow), + now: clock(), + includeResolutions: true, + }) +} +``` + +Add the row contracts and decoders used by those methods: + +```ts +interface D1InterruptRow { + interrupt_id: string + run_id: string + thread_id: string + generation: number + status: 'pending' | 'resolved' | 'cancelled' + requested_at: number + resolved_at: number | null + payload_json: string + binding_json: string + response_json: string | null +} + +interface D1InterruptBatchRow { + interrupted_run_id: string + thread_id: string + generation: number + expected_interrupt_ids_json: string + fingerprint: string + resolutions_json: string + continuation_run_id: string + committed_at: number +} + +function decodeD1InterruptRow(row: D1InterruptRow): InterruptRecord { + const payload: unknown = JSON.parse(row.payload_json) + const binding: unknown = JSON.parse(row.binding_json) + if ( + !binding || + typeof binding !== 'object' || + !('interruptId' in binding) || + binding.interruptId !== row.interrupt_id || + !('interruptedRunId' in binding) || + binding.interruptedRunId !== row.run_id || + !('generation' in binding) || + binding.generation !== row.generation + ) { + throw new InterruptStoreCorruptionError('Interrupt binding correlation mismatch.') + } + return cloneAndDeepFreezeJson({ + interruptId: row.interrupt_id, + runId: row.run_id, + threadId: row.thread_id, + generation: row.generation, + status: row.status, + requestedAt: row.requested_at, + ...(row.resolved_at !== null && { resolvedAt: row.resolved_at }), + payload, + binding, + ...(row.response_json !== null && { + response: JSON.parse(row.response_json) as unknown, + }), + }) as InterruptRecord +} + +function decodeD1BatchRow(row: D1InterruptBatchRow): InterruptBatchRecord { + const expected: unknown = JSON.parse(row.expected_interrupt_ids_json) + const resolutions: unknown = JSON.parse(row.resolutions_json) + if (!Array.isArray(expected) || !Array.isArray(resolutions)) { + throw new InterruptStoreCorruptionError('Invalid interrupt batch arrays.') + } + const resolutionIds = resolutions.map((entry) => { + if (!entry || typeof entry !== 'object' || !('interruptId' in entry)) { + throw new InterruptStoreCorruptionError('Invalid interrupt resolution.') + } + return String(entry.interruptId) + }) + if (!hasExactInterruptIds(expected.map(String), resolutionIds)) { + throw new InterruptStoreCorruptionError('Interrupt batch ID set mismatch.') + } + return cloneAndDeepFreezeJson({ + interruptedRunId: row.interrupted_run_id, + threadId: row.thread_id, + generation: row.generation, + expectedInterruptIds: expected.map(String), + fingerprint: row.fingerprint, + canonicalResolutions: row.resolutions_json, + resolutions, + continuationRunId: row.continuation_run_id, + committedAt: row.committed_at, + }) as InterruptBatchRecord +} +``` + +Recovery reads winner and rows in one D1 batch snapshot. Commit builds the dynamic `IN` list from the already unique expected IDs and submits this insert followed by one gated update per resolution: + +```sql +INSERT INTO interrupt_batches ( + interrupted_run_id, thread_id, generation, + expected_interrupt_ids_json, fingerprint, resolutions_json, + continuation_run_id, committed_at +) +SELECT ?, ?, ?, ?, ?, ?, ?, ? +WHERE ( + SELECT COUNT(*) FROM interrupts + WHERE run_id = ? AND status = 'pending' AND generation = ? +) = ? +AND NOT EXISTS ( + SELECT 1 FROM interrupts + WHERE run_id = ? AND status = 'pending' AND generation = ? + AND interrupt_id NOT IN (?, ?) +) +AND NOT EXISTS ( + SELECT 1 FROM interrupts + WHERE run_id = ? AND status = 'pending' AND generation = ? + AND json_extract(binding_json, '$.expiresAt') IS NOT NULL + AND unixepoch(json_extract(binding_json, '$.expiresAt')) * 1000 <= ? +); +``` + +Each update includes `WHERE interrupt_id = ? AND run_id = ? AND status = 'pending' AND generation = ? AND EXISTS (SELECT 1 FROM interrupt_batches WHERE interrupted_run_id = ? AND fingerprint = ? AND resolutions_json = ? AND continuation_run_id = ?)`. Inspect every batch result: the insert and every expected update must report one change. Zero insert changes or a uniqueness race triggers an authoritative batch read; replay requires both fingerprint and canonical `resolutions_json`, otherwise conflict. Any other zero update or statement error throws, and D1 rolls the batch back. The shared conformance suite plus D1 migration tests cover open, commit, all five recovery projections, expiry in the SQL predicate, exact replay, conflict, forced rollback, restart, and migrated pending/resolved rows. + +- [ ] **Step 5: Run green D1 checks and refactor** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-persistence-cloudflare test:lib -- --run tests/runtime.conformance.test.ts tests/migrations.test.ts tests/migration-cli.test.ts +pnpm --filter @tanstack/ai-persistence-cloudflare test:types +``` + +Expected: PASS; forced failure rolls back, exact replay returns the first continuation ID, and different writers observe one conflict. Extract D1 row decoding/recovery helpers and rerun both commands. + +### Task 8: Build the authoritative resume gateway and explicit recovery handler + +**Prerequisites:** Tasks 3-7. All built-in stores must satisfy conformance before server execution depends on them. + +**Files:** +- Modify: `packages/ai-persistence/src/middleware.ts:93-214,675-711,764-911` +- Modify: `packages/ai-persistence/src/capabilities.ts` +- Create: `packages/ai-persistence/src/recovery.ts` +- Modify: `packages/ai-persistence/src/index.ts` +- Modify: `packages/ai-persistence/tests/interrupts.test.ts` +- Modify: `packages/ai-persistence/tests/with-persistence.test.ts` +- Modify: `packages/ai-persistence/tests/persistence-validation.test.ts` +- Modify: `packages/ai/tests/interrupts.test.ts` + +- [ ] **Step 1: Write failing exhaustive-validation and recovery tests** + +Add one batch containing three independent invalid entries and prove all are reported while the store remains pending. Reuse the existing local `mockAdapter` and `collect` test utilities already defined at the top of `tests/interrupts.test.ts`. Define the tool and fixture in that file: + +```ts +const transferInput = z.object({ cents: z.number().int().positive() }) +const transferApprove = z.object({ note: z.string().min(1) }) +const transferReject = z.object({ reason: z.string().min(1) }) +const transferTool = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema: transferInput, + approvalSchema: { + approve: transferApprove, + reject: transferReject, + }, +}).server(async ({ cents }) => ({ receipt: `receipt-${cents}` })) + +async function openApprovalBatch( + persistence: AIPersistence, + interruptedRunId: string, +): Promise { + const store = persistence.stores.interrupts + if (!store) throw new Error('interrupt store missing') + const approval = normalizeApprovalSchema( + transferTool.approvalSchema, + transferTool.inputSchema, + ) + await store.openInterruptBatch({ + threadId: 'thread-1', + interruptedRunId, + descriptors: [ + { + id: 'approve', + reason: 'tool_call', + toolCallId: 'call-approve', + responseSchema: approval.responseSchema, + }, + { + id: 'deny', + reason: 'tool_call', + toolCallId: 'call-deny', + responseSchema: approval.responseSchema, + }, + { + id: 'generic', + reason: 'input_required', + responseSchema: { + type: 'object', + properties: { email: { type: 'string', format: 'email' } }, + required: ['email'], + additionalProperties: false, + }, + }, + ], + bindings: [ + { + kind: 'tool-approval', + interruptId: 'approve', + toolName: 'transfer', + toolCallId: 'call-approve', + originalArgs: { cents: 1 }, + inputSchemaHash: hashSchemaInput(transferTool.inputSchema), + approvalSchemaHash: approval.approvalSchemaHash, + responseSchemaHash: approval.responseSchemaHash, + }, + { + kind: 'tool-approval', + interruptId: 'deny', + toolName: 'transfer', + toolCallId: 'call-deny', + originalArgs: { cents: 2 }, + inputSchemaHash: hashSchemaInput(transferTool.inputSchema), + approvalSchemaHash: approval.approvalSchemaHash, + responseSchemaHash: approval.responseSchemaHash, + }, + { + kind: 'generic', + interruptId: 'generic', + responseSchemaHash: digestInterruptJson(canonicalInterruptJson({ + type: 'object', + properties: { email: { type: 'string', format: 'email' } }, + required: ['email'], + additionalProperties: false, + })), + }, + ], + }) +} +``` + +Obtain the no-execution adapter as `const { adapter: shouldNotRunAdapter, calls: adapterCalls } = mockAdapter([])` and assert `adapterCalls` rather than assuming `chatStream` is a spy: + +```ts +it('reports every invalid response and commits nothing', async () => { + const persistence = memoryPersistence() + const { + adapter: shouldNotRunAdapter, + calls: adapterCalls, + } = mockAdapter([]) + await openApprovalBatch(persistence, 'run-old') + const stream = chat({ + adapter: shouldNotRunAdapter, + messages: [], + threadId: 'thread-1', + runId: 'run-new', + parentRunId: 'run-old', + resume: [ + { + interruptId: 'approve', + status: 'resolved', + payload: { approved: true, editedArgs: { cents: 'wrong' } }, + }, + { + interruptId: 'deny', + status: 'resolved', + payload: { approved: false, editedArgs: { cents: 1 } }, + }, + { + interruptId: 'generic', + status: 'resolved', + payload: { email: 42 }, + }, + ], + tools: [transferTool], + middleware: [withChatPersistence(persistence)], + }) + const chunks = await collect(stream) + const errors = chunks.filter((chunk) => chunk.type === 'RUN_ERROR') + + expect(errors).toHaveLength(1) + expect(errors[0]?.['tanstack:interruptErrors']).toEqual([ + expect.objectContaining({ interruptId: 'approve', code: 'invalid-edited-args' }), + expect.objectContaining({ interruptId: 'deny', code: 'invalid-edited-args' }), + expect.objectContaining({ interruptId: 'generic', code: 'invalid-payload' }), + ]) + expect(await persistence.stores.interrupts?.listPendingByRun('run-old')).toHaveLength(3) + expect(adapterCalls).toHaveLength(0) +}) +``` + +Add the exact-set aggregate test: + +```ts +it('returns item and batch exact-set errors together', async () => { + const persistence = memoryPersistence() + const store = persistence.stores.interrupts + if (!store) throw new Error('interrupt store missing') + const responseSchema = { + type: 'object', + properties: { value: { type: 'string' } }, + required: ['value'], + additionalProperties: false, + } as const + await store.openInterruptBatch({ + threadId: 'thread-1', + interruptedRunId: 'run-old', + descriptors: ['a', 'b', 'c'].map((id) => ({ + id, + reason: 'input_required', + responseSchema, + })), + bindings: ['a', 'b', 'c'].map((interruptId) => ({ + kind: 'generic' as const, + interruptId, + responseSchemaHash: digestInterruptJson( + canonicalInterruptJson(responseSchema), + ), + })), + }) + const { adapter, calls } = mockAdapter([]) + const chunks = await collect(chat({ + adapter, + messages: [], + threadId: 'thread-1', + runId: 'run-new', + parentRunId: 'run-old', + resume: [ + { interruptId: 'a', status: 'resolved', payload: { value: 1 } }, + { interruptId: 'a', status: 'cancelled' }, + { interruptId: 'x', status: 'resolved', payload: { value: 'extra' } }, + ], + middleware: [withChatPersistence(persistence)], + })) + const runErrors = chunks.filter((chunk) => chunk.type === 'RUN_ERROR') + expect(runErrors).toHaveLength(1) + expect(runErrors[0]?.['tanstack:interruptErrors']).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + scope: 'item', + interruptId: 'a', + code: 'invalid-payload', + }), + expect.objectContaining({ + scope: 'item', + interruptId: 'x', + code: 'unknown-interrupt', + }), + expect.objectContaining({ + scope: 'batch', + code: 'protocol', + interruptIds: ['a'], + }), + expect.objectContaining({ + scope: 'batch', + code: 'protocol', + interruptIds: ['x'], + }), + expect.objectContaining({ + scope: 'batch', + code: 'incomplete-batch', + interruptIds: ['b', 'c'], + }), + ]), + ) + expect(await store.listPendingByRun('run-old')).toHaveLength(3) + expect(calls).toHaveLength(0) +}) +``` + +Add a table whose setup functions are concrete closures in the same test file: + +```ts +async function createInvalidGatewayFixture(interruptId: string) { + const persistence = memoryPersistence() + const store = persistence.stores.interrupts + if (!store) throw new Error('interrupt store missing') + const { adapter, calls: adapterCalls } = mockAdapter([]) + const approval = normalizeApprovalSchema( + transferTool.approvalSchema, + transferTool.inputSchema, + ) + const isTool = interruptId === 'tool' || interruptId === 'drift' + const descriptor: Interrupt = isTool + ? { + id: interruptId, + reason: 'tool_call', + toolCallId: `call-${interruptId}`, + responseSchema: approval.responseSchema, + ...(interruptId === 'expired' && { + expiresAt: '2026-07-13T09:59:00.000Z', + }), + } + : { + id: interruptId, + reason: 'confirmation', + responseSchema: { type: 'object' }, + ...(interruptId === 'expired' && { + expiresAt: '2026-07-13T09:59:00.000Z', + }), + } + const binding: UnopenedInterruptBinding = isTool + ? { + kind: 'tool-approval', + interruptId, + toolName: interruptId === 'tool' ? 'removed-tool' : 'transfer', + toolCallId: `call-${interruptId}`, + originalArgs: { cents: 1 }, + inputSchemaHash: + interruptId === 'drift' + ? 'sha256:deployed-different-schema' + : hashSchemaInput(transferTool.inputSchema), + approvalSchemaHash: approval.approvalSchemaHash, + responseSchemaHash: approval.responseSchemaHash, + } + : { + kind: 'generic', + interruptId, + responseSchemaHash: digestInterruptJson( + canonicalInterruptJson(descriptor.responseSchema ?? null), + ), + ...(descriptor.expiresAt !== undefined && { + expiresAt: descriptor.expiresAt, + }), + } + await store.openInterruptBatch({ + threadId: 'thread-1', + interruptedRunId: 'run-old', + descriptors: [descriptor], + bindings: [binding], + }) + return { + persistence, + store, + adapter, + adapterCalls, + threadId: 'thread-1', + interruptedRunId: 'run-old', + tools: interruptId === 'drift' ? [transferTool] : [], + } +} + +it.each([ + ['cancelled payload', { interruptId: 'a', status: 'cancelled', payload: {} }, 'protocol'], + ['expired', { interruptId: 'expired', status: 'cancelled' }, 'expired'], + ['unknown tool after deploy', { interruptId: 'tool', status: 'cancelled' }, 'stale'], + ['schema hash drift', { interruptId: 'drift', status: 'cancelled' }, 'stale'], +] as const)('%s performs no write or execution', async (_name, resumeItem, code) => { + const fixture = await createInvalidGatewayFixture(resumeItem.interruptId) + const chunks = await collect(chat({ + adapter: fixture.adapter, + messages: [], + threadId: fixture.threadId, + runId: 'run-new', + parentRunId: fixture.interruptedRunId, + resume: [resumeItem], + tools: fixture.tools, + middleware: [withChatPersistence(fixture.persistence)], + })) + expect(chunks).toContainEqual( + expect.objectContaining({ + type: 'RUN_ERROR', + 'tanstack:interruptErrors': expect.arrayContaining([ + expect.objectContaining({ code }), + ]), + }), + ) + expect(await fixture.store.listPendingByRun(fixture.interruptedRunId)).toHaveLength(1) + expect(fixture.adapterCalls).toHaveLength(0) +}) +``` + +Separate tests cover missing `parentRunId`, exact replay, a conflicting second batch, and a persistence object with no interrupt store, because those cases use different request/setup shapes. Each asserts one terminal, the named code, zero adapter calls, and unchanged pending rows. Recovery handler tests use the full/redacted authorization contract from Step 6. + +- [ ] **Step 2: Run the red gateway tests** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-persistence test:lib -- --run tests/interrupts.test.ts tests/with-persistence.test.ts tests/persistence-validation.test.ts +``` + +Expected: FAIL because middleware still uses fail-fast `validatePendingResumes`, sequential `applyPendingResumes`, and same-run lookup. + +- [ ] **Step 3: Provide the core capability and atomically open descriptors** + +In middleware `setup`, call `provideInterruptPersistence(ctx, persistence.stores.interrupts)` only after store validation proves all atomic methods exist. In `onChunk`, replace the per-descriptor `create` loop with one call: + +```ts +if ( + chunk.type === EventType.RUN_FINISHED && + chunk.outcome?.type === 'interrupt' +) { + const gateway = getInterruptPersistence(ctx) + const bindings = chunk.outcome.interrupts.map(readServerBinding) + const opened = await gateway.openInterruptBatch({ + threadId: ctx.threadId, + interruptedRunId: ctx.runId, + descriptors: chunk.outcome.interrupts, + bindings, + }) + await interruptRun(runs, ctx.runId) + return { + ...chunk, + outcome: { type: 'interrupt', interrupts: [...opened.descriptors] }, + } +} +``` + +`readServerBinding` accepts only the reserved server-produced metadata shape created in Task 9; missing/malformed bindings fail before this terminal chunk is yielded. `openInterruptBatch` is all-or-none, so an exception cannot leave a partial descriptor set. + +- [ ] **Step 4: Rebind configured tool validators and collect every error** + +Replace `validatePendingResumes` with an async exhaustive function. The authoritative lookup key is stored tool name plus schema hashes, never client metadata: + +```ts +async function validateInterruptResumeBatch(input: { + pending: readonly InterruptRecord[] + resume: readonly RunAgentResumeItem[] + tools: readonly Tool[] + now: number + correlation: InterruptCorrelation +}): Promise { + const itemErrors: ItemInterruptError[] = [] + const batchErrors: BatchInterruptError[] = [] + const pendingIds = new Set(input.pending.map((record) => record.interruptId)) + const resumeById = new Map() + for (const entry of input.resume) { + if (resumeById.has(entry.interruptId)) { + batchErrors.push(batchError('protocol', [entry.interruptId], 'Duplicate response.', input.correlation)) + continue + } + resumeById.set(entry.interruptId, entry) + if (!pendingIds.has(entry.interruptId)) { + itemErrors.push(itemError(entry.interruptId, 'unknown-interrupt', 'Unknown interrupt.', input.correlation)) + batchErrors.push(batchError('protocol', [entry.interruptId], 'Extra response.', input.correlation)) + } + } + const missing = [...pendingIds].filter((id) => !resumeById.has(id)).sort() + if (missing.length > 0) { + batchErrors.push(batchError('incomplete-batch', missing, 'Missing responses.', input.correlation)) + } + + for (const record of [...input.pending].sort((a, b) => + a.interruptId.localeCompare(b.interruptId), + )) { + const entry = resumeById.get(record.interruptId) + if (!entry) continue + const rebound = rebindAuthoritativeValidator(record.binding, input.tools) + if (rebound.status === 'stale') { + itemErrors.push(staleSchemaError(record, input.correlation)) + continue + } + itemErrors.push(...await validateOneResume(record, entry, rebound, input)) + } + itemErrors.sort((left, right) => left.interruptId.localeCompare(right.interruptId)) + return itemErrors.length > 0 || batchErrors.length > 0 + ? { ok: false, itemErrors, batchErrors, errors: [...itemErrors, ...batchErrors] } + : { ok: true, state: buildResumeState(input.pending, resumeById) } +} +``` + +`batchError` and `itemError` are constructors that copy the full correlation and set `source: 'server'`/`retryable: false`. `validateOneResume` is exhaustive: cancelled entries reject any payload; tool approval validates the response envelope, then edited arguments and selected nested payload; rejection forbids edits; client-tool execution validates output; generic validates the descriptor response schema. Standard Schema validators are awaited and their returned issue order is preserved. Only the outer `itemErrors` array is sorted by interrupt ID. Missing tool/hash drift returns one `stale` item error and never an unconstrained validator. + +- [ ] **Step 5: Commit once using `parentRunId` as the CAS key** + +In `onConfig`, reject `resume` without `ctx.parentRunId`; load only `listPendingByRun(ctx.parentRunId)`, validate every item, and compute the fingerprint only after validation. Commit exactly once with the candidate continuation `ctx.runId`: + +```ts +const candidate = canonicalizeInterruptResolutions(config.resume) +const committed = await store.commitInterruptResolutions({ + threadId: ctx.threadId, + interruptedRunId: ctx.parentRunId, + continuationRunId: ctx.runId, + expectedGeneration: pending[0]?.generation ?? 0, + expectedInterruptIds: pending.map((record) => record.interruptId), + resolutions: candidate.resolutions, + fingerprint: candidate.fingerprint, + canonicalResolutions: candidate.canonicalResolutions, +}) +``` + +On `committed`, return `resumeToolState` containing approvals with edited arguments/custom payload, client results, denied results, cancelled tool-call IDs, and generic resolutions. Later application middleware can read `genericInterrupts` and continue its own workflow without a tool event. On `replayed`, throw an internal replay signal carrying the winning continuation ID; Task 9 converts it to an accepted terminal replay result without executing tools twice. On `conflict`, throw a serializable interrupt submission failure with authoritative recovery. + +- [ ] **Step 6: Create a direct and authenticated HTTP recovery helper** + +Implement both an in-process function and an explicit handler factory: + +```ts +export async function getInterruptRecoveryState( + persistence: AIPersistence, + input: InterruptRecoveryQuery, +): Promise { + const store = persistence.stores.interrupts + if (!store) throw new Error('Interrupt persistence is not configured.') + return store.getInterruptRecoveryState(input) +} + +export function createInterruptRecoveryHandler(options: { + persistence: AIPersistence + authorize: ( + request: Request, + input: InterruptRecoveryQuery, + ) => RecoveryAuthorization | Promise +}): (request: Request) => Promise { + return async (request) => { + const input = parseInterruptRecoveryQuery(await request.json()) + const authorization = await options.authorize(request, input) + if (!authorization.allowed) { + return new Response('Forbidden', { status: 403 }) + } + const state = await getInterruptRecoveryState(options.persistence, input) + return Response.json(projectRecoveryState(state, authorization)) + } +} + +export type RecoveryAuthorization = + | { allowed: false } + | { allowed: true; includeResolutions: boolean } + +function projectRecoveryState( + state: InterruptRecoveryStateV1, + authorization: Extract, +): InterruptRecoveryStateV1 { + if ( + authorization.includeResolutions || + state.state !== 'committed' || + !state.committed + ) { + return cloneAndDeepFreezeJson(state) + } + const { resolutions: _redacted, ...committed } = state.committed + return cloneAndDeepFreezeJson({ ...state, committed }) +} +``` + +Test `{ allowed: false }` → 403, `{ allowed: true, includeResolutions: false }` → committed DTO without `resolutions`, and `includeResolutions: true` → full frozen resolutions. The helper does not register a route or infer a URL. + +- [ ] **Step 7: Run green gateway checks and refactor** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-persistence test:lib -- --run tests/interrupts.test.ts tests/memory.conformance.test.ts tests/with-persistence.test.ts tests/persistence-validation.test.ts +pnpm --filter @tanstack/ai-persistence test:types +``` + +Expected: PASS; three invalid entries arrive in one error, nothing writes or executes, and schema drift is stale. Split exact-set, validator-rebind, and resume-state builders into pure private helpers, rerun both commands, and confirm core still has no persistence-package import. + +### Task 9: Correct the AG-UI interrupt lifecycle and resumed tool audit + +**Prerequisites:** Task 8. + +**Files:** +- Modify: `packages/ai/src/activities/chat/index.ts:504-980,1409-1640,1673-1888,2821-2860` +- Modify: `packages/ai/src/activities/chat/tools/tool-calls.ts:425-448,681-890` +- Modify: `packages/ai/src/types.ts:590-735,980-1010,1081-1124` +- Modify: `packages/ai/src/utilities/chat-params.ts:45-104` +- Modify: `packages/ai/tests/interrupts.test.ts` +- Modify: `packages/ai/tests/chat.test.ts:423-870` +- Modify: `packages/ai/tests/chat-params.test.ts` +- Modify: `packages/ai/tests/stream-processor.test.ts:1197-1376,2177-2220` + +- [ ] **Step 1: Write failing wire-order and continuation-audit tests** + +Add a protocol test that records exact chunk order and a continuation test with approved, denied, and cancelled tool approvals: + +```ts +expect(interruptedChunks.map((chunk) => chunk.type).slice(-3)).toEqual([ + EventType.MESSAGES_SNAPSHOT, + EventType.STATE_SNAPSHOT, + EventType.RUN_FINISHED, +]) +expect(interruptOutcome.interrupts).toEqual([ + expect.objectContaining({ reason: 'tool_call' }), + expect.objectContaining({ reason: 'tanstack:client_tool_execution' }), +]) + +const resumedToolEvents = continuationChunks.filter((chunk) => + chunk.type.startsWith('TOOL_CALL_'), +) +expect(resumedToolEvents.map((chunk) => chunk.type)).toEqual([ + EventType.TOOL_CALL_RESULT, + EventType.TOOL_CALL_RESULT, +]) +expect(resumedToolEvents.map((chunk) => chunk.toolCallId)).toEqual([ + 'approved-call', + 'denied-call', +]) +expect(continuationChunks).not.toContainEqual( + expect.objectContaining({ toolCallId: 'cancelled-call' }), +) +``` + +Assert the continuation has a new run ID, its `parentRunId` equals the interrupted run, its thread is unchanged, approved edits replace rather than merge, denied history is `{ status: 'denied', payload }`, cancellation is persisted in the interrupt audit and marks the original call not-executed during provider-history reconstruction, generic continuation reaches application middleware, and `onError` runs once for interrupt submission failure. Add a run with an approval-producing tool but no interrupt persistence capability and assert one `RUN_ERROR` with `persistence-required`, no snapshots, and no interrupt `RUN_FINISHED`. + +- [ ] **Step 2: Run the red core lifecycle tests** + +Run: + +```powershell +pnpm --filter @tanstack/ai test:lib -- --run tests/interrupts.test.ts tests/chat.test.ts tests/chat-params.test.ts +``` + +Expected: FAIL on old reasons `approval_required` / `client_tool_input`, replayed start/args/end events, missing snapshots, and missing structured error handling. + +- [ ] **Step 3: Build canonical descriptors and reserved bindings** + +Change `buildActionableInterrupts` to use normalized tool schemas and server-owned binding metadata: + +```ts +const approvalSchema = normalizeApprovalSchema( + tool.approvalSchema, + tool.inputSchema, +) +interrupts.push({ + id: approval.approvalId, + reason: 'tool_call', + message: `Approval required to run ${approval.toolName}`, + toolCallId: approval.toolCallId, + responseSchema: approvalSchema.responseSchema, + metadata: { + 'tanstack:binding': { + kind: 'tool-approval', + toolName: approval.toolName, + toolCallId: approval.toolCallId, + originalArgs: approval.input, + responseSchemaHash: approvalSchema.responseSchemaHash, + inputSchemaHash: hashSchemaInput(tool.inputSchema), + approvalSchemaHash: approvalSchema.approvalSchemaHash, + }, + }, +}) +``` + +Client-tool descriptors use `tanstack:client_tool_execution`, the original call ID, and output/response schema identities. Generic reasons retain `input_required`, `confirmation`, or an unknown namespaced/custom reason. Invalid raw schemas fail before a descriptor is emitted. + +Before returning any nonempty actionable descriptor set, require the core-owned `InterruptPersistenceCapability` from Task 3. If the capability is absent, throw `InterruptSubmissionFailure` with batch code `persistence-required` before emitting snapshots or an interrupt terminal. This check lives in core lifecycle code; merely omitting `withChatPersistence` must never create an unresumable native interrupt. + +- [ ] **Step 4: Emit required snapshots immediately before the terminal** + +Add `state?: unknown` to chat text options and pass the parsed AG-UI state from `chatParamsFromRequestBody`. Before every interrupt terminal path, yield: + +```ts +yield* this.pipeThroughMiddleware({ + type: EventType.MESSAGES_SNAPSHOT, + messages: uiMessagesToWire(modelMessagesToUIMessages(this.messages)), + timestamp: Date.now(), +}) +if (this.params.state !== undefined) { + yield* this.pipeThroughMiddleware({ + type: EventType.STATE_SNAPSHOT, + snapshot: this.params.state, + timestamp: Date.now(), + }) +} +yield* this.pipeThroughMiddleware(this.buildInterruptFinishedChunk( + finishEvent, + executionResult.needsApproval, + executionResult.needsClientExecution, +)) +``` + +Because middleware processes a chunk before the generator yields it, Task 8's atomic descriptor open must finish before `RUN_FINISHED` becomes observable. If it fails, snapshots may precede the single `RUN_ERROR`, but no interrupt terminal is emitted. + +- [ ] **Step 5: Carry rich approval decisions into tool execution** + +Change `executeToolCalls` from `Map` to `ReadonlyMap`. Select the authoritative stored original input unless an approved decision contains `editedArgs`; use edits as a complete replacement and validate them before execution. A denial produces one synthetic result: + +```ts +if (!decision.approved) { + results.push({ + toolCallId: toolCall.id, + toolName, + result: { + status: 'denied', + ...(decision.payload !== undefined && { payload: decision.payload }), + }, + state: 'output-available', + }) + continue +} +input = decision.editedArgs === undefined ? input : decision.editedArgs +``` + +Filter `cancelledToolCallIds` before execution and result construction. Record those call IDs as not-executed in the interrupt audit/provider-history reconstruction path without producing a result chunk. Client-tool resolved output follows the existing result path; generic resolutions remain only in middleware continuation state and create no tool chunk. + +- [ ] **Step 6: Make resumed tool output result-only** + +Delete the `argsMap` branch that synthesizes `TOOL_CALL_START`, `TOOL_CALL_ARGS`, and `TOOL_CALL_END` in `buildToolResultChunks`. Always emit one `TOOL_CALL_RESULT` for actual approved/client results and synthetic denied results, correlated to the original `toolCallId`. Persist corresponding actual/denied history before the resumed model begins; cancelled calls remain in the interrupt audit record and are excluded from actionable provider history. + +- [ ] **Step 7: Emit exactly one structured terminal on gateway failures and replay** + +Introduce an internal `InterruptSubmissionFailure` carrying serialized errors/recovery and an `InterruptReplaySignal` carrying the winning continuation ID. In the `TextEngine.run()` catch, call middleware `onError` once, then yield one `RUN_ERROR` for submission failure and return without rethrowing or adding `RUN_FINISHED`: + +```ts +if (error instanceof InterruptSubmissionFailure) { + await this.runInterruptOnErrorOnce(error) + yield { + type: EventType.RUN_ERROR, + threadId: this.middlewareCtx.threadId, + runId: this.middlewareCtx.runId, + timestamp: Date.now(), + message: error.message, + code: 'INTERRUPT_SUBMISSION_FAILED', + 'tanstack:interruptErrors': error.errors, + ...(error.recovery && { + 'tanstack:interruptRecovery': error.recovery, + }), + } + return +} +``` + +For exact replay, emit a successful `RUN_FINISHED` whose standard `result` is canonical JSON `{ "type": "tanstack:interrupt-replay", "continuationRunId": "winning-run" }`; do not execute tools. The client recognizes this result and joins/reloads the winner. `runStreamingText` disposes MCP resources in `finally` but does not translate or duplicate these terminal signals. + +- [ ] **Step 8: Run green lifecycle checks and audit legacy emission** + +Run: + +```powershell +pnpm --filter @tanstack/ai test:lib -- --run tests/interrupts.test.ts tests/chat.test.ts tests/chat-params.test.ts tests/stream-processor.test.ts +pnpm --filter @tanstack/ai test:types +``` + +Expected: PASS; snapshots precede a nonempty interrupt outcome, continuations use a new run with `parentRunId`, and resumed tools emit results only. Confirm native server paths emit no `approval-requested` or `tool-input-available` custom event; deprecated readers remain tested for old streams. + +- [ ] **Step 9: Focused refactor and recheck** + +Extract descriptor/binding construction and terminal-error serialization into focused helpers, rerun the same commands, and verify every terminal path calls exactly one of `onFinish`, `onAbort`, or `onError`. + +### Task 10: Define and hydrate the typed bound interrupt union + +**Prerequisites:** Tasks 2, 3, and 9. + +**Files:** +- Create: `packages/ai-client/src/interrupt-manager.ts` +- Modify: `packages/ai-client/src/types.ts:1-110,397-585` +- Modify: `packages/ai-client/src/index.ts` +- Create: `packages/ai-client/tests/chat-client-interrupts.test.ts` +- Create: `packages/ai-client/tests/interrupts-types.test-d.ts` + +- [ ] **Step 1: Write failing type and hydration tests before implementation** + +Define an approval tool with distinct branches and assert the exact resolver calls: + +```ts +const transfer = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema: z.object({ cents: z.number() }), + outputSchema: z.object({ receipt: z.string() }), + approvalSchema: { + approve: z.object({ note: z.string() }), + reject: z.object({ reason: z.string() }), + }, +}).client() + +declare const interrupt: Extract< + ChatInterrupt, + { kind: 'tool-approval'; toolName: 'transfer' } +> + +interrupt.resolveInterrupt(true, { + editedArgs: { cents: 500 }, + payload: { note: 'Reviewed' }, +}) +interrupt.resolveInterrupt(false, { payload: { reason: 'Policy' } }) +// @ts-expect-error rejection never accepts editedArgs +interrupt.resolveInterrupt(false, { editedArgs: { cents: 1 }, payload: { reason: 'Policy' } }) +// @ts-expect-error custom data belongs under payload +interrupt.resolveInterrupt(true, { editedArgs: { cents: 1 }, note: 'Reviewed' }) + +const noInput = toolDefinition({ + name: 'noInput', + description: 'No editable input', + needsApproval: true, +}).client() +declare const noInputInterrupt: Extract< + ChatInterrupt, + { kind: 'tool-approval'; toolName: 'noInput' } +> +noInputInterrupt.resolveInterrupt(true) +// @ts-expect-error omitted inputSchema forbids editedArgs +noInputInterrupt.resolveInterrupt(true, { editedArgs: {} }) + +declare const generic: Extract< + ChatInterrupt, + { kind: 'generic' } +> +generic.resolveInterrupt({ applicationValue: 'runtime-validated' }) +``` + +In `chat-client-interrupts.test.ts`, construct the manager with the literal `transfer` tool and hydrate one descriptor whose metadata binding contains matching `toolCallId`, `interruptedRunId`, `generation`, tool name, and hashes. Clone the case while changing each field independently and assert only the fully matching case becomes `tool-approval`; every mismatch becomes `generic`. Add matching/mismatched client-tool, unknown reason, invalid generic schema, and expired descriptor cases before production code. + +- [ ] **Step 2: Run the red client type target** + +Run both suites red: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-client test:types +pnpm --filter @tanstack/ai-client test:lib -- --run tests/chat-client-interrupts.test.ts +``` + +Expected: FAIL because `ChatInterrupt`, bound variants, `InterruptManager`, and hydration do not exist. + +- [ ] **Step 3: Add the public base and mapped union** + +Define the public surface in `types.ts` and keep behavior private to the manager: + +```ts +export interface BoundInterruptBase { + id: string + reason: string + message?: string + toolCallId?: string + responseSchema?: JSONSchema + expiresAt?: string + metadata?: Record + kind: 'tool-approval' | 'client-tool-execution' | 'generic' + status: 'pending' | 'validating' | 'staged' | 'submitting' | 'error' + provenance: 'native' | 'legacy' + generation: number + errors: readonly ItemInterruptError[] + stagedResponse?: TDraft + clearResolution(): void + cancelInterrupt(): void +} + +export interface GenericAGUIInterrupt extends BoundInterruptBase< + | { status: 'resolved'; payload: unknown } + | { status: 'cancelled' } +> { + kind: 'generic' + toolName?: never + resolveInterrupt(payload: unknown): void +} + +export type BoundInterrupts< + TTools extends readonly AnyClientTool[], +> = readonly ChatInterrupt[] + +declare const noSchema: unique symbol +type NoSchema = typeof noSchema +type InputSchemaOf = TTool extends { inputSchema?: infer TSchema } + ? TSchema extends undefined ? NoSchema : TSchema + : NoSchema +type OutputSchemaOf = TTool extends { outputSchema?: infer TSchema } + ? TSchema extends undefined ? NoSchema : TSchema + : NoSchema +type ApproveSchemaOf = ApprovalSchemaOf extends { + approve?: infer TSchema +} ? TSchema extends undefined ? NoSchema : TSchema + : ApprovalSchemaOf extends undefined ? NoSchema + : ApprovalSchemaOf +type RejectSchemaOf = ApprovalSchemaOf extends { + reject?: infer TSchema +} ? TSchema extends undefined ? NoSchema : TSchema + : ApprovalSchemaOf extends undefined ? NoSchema + : ApprovalSchemaOf +type EditedArgsField = TSchema extends NoSchema + ? { editedArgs?: never } + : { editedArgs?: InferSchemaType } +type PayloadField = TSchema extends NoSchema + ? { payload?: never } + : TSchema extends JSONSchema + ? { payload: unknown } + : undefined extends InferSchemaType + ? { payload?: InferSchemaType } + : { payload: InferSchemaType } +type RequiredKeys = { + [TKey in keyof T]-?: Record extends Pick + ? never + : TKey +}[keyof T] +type OptionsTuple = [RequiredKeys] extends [never] + ? [options?: T] + : [options: T] + +export interface ToolApprovalInterrupt + extends BoundInterruptBase< + | ({ status: 'resolved'; approved: true } & + EditedArgsField> & + PayloadField>) + | ({ status: 'resolved'; approved: false } & + PayloadField>) + | { status: 'cancelled' } + > { + kind: 'tool-approval' + reason: 'tool_call' + toolName: TTool['name'] + toolCallId: string + resolveInterrupt( + approved: true, + ...options: OptionsTuple< + EditedArgsField> & PayloadField> + > + ): void + resolveInterrupt( + approved: false, + ...options: OptionsTuple>> + ): void +} + +export interface ClientToolExecutionInterrupt + extends BoundInterruptBase< + | { status: 'resolved'; payload: InferSchemaType> } + | { status: 'cancelled' } + > { + kind: 'client-tool-execution' + reason: 'tanstack:client_tool_execution' + toolName: TTool['name'] + toolCallId: string + resolveInterrupt(output: InferSchemaType>): void +} + +type ToolByName = + Extract +type ToolNames = TTools[number]['name'] +type ApprovalNames = { + [TName in ToolNames]: ApprovalCapabilityOf< + ToolByName + > extends true ? TName : never +}[ToolNames] +type ToolApprovalInterruptsFor = { + [TName in ApprovalNames]: ToolApprovalInterrupt> +}[ApprovalNames] +type ClientToolExecutionInterruptsFor = { + [TName in ToolNames]: ClientToolExecutionInterrupt> +}[ToolNames] +export type ChatInterrupt = + | ToolApprovalInterruptsFor + | ClientToolExecutionInterruptsFor + | GenericAGUIInterrupt +``` + +Use this contiguous type block verbatim and add the shared/branch/omitted/raw/optional assertions named in Step 1. `OutputSchemaOf` maps to `unknown` for client-tool execution rather than exposing the private sentinel. + +- [ ] **Step 4: Hydrate only trusted known-tool matches** + +Create `InterruptManager` with injected tools, submit, recovery, state-change, and draft-persistence callbacks. Hydration follows this decision: + +```ts +function classifyDescriptor( + descriptor: Interrupt, + registry: ToolRegistry, + correlation: { + interruptedRunId: string + generation: number + }, +): HydratedKind { + const binding = readUntrustedDescriptorBinding(descriptor.metadata) + if ( + descriptor.reason === 'tool_call' && + descriptor.toolCallId && + binding?.kind === 'tool-approval' && + descriptor.toolCallId === binding.toolCallId && + binding.interruptedRunId === correlation.interruptedRunId && + binding.generation === correlation.generation + ) { + const tool = registry.approvalTools.get(binding.toolName) + if (tool && registry.matchesHashes(tool, binding)) { + return { kind: 'tool-approval', tool, binding } + } + } + if ( + descriptor.reason === 'tanstack:client_tool_execution' && + descriptor.toolCallId && + binding?.kind === 'client-tool-execution' && + descriptor.toolCallId === binding.toolCallId && + binding.interruptedRunId === correlation.interruptedRunId && + binding.generation === correlation.generation + ) { + const tool = registry.tools.get(binding.toolName) + if (tool && registry.matchesHashes(tool, binding)) { + return { kind: 'client-tool-execution', tool, binding } + } + } + return { kind: 'generic' } +} +``` + +`InterruptManager.hydrate` passes its own authoritative correlation into this function; it never reads the expected run/generation from the descriptor alone. Bind methods to manager/item IDs, compile generic response schemas with Task 1, retain configured Standard Schema validators for known tools, and expose core-cloned/deep-frozen snapshots. + +- [ ] **Step 5: Run the already-red hydration tests green** + +Run the tests written in Step 1: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-client test:lib -- --run tests/chat-client-interrupts.test.ts +``` + +Expected: PASS; forged/mismatched descriptors are generic and an invalid generic schema can only be cancelled. + +- [ ] **Step 6: Run green type checks and refactor** + +Run: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-client test:types +pnpm --filter @tanstack/ai-client test:lib -- --run tests/chat-client-interrupts.test.ts +``` + +Expected: PASS. Extract type helpers (`NoSchema`, branch selectors, required-key tuple) into one contiguous section and registry hashing into one private class; rerun both commands. + +### Task 11: Implement staging, callback transactions, bulk operations, and retry + +**Prerequisites:** Task 10. + +**Files:** +- Modify: `packages/ai-client/src/interrupt-manager.ts` +- Modify: `packages/ai-client/src/types.ts` +- Modify: `packages/ai-client/tests/chat-client-interrupts.test.ts` + +- [ ] **Step 1: Add failing state-machine tests** + +Cover singleton resolve/cancel auto-submit; three-item wait; replacement; clear; async validator generations; invalid replacement preserving the previous valid draft; submitting mutation rejection; root boolean eligibility; callback visit-once/submit-once; throw/non-`undefined`/thenable/incomplete rollback; late post-await token use; transport frozen retry; mutation invalidating retry; server item/root errors; acceptance removal; and status/error notifications. + +Use a deferred async-callback regression: + +```ts +function deferred() { + let resolve!: (value: T | PromiseLike) => void + let reject!: (reason?: unknown) => void + const promise = new Promise((resolvePromise, rejectPromise) => { + resolve = resolvePromise + reject = rejectPromise + }) + return { promise, resolve, reject } +} + +const testTools = [transferTool, auditTool, notificationTool] as const +type TestTools = typeof testTools + +function createThreeInterruptManager(options: { + submit: InterruptManagerOptions['submit'] +}) { + const manager = new InterruptManager({ + tools: testTools, + submit: options.submit, + onStateChange: vi.fn(), + }) + manager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-old', + generation: 1, + descriptors: threeApprovalDescriptors, + }) + return manager +} + +it('invalidates a callback token before a late async item call', async () => { + const gate = deferred() + const submit = vi.fn() + const manager = createThreeInterruptManager({ submit }) + let captured: ChatInterrupt | undefined + + await expect( + manager.resolveInterrupts(((interrupt: ChatInterrupt) => { + captured = interrupt + return gate.promise.then(() => { + interrupt.cancelInterrupt() + }) + }) as (interrupt: ChatInterrupt) => undefined), + ).rejects.toMatchObject({ code: 'async-resolver' }) + + gate.resolve() + await gate.promise + captured?.cancelInterrupt() + expect(submit).not.toHaveBeenCalled() + expect(manager.interruptErrors).toContainEqual( + expect.objectContaining({ code: 'inactive-transaction' }), + ) +}) +``` + +Define `transferTool`, `auditTool`, and `notificationTool` immediately above this block with literal names and payload-less approval branches. Define `threeApprovalDescriptors` as three native `tool_call` descriptors whose reserved bindings and schema hashes match those tools. `InterruptManagerOptions`, `hydrate`, and `onStateChange` are the concrete public-to-package-private constructor seam established in Task 10; use these exact names in implementation and tests so the fixture does not depend on hidden maps. + +- [ ] **Step 2: Run the red state-machine test** + +Run: + +```powershell +pnpm --filter @tanstack/ai-client test:lib -- --run tests/chat-client-interrupts.test.ts +``` + +Expected: FAIL because bound methods do not yet stage, gate, roll back, freeze, or retry. + +- [ ] **Step 3: Implement synchronous candidate staging with async validation generations** + +Item methods return `void`. Each candidate increments an item generation, sets `validating` when a Standard Schema result is thenable, and ignores a result whose captured generation is stale. A valid candidate replaces the draft; an invalid candidate keeps the last valid `stagedResponse`, sets item errors/status, and blocks submission. Cancellation is payloadless and needs no schema validation. `clearResolution` removes draft/errors and returns to pending. Replacement/clear/cancel while submitting throws a non-retryable protocol error. + +- [ ] **Step 4: Freeze and submit only a complete valid batch** + +Implement one submission gate: + +```ts +private maybeSubmit(): void { + if (this.transaction || this.submission || this.items.size === 0) return + const items = [...this.items.values()] + if (items.some((item) => item.status === 'validating' || item.status === 'error')) return + if (items.some((item) => item.draft === undefined)) return + const entries = items + .map((item) => item.draft?.response) + .filter((entry): entry is RunAgentResumeItem => entry !== undefined) + .sort((a, b) => a.interruptId.localeCompare(b.interruptId)) + const candidate = canonicalizeInterruptResolutions(entries) + const frozen = cloneAndDeepFreezeJson({ + generation: this.batchGeneration, + entries: candidate.resolutions, + fingerprint: candidate.fingerprint, + canonicalResolutions: candidate.canonicalResolutions, + }) + void this.submitFrozenBatch(frozen) +} +``` + +A singleton reaches this gate after its first valid decision. Multiple items wait for complete coverage. During submission every item is `submitting`. Acceptance removes active descriptors/drafts. A retryable transport/availability failure restores `staged`, retains the exact frozen batch, and appends a root transport/server error. + +- [ ] **Step 5: Implement root boolean and synchronous callback transactions** + +The boolean overload first verifies every open item is a tool approval and the chosen branch has no required payload; only then stage all and submit once. For callbacks, capture one stable item snapshot, clone all drafts/errors/statuses, install a unique token, suppress auto-submit, and invoke once per item. Require literal runtime return `undefined`; reject thenables and other values. On throw, invalid return, or incomplete coverage, restore the clone, seal the token, increment generation, and start no request. Transaction-scoped item methods check the active token and generation; late calls add `inactive-transaction` without mutation. + +- [ ] **Step 6: Implement cancel-all, exact retry, and error mapping** + +`cancelInterrupts()` stages payloadless cancellation for the stable full set and submits once. `retryInterrupts()` is allowed only while an intact retryable frozen batch exists. Any edit/clear/cancel after failure increments batch generation, clears the frozen fingerprint and superseded transport error, and disables retry. Map item server errors by `interruptId`; map batch/transport errors only to `interruptErrors`. Expired/stale/conflict results invoke recovery instead of retry. + +- [ ] **Step 7: Run green state-machine tests and refactor** + +Run: + +```powershell +pnpm --filter @tanstack/ai-client test:lib -- --run tests/chat-client-interrupts.test.ts +pnpm --filter @tanstack/ai-client test:types +``` + +Expected: PASS; one interrupt submits immediately, three wait, callback failures roll back, and late async calls cannot submit. Refactor snapshots, token sealing, and status notifications into focused private helpers and rerun both commands. + +### Task 12: Integrate the manager into ChatClient and legacy compatibility + +**Prerequisites:** Task 11. + +**Files:** +- Modify: `packages/ai-client/src/chat-client.ts:41-146,191-275,510-706,967-1070,1403-1585` +- Modify: `packages/ai-client/src/connection-adapters.ts:353-450,1171-1205` +- Modify: `packages/ai-client/src/types.ts:397-585` +- Modify: `packages/ai-client/src/index.ts` +- Modify: `packages/ai-client/tests/chat-client-interrupts.test.ts` +- Modify: `packages/ai-client/tests/chat-client-resume.test.ts` +- Modify: `packages/ai-client/tests/chat-client-client-tool-status.test.ts` +- Modify: `packages/ai-client/tests/connection-adapters.test.ts` + +- [ ] **Step 1: Add failing transport/integration/compatibility tests** + +Prove that a native resume sends no messages, uses a fresh run ID with `parentRunId` equal to the interrupted run, parses item/root error extensions, recognizes the replay terminal, and blocks unrelated input. Add legacy approval/custom-tool events and assert one history request with no native resume. Add the sandbox collision case: + +```ts +client.processChunk({ + type: EventType.CUSTOM, + name: 'approval-requested', + value: { approvalId: 'sandbox-approval', operation: 'process.exec' }, + timestamp: Date.now(), +}) +expect(client.getInterrupts()).toEqual([]) +``` + +Also assert `addToolApprovalResponse({ approved: false })` stages a resolved denial, `cancelInterrupt` emits `cancelled`, `addToolResult` delegates to a matching client-tool interrupt, mixed provenance stops, and native server events never cause a legacy follow-up. + +- [ ] **Step 2: Run the red integration tests** + +Run: + +```powershell +pnpm --filter @tanstack/ai-client test:lib -- --run tests/chat-client-interrupts.test.ts tests/chat-client-resume.test.ts tests/chat-client-client-tool-status.test.ts tests/connection-adapters.test.ts +``` + +Expected: FAIL because ChatClient still owns raw `pendingInterrupts`, reuses the interrupted run ID, and sends compatibility responses directly. + +- [ ] **Step 3: Make InterruptManager the single normalized owner** + +Replace `private pendingInterrupts` and `pendingInterruptRunId` with one manager. `observeInterruptState` hydrates native outcomes with `{ threadId, interruptedRunId, generation }`; RUN_ERROR extensions flow to manager error/recovery handling. Expose: + +```ts +getInterrupts(): BoundInterrupts +getInterruptErrors(): readonly InterruptSubmissionError[] +getIsResuming(): boolean + +/** @deprecated Use getInterrupts(). Removed in 1.0. */ +getPendingInterrupts(): BoundInterrupts + +resolveInterrupts( + decisionOrResolver: boolean | ((interrupt: ChatInterrupt) => undefined), +): Promise + +cancelInterrupts(): Promise +retryInterrupts(): Promise +resumeInterruptsUnsafe( + entries: readonly RunAgentResumeItem[], + state?: ChatResumeState, +): Promise + +/** @deprecated Use resumeInterruptsUnsafe(). Removed in 1.0. */ +resumeInterrupts( + entries: readonly RunAgentResumeItem[], + state?: ChatResumeState, +): Promise +``` + +Add one immutable state notification so frameworks never compose three independently timed reads: + +```ts +export interface ChatInterruptState { + interrupts: BoundInterrupts + interruptErrors: readonly InterruptSubmissionError[] + isResuming: boolean +} + +onInterruptStateChange?: ( + state: ChatInterruptState, +) => void +``` + +Emit it after every hydration, draft/error/status transition, submission start/end, recovery replacement, reset, and disposal. `getInterrupts`, `getInterruptErrors`, and `getIsResuming` return the exact current immutable snapshot fields. + +Create one fresh deeply frozen `BoundInterrupts` array per notification. `getInterrupts()`, deprecated `getPendingInterrupts()`, hook `interrupts`, and hook `pendingInterrupts` all return that exact array object. Do not maintain a second raw descriptor store or project legacy DTOs into the deprecated alias. + +- [ ] **Step 4: Send continuations with new run correlation** + +When manager submits a native batch, generate a new run ID, keep the same thread, set `parentRunId` to the interrupted run, send the complete `resume` array, and send `[]` messages. `resumeInterruptsUnsafe` uses the same correlation/submission path but accepts raw entries. Keep expiry, exact set, server validation, and CAS enforced. + +- [ ] **Step 5: Handle replay and structured failures** + +Parse `tanstack:interruptErrors` and `tanstack:interruptRecovery` from RUN_ERROR. For a committed recovery DTO or canonical interrupt-replay result, attach only to the winning continuation: + +```ts +private async attachWinningContinuation( + threadId: string, + continuationRunId: string, +): Promise { + const controller = new AbortController() + const chunks = + this.connection && 'joinRun' in this.connection + ? this.connection.joinRun(continuationRunId, controller.signal) + : this.options.continuationLoader + ? this.options.continuationLoader( + { threadId, continuationRunId }, + { signal: controller.signal }, + ) + : undefined + if (!chunks) { + this.interruptManager.quarantineForRecovery({ + code: 'recovery-unavailable', + message: 'The winning continuation cannot be joined by this client.', + continuationRunId, + }) + return + } + await this.consumeReadOnlyContinuation(chunks, continuationRunId) +} + +private processIncomingChunk(chunk: StreamChunk): void { + this.callbacksRef.current.onChunk(chunk) + this.devtoolsBridge.observeChunk(chunk) + this.processor.processChunk(chunk) + this.updateRunLifecycle(chunk) + this.observeInterruptState(chunk) +} + +private async consumeReadOnlyContinuation( + chunks: AsyncIterable, + continuationRunId: string, +): Promise { + for await (const chunk of chunks) { + const chunkRunId = getChunkRunId(chunk) + if (chunkRunId !== undefined && chunkRunId !== continuationRunId) { + throw new Error('Continuation loader returned a different run.') + } + this.processIncomingChunk(chunk) + } +} +``` + +Replace the duplicated incoming-chunk statements in the existing connect/subscription loops with `processIncomingChunk(chunk)`. The read-only loader never calls `connect`, `sendMessage`, `reload`, the model endpoint, or the interrupt recovery loader. Tests provide a non-resumable connection plus a spy `continuationLoader`, assert the loader receives the stored winning ID, and assert no new model run is scheduled. A client with neither `joinRun` nor `continuationLoader` remains quarantined with `recovery-unavailable`. + +- [ ] **Step 6: Normalize legacy readers without a dual writer** + +Recognize legacy `approval-requested` only when its value has `toolCallId`, `toolName`, and `approval: { id, needsApproval: true }`; recognize `tool-input-available` only with its full historical tool shape. This prevents sandbox's path-only approval event from colliding. Hydrate `provenance: 'legacy'` and reject edits, custom payloads, generic responses, native cancellation, expiry, and conflict recovery as `legacy-unsupported`. + +When all legacy items are covered, clone current message history, update every matching approval/result part in the clone, and send one follow-up history request with no `resume`. On construction failure, change no live message. On transport failure, keep staged responses and report `legacy-submit-failed`. Reject mixed native/legacy sets before any mutation. + +- [ ] **Step 7: Deprecate approval APIs and preserve client-tool results** + +Add removal-at-1.0 JSDoc to `addToolApprovalResponse`, `pendingInterrupts`/`getPendingInterrupts`, raw `resumeInterrupts`, and legacy custom-event reader types. Implement `addToolApprovalResponse` by finding the bound item and calling `resolveInterrupt(response.approved)`; false is denial, not cancellation. Implement old `resumeInterrupts` as a compatibility wrapper over `resumeInterruptsUnsafe` while retaining its historical return type. Keep `addToolResult` supported and delegate to `resolveInterrupt(output)` for a matching native client-tool item or the legacy history backend for an old event. + +- [ ] **Step 8: Run green integration checks and refactor** + +Run: + +```powershell +pnpm --filter @tanstack/ai-client test:lib -- --run tests/chat-client-interrupts.test.ts tests/chat-client-resume.test.ts tests/chat-client-client-tool-status.test.ts tests/connection-adapters.test.ts +pnpm --filter @tanstack/ai-client test:types +``` + +Expected: PASS; native continuation context has a new run ID/old parent, denial and cancellation differ, sandbox events do not hydrate, and `addToolResult` remains supported. Remove duplicated approval batching logic from ChatClient and rerun both commands. + +### Task 13: Persist raw V2 drafts and recover through explicit adapters + +**Prerequisites:** Tasks 11 and 12. + +**Files:** +- Modify: `packages/ai-client/src/chat-persistence-controller.ts` +- Modify: `packages/ai-client/src/connection-adapters.ts:353-450,574-700,1171-1205` +- Modify: `packages/ai-client/src/types.ts:18-110,397-435` +- Modify: `packages/ai-client/src/chat-client.ts:250-275,485-520,698-706,1393-1410` +- Modify: `packages/ai-client/src/index.ts` +- Modify: `packages/ai-client/tests/chat-persistence-controller.test.ts` +- Modify: `packages/ai-client/tests/chat-client-resume.test.ts` +- Modify: `packages/ai-client/tests/connection-adapters.test.ts` + +- [ ] **Step 1: Add failing V1/V2 reconciliation and recovery tests** + +Test a partial one-of-three draft reload, V1 migration, mismatched thread/run/generation/ID/schema hash, expiry, committed winner, accepted tombstone after failed removal, reset/dispose stale-write suppression, event-extension recovery, fallback fetch, unauthorized/newer generation checks, and no-loader quarantine. + +Use an assertion that persisted values contain no functions or hydrated fields: + +```ts +const stored = JSON.parse(JSON.stringify(resumeAdapter.setItem.mock.calls.at(-1)?.[1])) +expect(stored.schemaVersion).toBe(2) +expect(stored.interruptState.drafts).toHaveLength(1) +expect(JSON.stringify(stored)).not.toContain('resolveInterrupt') +expect(JSON.stringify(stored)).not.toContain('errors') +expect(JSON.stringify(stored)).not.toContain('kind') +``` + +- [ ] **Step 2: Run the red durability tests** + +Run: + +```powershell +pnpm --filter @tanstack/ai-client test:lib -- --run tests/chat-client-resume.test.ts tests/chat-persistence-controller.test.ts tests/connection-adapters.test.ts +``` + +Expected: FAIL because snapshots are unversioned, store raw pending descriptors only, and no recovery loader exists. + +- [ ] **Step 3: Define the raw V2 snapshot and V1 migration** + +Add the exact JSON DTO: + +```ts +export interface ChatResumeSnapshotV2 { + schemaVersion: 2 + resumeState: ChatResumeState + interruptState?: { + provenance: 'native' | 'legacy' + phase: 'pending' | 'accepted' + interruptedRunId: string + generation: number + descriptorsHash: string + descriptors: readonly Interrupt[] + drafts: readonly { + interruptId: string + responseSchemaHash: string + response: RunAgentResumeItem + savedAt: string + }[] + continuationRunId?: string + } +} +``` + +Treat the PR-head unversioned `{ resumeState, pendingInterrupts }` shape as V1. Convert it to V2 generation `0`, native descriptors, no drafts, and `requiresRecovery: true` internally. Remove invalid/unrelatable V1 storage and report a root protocol error. Never serialize methods, validators, configured tools, kind, or error objects. + +- [ ] **Step 4: Reconcile drafts only after authoritative recovery** + +On native hydration, call recovery before binding. Accept only matching `(threadId, interruptedRunId)` and a generation not older than local. For `pending`, bind returned descriptors against the current tool registry, then retain a draft only when generation, interrupt ID, canonical response-schema hash, and expiry all match. For `committed`, clear drafts and attach/join the winning continuation. For `expired`, `missing`, and `legacy-committed`, clear unsafe retry state and expose the specified non-retryable error. Without extension or loader, quarantine drafts, disable submission, and expose `recovery-unavailable`. + +- [ ] **Step 5: Add explicit connection and fetcher recovery contracts** + +Extend connection adapters with an optional method: + +```ts +export interface InterruptRecoveryConnection { + loadInterruptState( + input: InterruptRecoveryQuery, + options: { signal: AbortSignal }, + ): Promise +} + +export type InterruptContinuationLoader = ( + input: { + threadId: string + continuationRunId: string + }, + options: { signal: AbortSignal }, +) => AsyncIterable +``` + +Add `interruptStateFetcher?: InterruptStateFetcher` and `continuationLoader?: InterruptContinuationLoader` beside `fetcher` in fetcher mode. Provide explicit helpers that require caller-supplied URLs: + +```ts +export function createInterruptStateFetcher( + url: string | (() => string), + options: FetchConnectionOptions = {}, +): InterruptStateFetcher { + return async (input, { signal }) => { + const response = await (options.fetchClient ?? fetch)( + typeof url === 'function' ? url() : url, + { + method: 'POST', + headers: { 'Content-Type': 'application/json', ...mergeHeaders(options.headers) }, + credentials: options.credentials ?? 'same-origin', + signal, + body: JSON.stringify(input), + }, + ) + if (!response.ok) throw new Error(`Interrupt recovery failed: ${response.status}`) + return parseInterruptRecoveryState(await response.json()) + } +} + +export function createInterruptContinuationLoader( + url: string | (() => string), + options: FetchConnectionOptions = {}, +): InterruptContinuationLoader { + return async function* loadContinuation(input, { signal }) { + const requestUrl = withSearchParams( + typeof url === 'function' ? url() : url, + { + runId: input.continuationRunId, + offset: '-1', + }, + ) + yield* resumableServerSentEvents( + options.fetchClient ?? fetch, + requestUrl, + { + method: 'GET', + headers: mergeHeaders(options.headers), + credentials: options.credentials ?? 'same-origin', + }, + signal, + ) + } +} +``` + +Implement `parseInterruptRecoveryState(value: unknown)` beside the fetcher with explicit object, version, state, correlation, generation, descriptor, and committed-shape checks; it returns `InterruptRecoveryStateV1` only after narrowing and throws a protocol error otherwise. `mergeHeaders`, `withSearchParams`, and `resumableServerSentEvents` are existing functions in this module. `fetchServerSentEvents` already provides `joinRun`; raw-stream/fetcher adapters expose only explicitly supplied `loadInterruptState` and `continuationLoader`. They never derive either path from the chat URL, and continuation loading is always a read-only GET of the winning run. + +- [ ] **Step 6: Write accepted tombstones before cleanup** + +After acceptance, remove active descriptors/drafts, persist `{ phase: 'accepted', continuationRunId, drafts: [] }`, then queue `removeItem`. If removal fails, the tombstone prevents stale pending UI on reload and cleanup retries after authoritative committed recovery. Thread reset and controller disposal increment generations so old async reads/writes cannot restore state. + +- [ ] **Step 7: Run green durability checks and refactor** + +Run: + +```powershell +pnpm --filter @tanstack/ai-client test:lib -- --run tests/chat-client-interrupts.test.ts tests/chat-client-resume.test.ts tests/chat-persistence-controller.test.ts tests/connection-adapters.test.ts +pnpm --filter @tanstack/ai-client test:types +``` + +Expected: PASS; one draft survives a safe reload, drifted drafts are removed, committed state attaches to the winning run, and no URL is guessed. Extract V1 parsing, V2 cloning, and reconciliation predicates into pure functions and rerun both commands. + +### Task 14: Write every framework behavior and type contract red + +**Prerequisites:** Tasks 12 and 13. This task changes tests only; no framework source or public type is modified until every runtime and type suite has failed for the intended missing API. + +**Files:** + +- Modify: `packages/ai-react/tests/use-chat.test.ts` +- Modify: `packages/ai-react/tests/use-chat-types.test.ts` +- Modify: `packages/ai-preact/tests/use-chat.test.ts` +- Modify: `packages/ai-preact/tests/use-chat-types.test.ts` +- Modify: `packages/ai-solid/tests/use-chat.test.ts` +- Modify: `packages/ai-solid/tests/use-chat-types.test.ts` +- Modify: `packages/ai-vue/tests/use-chat.test.ts` +- Modify: `packages/ai-vue/tests/use-chat-types.test.ts` +- Modify: `packages/ai-svelte/tests/use-chat.test.ts` +- Modify: `packages/ai-svelte/tests/create-chat-types.test.ts` +- Modify: `packages/ai-angular/tests/inject-chat.test.ts` +- Modify: `packages/ai-angular/tests/inject-chat-types.test.ts` + +- [ ] **Step 1: Write the bound-state behavior contract in all six runtime suites** + +In each framework test utility, capture the `ChatClientOptions.onInterruptStateChange` callback and return root spies for `resolveInterrupts`, `cancelInterrupts`, `retryInterrupts`, and `resumeInterruptsUnsafe`. The callback accepts one `ChatInterruptState` object. Use two bound items so calling one synchronous item resolver cannot represent the singleton auto-submit case: + +```ts +it('reactively projects immutable interrupt state and delegates root methods', async () => { + const { result } = renderHook(() => useChat({ tools: [transferTool] })) + const first = createBoundToolInterrupt({ + id: 'approval-1', + status: 'staged', + stagedResponse: { + status: 'resolved', + approved: true, + editedArgs: { amount: 10 }, + }, + }) + const second = createBoundToolInterrupt({ + id: 'approval-2', + status: 'pending', + }) + const staged: ChatInterruptState = { + interrupts: Object.freeze([first, second]), + interruptErrors: [], + isResuming: false, + } + act(() => chatClientCallbacks.onInterruptStateChange?.(staged)) + + expect(result.current.interrupts).toBe(staged.interrupts) + expect(result.current.pendingInterrupts).toBe(staged.interrupts) + expect(result.current.interrupts[0]?.resolveInterrupt(true)).toBeUndefined() + expect(first.resolveInterrupt).toHaveBeenCalledWith(true) + + const invalid = createBoundToolInterrupt({ + id: 'approval-1', + status: 'error', + stagedResponse: first.stagedResponse, + errors: [ + itemError('approval-1', 'invalid-edited-args', ['amount']), + ], + }) + const transport = batchError('transport', ['approval-1', 'approval-2']) + const failed: ChatInterruptState = { + interrupts: Object.freeze([invalid, second]), + interruptErrors: [transport], + isResuming: true, + } + act(() => chatClientCallbacks.onInterruptStateChange?.(failed)) + expect(result.current.interrupts[0]?.status).toBe('error') + expect(result.current.interrupts[0]?.errors).toEqual(invalid.errors) + expect(result.current.interruptErrors).toEqual([transport]) + expect(result.current.isResuming).toBe(true) + + const resolver = (item: ChatInterrupt) => { + item.cancelInterrupt() + return undefined + } + await act(() => result.current.resolveInterrupts(resolver)) + expect(chatClientRootSpies.resolveInterrupts).toHaveBeenCalledWith(resolver) + expect(resolver(second)).toBeUndefined() + await act(() => result.current.cancelInterrupts()) + await act(() => result.current.retryInterrupts()) + expect(chatClientRootSpies.cancelInterrupts).toHaveBeenCalledOnce() + expect(chatClientRootSpies.retryInterrupts).toHaveBeenCalledOnce() +}) +``` + +Define `transferTool` with `toolDefinition({ name: 'transfer', needsApproval: true, inputSchema: z.object({ amount: z.number() }) })`. The test-only `createBoundToolInterrupt`, `itemError`, and `batchError` functions return fully correlated frozen public DTOs; their bound methods are `vi.fn<() => void>()`. They live in each package's existing `tests/test-utils.ts`, not in production. React and Preact use the displayed `renderHook`/their imported `act`; Solid reads `result.current.interrupts()`; Vue reads `result.current.interrupts.value` after `flushPromises()`; Svelte reads `chat.interrupts` after `await tick()`; Angular reads `result.interrupts()` after `TestBed.flushEffects()`. Each file performs every displayed identity, status, error, draft, root-error, and root-spy assertion with that exact native accessor. + +- [ ] **Step 2: Add branch-aware type failures to all six type suites** + +In each type-test file define this local tool and run the assertions against that framework's public result: + +```ts +const transferTool = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + inputSchema: z.object({ amount: z.number(), recipient: z.string() }), + needsApproval: true, + approvalSchema: { + approve: z.object({ note: z.string() }), + reject: z.object({ reason: z.string() }), + }, +}) +const tools = [transferTool] as const + +function assertInterruptTypes( + interrupt: ChatInterrupt | undefined, +): void { + if (!interrupt) return + if (interrupt.kind === 'tool-approval' && interrupt.toolName === 'transfer') { + interrupt.resolveInterrupt(true, { + editedArgs: { amount: 12, recipient: 'Ada' }, + payload: { note: 'Reviewed' }, + }) + interrupt.resolveInterrupt(false, { + payload: { reason: 'Limit exceeded' }, + }) + // @ts-expect-error approval cannot use the rejection payload + interrupt.resolveInterrupt(true, { payload: { reason: 'wrong branch' } }) + // @ts-expect-error editedArgs preserve the input schema + interrupt.resolveInterrupt(true, { + editedArgs: { amount: '12', recipient: 'Ada' }, + payload: { note: 'Reviewed' }, + }) + // @ts-expect-error rejection never accepts editedArgs + interrupt.resolveInterrupt(false, { + editedArgs: { amount: 12, recipient: 'Ada' }, + payload: { reason: 'No' }, + }) + } + if (interrupt.kind === 'generic') { + type GenericPayload = Parameters[0] + expectTypeOf().toEqualTypeOf() + interrupt.resolveInterrupt({ applicationValue: 'runtime validated' }) + } +} +``` + +Call `assertInterruptTypes` with these exact public expressions: React `renderHook(() => useChat({ tools })).result.current.interrupts[0]`; Preact `renderHook(() => useChat({ tools })).result.current.interrupts[0]`; Solid `useChat({ tools }).interrupts()[0]` inside its existing root; Vue `useChat({ tools }).interrupts.value[0]` inside its effect scope; Svelte `createChat({ tools }).interrupts[0]` inside its rune fixture; Angular `injectChat({ tools }).interrupts()[0]` inside its injection context. + +- [ ] **Step 3: Run all twelve framework suites red** + +Run sequentially to cap native Windows memory use: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-react test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-react test:types +pnpm --filter @tanstack/ai-preact test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-preact test:types +pnpm --filter @tanstack/ai-solid test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-solid test:types +pnpm --filter @tanstack/ai-vue test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-vue test:types +pnpm --filter @tanstack/ai-svelte test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-svelte test:types +pnpm --filter @tanstack/ai-angular test:lib -- --run tests/inject-chat.test.ts +pnpm --filter @tanstack/ai-angular test:types +``` + +Expected: all runtime and type suites FAIL because `interrupts`, `resolveInterrupts`, `cancelInterrupts`, `retryInterrupts`, `interruptErrors`, `isResuming`, `resumeInterruptsUnsafe`, and branch-specific inference are not exposed. Save the failing command/output in the task transcript; do not modify framework source in Task 14. + +- [ ] **Step 4: Confirm the red commit boundary** + +Run `git diff --name-only` and confirm Task 14 changed only the twelve named runtime/type test files and the six existing test-utility files. Do not commit. Framework source and public type declarations remain untouched until Task 15, so every type-red assertion predates the implementation it drives. + +### Task 15: Implement reactive and typed framework parity + +**Prerequisites:** Task 14. + +**Files:** + +- Modify: `packages/ai-react/src/use-chat.ts:52,87-212,426-444` +- Modify: `packages/ai-react/src/types.ts:102-173` +- Modify: `packages/ai-preact/src/use-chat.ts:51,86-190,349-367` +- Modify: `packages/ai-preact/src/types.ts:102-173` +- Modify: `packages/ai-solid/src/use-chat.ts:60-70,83-166,323-344` +- Modify: `packages/ai-solid/src/types.ts:8-76,138-159` +- Modify: `packages/ai-vue/src/use-chat.ts:60-64,83-162,328-349` +- Modify: `packages/ai-vue/src/types.ts:116-159` +- Modify: `packages/ai-svelte/src/create-chat.svelte.ts:79-82,98-178,333-376` +- Modify: `packages/ai-svelte/src/types.ts:116-161` +- Modify: `packages/ai-angular/src/inject-chat.ts:57-65,75-114,250-268` +- Modify: `packages/ai-angular/src/types.ts:98-141` + +- [ ] **Step 1: Add the public generic control fields** + +Import all interrupt types from `@tanstack/ai-client` and add this shape to each framework return type: + +```ts +export interface InterruptControls { + interrupts: BoundInterrupts + /** @deprecated Use `interrupts`. Removed in 1.0. */ + pendingInterrupts: BoundInterrupts + resolveInterrupts: ResolveInterrupts + cancelInterrupts(): Promise + retryInterrupts(): Promise + interruptErrors: readonly InterruptSubmissionError[] + isResuming: boolean + resumeInterruptsUnsafe( + entries: readonly RunAgentResumeItem[], + state?: ChatResumeState, + ): Promise + /** @deprecated Use a bound resolver. Removed in 1.0. */ + addToolApprovalResponse(response: { + id: string + approved: boolean + }): Promise + /** @deprecated Use `resumeInterruptsUnsafe`. Removed in 1.0. */ + resumeInterrupts( + entries: readonly RunAgentResumeItem[], + state?: ChatResumeState, + ): Promise +} +``` + +React/Preact expose the plain fields. Solid wraps `interrupts`, `pendingInterrupts`, `interruptErrors`, and `isResuming` in accessors. Vue uses readonly refs, Svelte uses getters, and Angular uses readonly signals. Root methods remain functions in every package. No framework declares its own interrupt conditional type. + +- [ ] **Step 2: Wire React in the existing instance-holder order** + +Place state before `useMemo`, put the callback inside the existing `new ChatClient` options after `getActiveInstance` exists, and synchronize after the instance is active: + +```ts +const [interruptState, setInterruptState] = useState< + ChatInterruptState +>(() => ({ + interrupts: Object.freeze([]), + interruptErrors: Object.freeze([]), + isResuming: false, +})) + +const syncInterruptState = useCallback((target: ChatClient) => { + setInterruptState(target.getInterruptState()) +}, []) + +// Inside the existing useMemo, immediately after getActiveInstance: +const handleInterruptStateChange = (next: ChatInterruptState) => { + if (!getActiveInstance()) return + setInterruptState(next) + optionsRef.current.onInterruptStateChange?.(next) +} + +// Inside the existing new ChatClient options, beside onResumeStateChange: +onInterruptStateChange: handleInterruptStateChange, + +useEffect(() => { + syncInterruptState(client) +}, [client, syncInterruptState]) + +const resolveInterrupts = useCallback>( + (decisionOrResolver) => client.resolveInterrupts(decisionOrResolver), + [client], +) +const cancelInterrupts = useCallback(() => client.cancelInterrupts(), [client]) +const retryInterrupts = useCallback(() => client.retryInterrupts(), [client]) +const resumeInterruptsUnsafe = useCallback( + (entries: readonly RunAgentResumeItem[], state?: ChatResumeState) => + client.resumeInterruptsUnsafe(entries, state), + [client], +) +``` + +Keep the existing `instanceHolder.current = instance` and `activeClientRef.current = instance` assignments before the memo returns. Add `syncInterruptState` to the memo dependency list. Return `interrupts: interruptState.interrupts` and `pendingInterrupts: interruptState.interrupts`, preserving exact array identity. + +- [ ] **Step 3: Wire Preact, Solid, Vue, Svelte, and Angular with concrete native state** + +In Preact, add the callback to its existing constructor closure and set the initial snapshot in the existing post-construction effect: + +```ts +const [interruptState, setInterruptState] = useState>({ + interrupts: Object.freeze([]), + interruptErrors: Object.freeze([]), + isResuming: false, +}) +// inside the existing new ChatClient options: +onInterruptStateChange: (next) => { + if (activeClientRef.current !== instance) return + setInterruptState(next) + optionsRef.current.onInterruptStateChange?.(next) +}, +// inside the existing useEffect([client]): +setInterruptState(client.getInterruptState()) +``` + +In Solid, add one signal before `createMemo`, update it in the constructor callback, and initialize it after `client()` exists: + +```ts +const [interruptState, setInterruptState] = + createSignal>({ + interrupts: Object.freeze([]), + interruptErrors: Object.freeze([]), + isResuming: false, + }) +// inside new ChatClient: +onInterruptStateChange: (next) => { + setInterruptState(next) + options.onInterruptStateChange?.(next) +}, +setInterruptState(client().getInterruptState()) +const resolveInterrupts: ResolveInterrupts = (value) => + client().resolveInterrupts(value) +``` + +In Vue, Svelte, and Angular use these exact declarations/callbacks: + +```ts +// Vue +const interruptState = shallowRef>({ + interrupts: Object.freeze([]), + interruptErrors: Object.freeze([]), + isResuming: false, +}) +// in new ChatClient: +onInterruptStateChange: (next) => { + interruptState.value = next + options.onInterruptStateChange?.(next) +}, +interruptState.value = client.getInterruptState() + +// Svelte +let interruptState = $state>({ + interrupts: Object.freeze([]), + interruptErrors: Object.freeze([]), + isResuming: false, +}) +// in new ChatClient: +onInterruptStateChange: (next) => { + interruptState = next + options.onInterruptStateChange?.(next) +}, +interruptState = client.getInterruptState() + +// Angular +const interruptState = signal>({ + interrupts: Object.freeze([]), + interruptErrors: Object.freeze([]), + isResuming: false, +}) +// in new ChatClient: +onInterruptStateChange: (next) => { + interruptState.set(next) + options.onInterruptStateChange?.(next) +}, +interruptState.set(client.getInterruptState()) +``` + +For each constructor, add `onInterruptStateChange` beside the existing `onResumeStateChange`; preserve every existing callback. Delegate `resolveInterrupts`, `cancelInterrupts`, `retryInterrupts`, and `resumeInterruptsUnsafe` directly to that package's existing `client`/`client()`. Vue returns `readonly(interruptState)` projections, Svelte getters return fields from `interruptState`, and Angular uses `computed(() => interruptState().interrupts)`; deprecated `pendingInterrupts` reads the identical field in all cases. + +- [ ] **Step 4: Run all runtime and type suites green** + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-react test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-react test:types +pnpm --filter @tanstack/ai-preact test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-preact test:types +pnpm --filter @tanstack/ai-solid test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-solid test:types +pnpm --filter @tanstack/ai-vue test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-vue test:types +pnpm --filter @tanstack/ai-svelte test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-svelte test:types +pnpm --filter @tanstack/ai-angular test:lib -- --run tests/inject-chat.test.ts +pnpm --filter @tanstack/ai-angular test:types +``` + +Expected: every command PASS; each `@ts-expect-error` is consumed, generic payload is `unknown`, callbacks observe staged/error/replaced states, root callback side effects return `undefined`, and `pendingInterrupts === interrupts` in plain-value packages (or their unwrapped values are strictly equal in reactive packages). + +### Task 16: Exercise the full interrupt lifecycle in a deterministic browser fixture + +**Prerequisites:** Tasks 5-9 and 13-15. + +**Files:** + +- Create: `testing/e2e/src/lib/interrupts-v2-fixture.ts` +- Create: `testing/e2e/src/routes/api.interrupts-v2.ts` +- Create: `testing/e2e/src/routes/api.interrupts-v2.recovery.ts` +- Create: `testing/e2e/src/routes/interrupts-v2.tsx` +- Create: `testing/e2e/tests/interrupts.spec.ts` +- Modify (generated by the router command): `testing/e2e/src/routeTree.gen.ts` + +- [ ] **Step 1: Write the browser scenarios against stable selectors** + +Create one serial `interrupts.spec.ts` describe block only where a test deliberately shares a thread; all unrelated tests retain Playwright's normal parallelism and use a unique ID: + +```ts +import { expect, test } from '@playwright/test' + +function interruptUrl(testId: string, scenario: string): string { + const search = new URLSearchParams({ testId, scenario }) + return `/interrupts-v2?${search.toString()}` +} + +test('waits for all three cards before one atomic submission', async ({ page }) => { + const testId = `interrupts-v2-batch-${test.info().workerIndex}-${Date.now()}` + await page.goto(interruptUrl(testId, 'three-tool-approvals')) + await page.getByTestId('start-run').click() + await expect(page.getByTestId('interrupt-card')).toHaveCount(3) + + await page.getByTestId('interrupt-card').nth(0).getByRole('button', { name: 'Approve' }).click() + await expect(page.getByTestId('continuation-count')).toHaveText('0') + await page.getByTestId('interrupt-card').nth(1).getByRole('button', { name: 'Deny' }).click() + await expect(page.getByTestId('continuation-count')).toHaveText('0') + await page.getByTestId('interrupt-card').nth(2).getByRole('button', { name: 'Cancel' }).click() + + await expect(page.getByTestId('continuation-count')).toHaveText('1') + await expect(page.getByTestId('submitted-decisions')).toHaveText( + 'approve,deny,cancel', + ) +}) +``` + +Add the singleton omitted/edit, mixed audit history, reload, race, and payloadless scenarios with selectors for decision, staged status, continuation count, result-only event names, and stored history. Add these four tests verbatim for the previously ambiguous cases: + +```ts +test('root callback performs synchronous side effects and returns undefined', async ({ page }) => { + const testId = `interrupt-callback-${test.info().workerIndex}-${Date.now()}` + await page.goto(interruptUrl(testId, 'heterogeneous-callback')) + await page.getByTestId('start-run').click() + await expect(page.getByTestId('interrupt-card')).toHaveCount(3) + await page.getByTestId('resolve-callback').click() + await expect(page.getByTestId('callback-return-values')).toHaveText( + 'undefined,undefined,undefined', + ) + await expect(page.getByTestId('continuation-count')).toHaveText('1') + await expect(page.getByTestId('submitted-decisions')).toHaveText( + 'approve,deny,generic', + ) +}) + +test('shows every item error and the root aggregate in one response', async ({ page }) => { + const testId = `interrupt-errors-${test.info().workerIndex}-${Date.now()}` + await page.goto(interruptUrl(testId, 'two-invalid')) + await page.getByTestId('start-run').click() + await page.getByTestId('invalid-first').click() + await page.getByTestId('invalid-second').click() + await expect(page.getByTestId('interrupt-error-first')).toContainText( + 'invalid-edited-args', + ) + await expect(page.getByTestId('interrupt-error-second')).toContainText( + 'invalid-payload', + ) + await expect(page.getByTestId('interrupt-errors-root')).toContainText( + 'item-validation-failed', + ) + await expect(page.getByTestId('interrupt-errors-root')).toContainText( + 'approval-1,question-1', + ) + await page.getByTestId('correct-first').click() + await expect(page.getByTestId('interrupt-error-first')).toHaveText('') + await expect(page.getByTestId('interrupt-error-second')).toContainText( + 'invalid-payload', + ) + await page.getByTestId('correct-second').click() + await expect(page.getByTestId('continuation-count')).toHaveText('1') +}) + +test('joins the original continuation after a committed response is truncated', async ({ + page, + request, +}) => { + const testId = `interrupt-retry-${test.info().workerIndex}-${Date.now()}` + await page.goto(interruptUrl(testId, 'commit-then-truncate')) + await page.getByTestId('start-run').click() + await page.getByRole('button', { name: 'Approve' }).click() + await expect(page.getByTestId('retry-banner')).toBeVisible() + const committed = await request.get( + `/api/interrupts-v2?testId=${encodeURIComponent(testId)}&stats=1`, + ) + expect(await committed.json()).toMatchObject({ + continuationCount: 1, + continuationRunIds: [expect.any(String)], + truncatedResponses: 1, + }) + await page.getByTestId('retry-interrupts').click() + await expect(page.getByTestId('continuation-count')).toHaveText('1') + const replayed = await request.get( + `/api/interrupts-v2?testId=${encodeURIComponent(testId)}&stats=1`, + ) + const replayStats = await replayed.json() + expect(replayStats).toMatchObject({ + continuationCount: 1, + replayCount: 1, + truncatedResponses: 1, + }) + expect(replayStats.joinedContinuationRunId).toBe( + replayStats.continuationRunIds[0], + ) +}) + +test('keeps client-tool execution typed, visible, and resolvable by its bound method', async ({ + page, +}) => { + const testId = `interrupt-client-tool-${test.info().workerIndex}-${Date.now()}` + await page.goto(interruptUrl(testId, 'client-tool-and-approval')) + await page.getByTestId('start-run').click() + await expect( + page.getByTestId('interrupt-kind-client-tool-execution'), + ).toHaveCount(1) + await expect(page.getByTestId('interrupt-kind-tool-approval')).toHaveCount(1) + await page.getByTestId('resolve-client-tool-bound').click() + await expect(page.getByTestId('client-tool-status')).toHaveText('staged') + await page.getByTestId('approve-server-tool').click() + await expect(page.getByTestId('continuation-count')).toHaveText('1') + await expect(page.getByTestId('client-tool-output')).toHaveText( + '{"browserValue":"done"}', + ) +}) +``` + +The retry scenario does not abort the request before the server sees it. The route consumes the first resume through authoritative commit and continuation scheduling, then intentionally returns a truncated stream once. The second request is an exact replay and the client joins the stored continuation ID. + +- [ ] **Step 2: Run the focused E2E test red** + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai-e2e test:e2e -- tests/interrupts.spec.ts +``` + +Expected: FAIL because `/interrupts-v2`, its chat endpoint, its recovery endpoint, and all stable selectors do not exist. + +- [ ] **Step 3: Build a deterministic adapter and per-test authoritative store** + +In `interrupts-v2-fixture.ts`, create a synthetic `AnyTextAdapter` that emits exact AG-UI chunks rather than depending on a live provider or timing. Its first run emits the scenario descriptors; its resumed run records response order and emits one final text result. Key server state by `testId` so Playwright workers and aimock sequences cannot collide: + +```ts +import { memoryPersistence } from '@tanstack/ai-persistence' +import type { AIPersistence } from '@tanstack/ai-persistence' + +type InterruptFixture = { + persistence: AIPersistence + continuationCount: number + continuationRunIds: string[] + decisions: string[] + truncateCommittedResponseOnce: boolean + truncatedResponses: number + replayCount: number + joinedContinuationRunId?: string +} + +const fixtures = new Map() + +export function getInterruptFixture(testId: string): InterruptFixture { + const existing = fixtures.get(testId) + if (existing) return existing + const created: InterruptFixture = { + persistence: memoryPersistence(), + continuationCount: 0, + continuationRunIds: [], + decisions: [], + truncateCommittedResponseOnce: true, + truncatedResponses: 0, + replayCount: 0, + } + fixtures.set(testId, created) + return created +} +``` + +Define the fixture's server tools with real `toolDefinition` schemas: one shared approval schema, one `{ approve, reject }` branch schema, one editable input schema, and one client tool. Cancellation remains payloadless. Define a generic descriptor with its own Draft 2020-12 `responseSchema`. The synthetic stream must include stable `threadId`, `runId`, `interruptId`, `toolCallId`, `parentRunId`, and three distinct responses so the test asserts wire semantics rather than UI-only state. + +- [ ] **Step 4: Add separate chat and recovery endpoints** + +`api.interrupts-v2.ts` uses the real `AIPersistence` middleware and exposes POST chat, GET continuation replay, and GET test statistics: + +```ts +import { + chat, + chatParamsFromRequestBody, + toServerSentEventsResponse, + type StreamChunk, +} from '@tanstack/ai' +import { withChatPersistence } from '@tanstack/ai-persistence' +import { createFileRoute } from '@tanstack/react-router' +import { + createInterruptFixtureAdapter, + getInterruptFixture, + interruptFixtureTools, +} from '../lib/interrupts-v2-fixture' + +async function collectChunks( + stream: AsyncIterable, +): Promise { + const chunks: StreamChunk[] = [] + for await (const chunk of stream) chunks.push(chunk) + return chunks +} + +function replayChunks(chunks: readonly StreamChunk[]): AsyncIterable { + return { + async *[Symbol.asyncIterator]() { + for (const chunk of chunks) yield chunk + }, + } +} + +function truncatedSseResponse(): Response { + const encoder = new TextEncoder() + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode('event: RUN_STARTED\ndata: {')) + queueMicrotask(() => controller.error(new Error('fixture truncation'))) + }, + }), + { headers: { 'Content-Type': 'text/event-stream' } }, + ) +} + +export const Route = createFileRoute('/api/interrupts-v2')({ + server: { + handlers: { + GET: async ({ request }) => { + const url = new URL(request.url) + const testId = url.searchParams.get('testId') + if (!testId) return new Response('Missing testId', { status: 400 }) + const fixture = getInterruptFixture(testId) + if (url.searchParams.get('stats') === '1') { + return Response.json({ + continuationCount: fixture.continuationCount, + continuationRunIds: fixture.continuationRunIds, + truncatedResponses: fixture.truncatedResponses, + replayCount: fixture.replayCount, + joinedContinuationRunId: fixture.joinedContinuationRunId, + }) + } + const runId = url.searchParams.get('runId') + if (!runId) return new Response('Missing runId', { status: 400 }) + const chunks = fixture.continuationChunks.get(runId) + if (!chunks) return new Response('Unknown run', { status: 404 }) + fixture.joinedContinuationRunId = runId + return toServerSentEventsResponse(replayChunks(chunks)) + }, + POST: async ({ request }) => { + const url = new URL(request.url) + const testId = url.searchParams.get('testId') + const scenario = url.searchParams.get('scenario') + if (!testId || !scenario) { + return new Response('Missing fixture correlation', { status: 400 }) + } + const fixture = getInterruptFixture(testId) + const params = await chatParamsFromRequestBody(await request.json()) + const stream = chat({ + ...params, + adapter: createInterruptFixtureAdapter(scenario, fixture), + tools: interruptFixtureTools, + middleware: [withChatPersistence(fixture.persistence)], + }) as AsyncIterable + + if ( + scenario === 'commit-then-truncate' && + params.resume && + fixture.truncateCommittedResponseOnce + ) { + fixture.truncateCommittedResponseOnce = false + const chunks = await collectChunks(stream) + fixture.truncatedResponses += 1 + fixture.continuationCount += 1 + fixture.continuationRunIds.push(params.runId) + fixture.continuationChunks.set(params.runId, chunks) + return truncatedSseResponse() + } + return toServerSentEventsResponse(stream) + }, + }, + }, +}) +``` + +When the authoritative gateway returns the replay terminal, `createInterruptFixtureAdapter` increments `replayCount` but does not increment `continuationCount` or append another continuation run. Add `continuationChunks: Map` to `InterruptFixture`. + +`api.interrupts-v2.recovery.ts` accepts only `InterruptRecoveryQuery` and uses the fixture persistence selected by the explicit `testId` query: + +```ts +import { createInterruptRecoveryHandler } from '@tanstack/ai-persistence' +import { createFileRoute } from '@tanstack/react-router' +import { getInterruptFixture } from '../lib/interrupts-v2-fixture' + +export const Route = createFileRoute('/api/interrupts-v2/recovery')({ + server: { + handlers: { + POST: async ({ request }) => { + const testId = new URL(request.url).searchParams.get('testId') + if (!testId) return new Response('Missing testId', { status: 400 }) + const fixture = getInterruptFixture(testId) + const handleRecovery = createInterruptRecoveryHandler({ + persistence: fixture.persistence, + authorize: () => ({ allowed: true, includeResolutions: true }), + }) + return handleRecovery(request) + }, + }, + }, +}) +``` + +Never infer `/recovery` or a continuation replay URL from the chat URL in library code; the test page explicitly configures recovery and the SSE connection itself supplies `joinRun`. + +- [ ] **Step 5: Build the test page using only public APIs** + +Configure `useChat` with the fixture's typed tool tuple, `fetchServerSentEvents(() => `/api/interrupts-v2?testId=${encodeURIComponent(testId)}&scenario=${encodeURIComponent(scenario)}`)`, `createInterruptStateFetcher(() => `/api/interrupts-v2/recovery?testId=${encodeURIComponent(testId)}`)`, and a local draft store. Read `testId` and `scenario` from the page search params before constructing the adapters. This is an explicit caller-supplied chat/recovery correlation; `testId` is fixture routing data and is not added to `InterruptRecoveryQuery`. Because `fetchServerSentEvents` preserves existing query parameters when it appends `runId`, the same explicit connection also supplies `joinRun`. Render every `interrupts` item with its `kind`, `errors`, current staged response, edit fields, and per-item resolve/clear controls. Render root controls for callback resolution, retry, and approve/deny/cancel all. Use only public bound methods: + +```tsx +{interrupts.map((interrupt) => ( +
+ + {interrupt.errors.map((issue) => issue.message).join('|')} + + {interrupt.kind === 'tool-approval' ? ( + + ) : ( + + )} +
+))} +``` + +Declare `genericDraft` as `unknown` from the page's JSON editor state and let the built-in validator surface errors. Resolve the server-emitted `client-tool-execution` item through its typed bound `resolveInterrupt(output)` method in the scenario above. Add one separate assertion that the retained `addToolResult` API delegates to the same staged item, rather than creating a duplicate interrupt or a second submission. + +- [ ] **Step 6: Regenerate the route tree and run E2E green** + +The E2E package has no standalone route-generator script; its one-shot Vite build invokes the configured TanStack Router plugin and refreshes `routeTree.gen.ts`. Run: + +```powershell +pnpm --filter @tanstack/ai-e2e build +pnpm --filter @tanstack/ai-e2e test:types +$env:CI='true' +pnpm --filter @tanstack/ai-e2e test:e2e -- tests/interrupts.spec.ts +``` + +Expected: the build exits 0 and changes only `routeTree.gen.ts` as generated route output; type checks and all 11 scenarios PASS with no retries required locally. As the focused refactor, remove only duplicated fixture builders, keep selectors and assertions explicit, and rerun all three commands. + +### Task 17: Publish the API, validation recipe, migration guide, and changesets + +**Prerequisites:** Tasks 1-16. Document only the settled and tested public surface. + +**Files:** + +- Create: `docs/chat/interrupts.md` +- Create: `docs/migration/interrupts.md` +- Modify: `docs/tools/tool-approval.md` +- Modify: `docs/tools/client-tools.md` +- Modify: `docs/chat/persistence.md` +- Modify: `docs/persistence/custom-stores.md` +- Modify: `docs/persistence/migrations.md` +- Modify: `docs/persistence/delivery-durability.md` +- Modify: `docs/architecture/approval-flow-processing.md` +- Modify: `docs/reference/functions/toolDefinition.md` +- Modify: `docs/config.json` +- Create: `.changeset/ag-ui-interrupts.md` + +- [ ] **Step 1: Make the documentation link test fail for the planned pages** + +Add `chat/interrupts` and `migration/interrupts` to the appropriate `docs/config.json` groups before creating the files. Also add a canonical `/chat/interrupts` link to `docs/tools/tool-approval.md` and a `/migration/interrupts` link to `docs/chat/persistence.md`; do not create the destination files yet. Then run: + +```powershell +pnpm test:docs +``` + +Expected: FAIL with missing-page/link errors for both new destinations. `scripts/verify-links.ts` validates Markdown links, not navigation-only `docs/config.json` entries, so the two deliberate links are required for this red step. Do not weaken the verifier. + +- [ ] **Step 2: Write the end-to-end interrupt guide with server and client halves** + +`docs/chat/interrupts.md` must cover: + +- AG-UI descriptor/result lifecycle and the PR-head breaking change; +- tool versus generic interrupt discrimination; +- `responseSchema` as the wire/runtime-validation foundation; +- `approvalSchema`, `{ approve, reject }`, optional `editedArgs`, grouped `payload`, and payloadless cancellation; +- one-item auto-submit, multi-item staging, root callback transactions, bulk operations, `clearResolution`, retry, and `interrupt.errors`/root `interruptErrors`; +- native recovery, V2 raw-draft persistence, idempotency, and tab conflicts; +- the separate client-tool `addToolResult` path. + +Use a current model only when a provider-backed server example is necessary: + +```ts +const stream = chat({ + adapter: openaiText('gpt-5.5'), + messages, + tools: [transferTool], + middleware: [withChatPersistence(persistence)], + threadId, +}) +``` + +The matching client example must show the inferred branch payload and no type casts: + +```ts +const { + interrupts, + resolveInterrupts, + cancelInterrupts, + retryInterrupts, + interruptErrors, + isResuming, + resumeInterruptsUnsafe, +} = useChat({ tools: [transferTool] as const }) + +const interrupt = interrupts.find( + (item) => item.kind === 'tool-approval' && item.toolName === 'transfer', +) +if ( + interrupt?.kind === 'tool-approval' && + interrupt.toolName === 'transfer' +) { + interrupt.resolveInterrupt(true, { + editedArgs: { amount: 12, recipient: 'Ada' }, + payload: { note: 'Reviewed' }, + }) +} +``` + +Keep `as const` (a const assertion) but use no `as SomeType` assertion in docs. + +- [ ] **Step 3: Document converting received JSON Schema before app-level validation** + +Generic payloads stay `unknown` statically. Show users how to turn the received Draft 2020-12 schema into their chosen runtime validator before resolving. Use a real Zod 4 conversion example and narrow the conversion result without a type assertion: + +```ts +import { z } from 'zod' +import type { GenericAGUIInterrupt } from '@tanstack/ai-client' + +function resolveGenericEditorValue( + interrupt: GenericAGUIInterrupt, + editorValue: string, +): readonly string[] { + if (!interrupt.responseSchema) return ['This interrupt has no response schema.'] + + let candidateResponse: unknown + try { + candidateResponse = JSON.parse(editorValue) + } catch { + return ['Enter valid JSON.'] + } + + const responseValidator = z.fromJSONSchema(interrupt.responseSchema) + const parsed = responseValidator.safeParse(candidateResponse) + if (!parsed.success) { + return parsed.error.issues.map( + (issue) => `${issue.path.join('.')}: ${issue.message}`, + ) + } + interrupt.resolveInterrupt(parsed.data) + return [] +} +``` + +Explain that the conversion library returns a runtime schema rather than a wire-derived static application type, the client still performs canonical built-in Draft 2020-12 validation, and server validation remains authoritative. + +- [ ] **Step 4: Write the migration and specialized pages** + +`docs/migration/interrupts.md` maps every deprecated surface to the new API: + +```ts +// Deprecated +await addToolApprovalResponse({ id: approval.id, approved: true }) + +// Current +interrupt.resolveInterrupt(true) +``` + +State explicitly that there is no codemod. Add recipes for batched staging, branch payloads, optional edits, generic unknown payloads, legacy persistence behavior, and opt-in native recovery. + +Update the specialized pages as follows: + +| Page | Required content | +| --- | --- | +| `tools/tool-approval.md` | shorthand, approve/reject payloads, and payloadless cancel examples | +| `tools/client-tools.md` | approval interrupts and client execution are separate axes | +| `chat/persistence.md` | raw drafts only, V2 envelope, recovery before rebinding | +| `persistence/custom-stores.md` | atomic `acceptInterruptBatch` capability and receipt contract | +| `persistence/migrations.md` | new table/columns/indexes for every adapter | +| `persistence/delivery-durability.md` | idempotency key, accepted tombstone, replay rules | +| `architecture/approval-flow-processing.md` | descriptor → validate-all → CAS → continuation → history | +| `reference/functions/toolDefinition.md` | conditional `approvalSchema` signature and inferred overloads | + +- [ ] **Step 5: Update docs metadata exactly once** + +In `docs/config.json`: + +- give both new pages `addedAt: "2026-07-13"`; +- set `updatedAt: "2026-07-13"` on each existing page receiving a content change while preserving its existing `addedAt`; +- add the currently unlisted `architecture/approval-flow-processing` entry under the existing `Advanced` section with its repository creation date `addedAt: "2026-04-15"` and `updatedAt: "2026-07-13"`; +- do not change timestamps for unrelated pages or factual-only fixes. + +- [ ] **Step 6: Add one release changeset and run docs green** + +Create `.changeset/ag-ui-interrupts.md` with minor bumps for `@tanstack/ai`, `@tanstack/ai-client`, `@tanstack/ai-react`, `@tanstack/ai-preact`, `@tanstack/ai-solid`, `@tanstack/ai-vue`, `@tanstack/ai-svelte`, and `@tanstack/ai-angular`; add patch bumps for `@tanstack/ai-persistence`, `@tanstack/ai-persistence-drizzle`, `@tanstack/ai-persistence-prisma`, and `@tanstack/ai-persistence-cloudflare`. Describe the breaking migration and deprecated compatibility surfaces. Do not edit existing changesets. + +Run: + +```powershell +pnpm test:docs +pnpm test:types +``` + +Expected: PASS; all links resolve, examples are syntactically valid under doc checks, timestamps are valid, and public declarations build. Refactor repeated explanation into links to `chat/interrupts` without removing the server/client halves from the main guide. + +### Task 18: Converge review to zero, run CI parity, and make the only implementation commit + +**Prerequisites:** Tasks 0-17 are green. Do not start this task while an implementation or test process is still active. + +**Files:** + +- Review: every file changed by Tasks 1-17 +- Modify: only files required to fix findings + +- [ ] **Step 1: Format the actual diff and inspect its shape** + +Use the repository formatter on changed supported files only, then inspect rather than blindly accepting bulk edits: + +```powershell +$changed = git ls-files --modified --others --exclude-standard +pnpm exec prettier --experimental-cli --ignore-unknown --write $changed +git diff --check +git status --short +git diff --stat +``` + +Expected: Prettier succeeds, `git diff --check` is silent, and the status contains only the paths named in this plan plus generated lockfile, migration, and route-tree files. If PowerShell argument expansion exceeds its limit, format one package group at a time; do not switch shells. + +- [ ] **Step 2: Perform the focused self-review ledger** + +Review the diff against the approved spec and record a pass/fix entry for every row below in the task transcript; do not commit a report file: + +| Ledger item | Required evidence | +| --- | --- | +| Protocol | canonical reasons, descriptor/result/error DTOs, one structured failure event | +| Types | conditional schema availability, three branches, optional edits, generic `unknown`, builder preservation | +| Client | singleton auto-submit, batch wait, replacement/clear, callback rollback and late-call guard, retry statuses | +| Validation | same Draft 2020-12 compiler client/server, all items validated, deterministic paths, no mutation | +| Persistence | exact-set CAS, first writer wins, idempotent replay, committed recovery, migrations/assets | +| Recovery | explicit handler/fetcher, current tool/schema rebinding, V2 draft quarantine and tombstones | +| Continuation | same thread, new run, `parentRunId`, generic middleware seam, result-only history | +| Compatibility | all three old APIs deprecated, legacy read only, no dual writer, sandbox guard, no codemod | +| Frameworks | runtime and type parity across all six packages | +| Documentation | server/client halves, schema conversion, no assertion casts, dates, current model, changeset | + +Search for omissions and accidental placeholders: + +```powershell +rg -n "TO[D]O|FIXME|IMPLEMENT_ME|throw new Error\(['\"]Not implemented" packages testing/e2e docs +rg -n "addToolApprovalResponse|pendingInterrupts|resumeInterrupts" packages/ai-client packages/ai-react packages/ai-preact packages/ai-solid packages/ai-vue packages/ai-svelte packages/ai-angular docs +rg -n "responseSchema|approvalSchema|parentRunId|interruptErrors|clearResolution|retryInterrupts" packages testing/e2e docs +``` + +Expected: no implementation placeholders; every legacy hit has a deprecation/compatibility purpose; every required new symbol has implementation, test, and documentation hits. Also search touched docs/model examples for repository-disallowed model names and replace any such existing line being edited with `gpt-5.5`. + +- [ ] **Step 3: Run the execution-time test-hygiene and debugging audits** + +Ask a `test-hygiene` peer to review new/modified tests for reusable fixtures, type-safe mocks, semantic assertions, deterministic timing, and absence of sleep-based synchronization. Fix every actionable finding and rerun the narrow owning test. + +Ask a `debugging-discipline` peer to audit any unexpected-failure entries retained from Task 0 onward. Each entry must show reproduced symptom, repository prior-art search, root cause, smallest fix, and regression test. If there were no unexpected failures, record that fact in the task transcript; do not fabricate an investigation. + +- [ ] **Step 4: Run the unbiased code-review/fix loop to zero** + +Use the `cr-loop` skill after implementation and self-review. Give reviewers the approved spec, this plan, and the complete diff; require coverage across protocol, type inference, browser state, security/validation, transactional persistence, concurrency/replay, framework parity, tests, and docs. Use xhigh reasoning for investigation and review agents as required by repository instructions. + +Classify every finding with file/line evidence. Fix all actionable findings, rerun the smallest owning command after each fix batch, and repeat independent review until the convergence report has zero actionable findings. Do not dismiss a finding merely because an existing test passes. + +- [ ] **Step 5: Run the complete focused package matrix sequentially** + +Set CI once and run native PowerShell one-shot commands. Await each process before starting the next: + +```powershell +$env:CI='true' +pnpm --filter @tanstack/ai test:lib -- --run tests/interrupts.test.ts tests/tool-definition.test.ts tests/chat-resume-interrupts.test.ts tests/ag-ui-event-order.test.ts +pnpm --filter @tanstack/ai test:types +pnpm --filter @tanstack/ai-persistence test:lib -- --run tests/interrupts.test.ts tests/interrupt-conformance.test.ts tests/middleware.test.ts +pnpm --filter @tanstack/ai-persistence test:types +pnpm --filter @tanstack/ai-persistence-drizzle test:lib +pnpm --filter @tanstack/ai-persistence-drizzle test:types +pnpm --filter @tanstack/ai-persistence-prisma test:lib +pnpm --filter @tanstack/ai-persistence-prisma test:types +pnpm --filter @tanstack/ai-persistence-cloudflare test:lib +pnpm --filter @tanstack/ai-persistence-cloudflare test:types +pnpm --filter @tanstack/ai-client test:lib -- --run tests/interrupt-manager.test.ts tests/chat-client-interrupts.test.ts tests/chat-client-resume.test.ts tests/chat-persistence-controller.test.ts tests/connection-adapters.test.ts +pnpm --filter @tanstack/ai-client test:types +pnpm --filter @tanstack/ai-react test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-react test:types +pnpm --filter @tanstack/ai-preact test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-preact test:types +pnpm --filter @tanstack/ai-solid test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-solid test:types +pnpm --filter @tanstack/ai-vue test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-vue test:types +pnpm --filter @tanstack/ai-svelte test:lib -- --run tests/use-chat.test.ts +pnpm --filter @tanstack/ai-svelte test:types +pnpm --filter @tanstack/ai-angular test:lib -- --run tests/inject-chat.test.ts +pnpm --filter @tanstack/ai-angular test:types +pnpm test:docs +pnpm --filter @tanstack/ai-e2e test:types +pnpm --filter @tanstack/ai-e2e test:e2e -- tests/interrupts.spec.ts +``` + +Expected: every command exits 0. On any unexpected failure, stop the matrix, apply `debugging-discipline`, fix the root cause, rerun the owning command, then restart from the first command whose output could be affected. + +- [ ] **Step 6: Run the mandatory repository-wide pre-PR gates** + +These commands are mandatory even after the focused matrix: + +```powershell +$env:CI='true' +pnpm test:pr +pnpm --filter @tanstack/ai-e2e test:e2e +``` + +Expected: both exit 0. `pnpm test:pr` is the canonical CI-equivalent Nx affected target set, and the full E2E command is mandatory for this behavior change. Do not commit, push, or rely on remote CI if either fails. + +- [ ] **Step 7: Complete the verification checklist and inspect process hygiene** + +Use the `completion-checklist` and `verification-before-completion` skills against fresh Step 5/6 output. Then, at quiescence only, inspect for task-owned processes: + +```powershell +Get-Process node,esbuild,vitest -ErrorAction SilentlyContinue | + Select-Object Id, ProcessName, StartTime, Path +git diff --check +git status --short --branch +``` + +Expected: no task-owned orphan process remains, `git diff --check` is silent, and the branch is `codex/interrupts-full-spec-design`. Never blanket-kill `node.exe`; identify ownership and wait for active commands. Do not run `wsl --shutdown` as part of this task. + +- [ ] **Step 8: Stage and create the single final implementation commit** + +Only after every review and gate is green, stage the planned implementation scope and inspect the staged diff: + +```powershell +$implementationFiles = git ls-files --modified --others --exclude-standard +$implementationFiles +git add -- $implementationFiles +git diff --cached --check +git diff --cached --stat +git status --short +git commit -m "feat: implement durable AG-UI interrupts" +git status --short --branch +``` + +Before `git add`, compare the printed list with the reviewed Task 1-17 status and remove any preserved user-owned path from `$implementationFiles`. Expected: one final commit succeeds, no generated plan/review/report artifact is accidentally staged, and the working tree is clean except for any explicitly preserved user-owned file. Do not add a co-author trailer. Do not push and do not open a PR in this execution; hand the verified local branch and commit back to the user. From f6558ee9b51696b0ace279551296c8594dd47287 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 14 Jul 2026 14:16:37 +0200 Subject: [PATCH 3/6] feat: implement AG-UI interrupt lifecycle --- .changeset/ag-ui-interrupts.md | 24 + docs/architecture/approval-flow-processing.md | 172 +- docs/chat/connection-adapters.md | 3 + docs/chat/interrupts.md | 437 +++++ docs/chat/persistence.md | 107 +- docs/config.json | 31 +- docs/migration/interrupts.md | 323 ++++ docs/persistence/chat-persistence.md | 2 +- docs/persistence/cloudflare.md | 14 +- docs/persistence/controls.md | 13 +- docs/persistence/custom-stores.md | 75 +- docs/persistence/delivery-durability.md | 43 +- docs/persistence/drizzle.md | 11 +- docs/persistence/generation-persistence.md | 6 +- docs/persistence/migrations.md | 43 + docs/persistence/overview.md | 7 +- docs/persistence/sandbox-persistence.md | 9 +- docs/persistence/sql-backends.md | 7 +- docs/reference/functions/toolDefinition.md | 70 +- docs/tools/client-tools.md | 51 +- docs/tools/tool-approval.md | 76 +- .../_database-demo/api.database-demo.ts | 28 +- packages/ai-acp/tests/compatible.test.ts | 2 +- .../ai-acp/tests/sandbox-provisioning.test.ts | 2 +- packages/ai-angular/src/inject-chat.ts | 45 + packages/ai-angular/src/types.ts | 23 + .../tests/inject-chat-types.test.ts | 126 +- packages/ai-angular/tests/inject-chat.test.ts | 82 +- packages/ai-angular/tests/test-utils.ts | 75 + .../tests/anthropic-adapter.test.ts | 42 + packages/ai-client/src/chat-client.ts | 810 +++++++-- .../src/chat-persistence-controller.ts | 15 + packages/ai-client/src/connection-adapters.ts | 257 +++ packages/ai-client/src/index.ts | 23 + packages/ai-client/src/interrupt-manager.ts | 1396 ++++++++++++++++ packages/ai-client/src/types.ts | 183 +- .../tests/chat-client-interrupts.test.ts | 1486 +++++++++++++++++ .../tests/chat-client-resume.test.ts | 74 +- .../tests/connection-adapters.test.ts | 57 + .../tests/interrupt-recovery-adapters.test.ts | 111 ++ .../tests/interrupts-types.test-d.ts | 217 +++ .../ai-durable-stream/src/durable-stream.ts | 12 +- .../tests/text-interactions-adapter.test.ts | 102 +- .../openrouter-responses-adapter.test.ts | 66 +- .../0001_tanstack_ai_interrupt_batches.sql | 68 + .../0001_tanstack_ai_interrupt_batches.sql | 68 + packages/ai-persistence-cloudflare/src/d1.ts | 1001 ++++++++++- .../ai-persistence-cloudflare/src/index.ts | 3 +- .../src/migrations.ts | 6 + .../tests/api-types.test-d.ts | 9 + .../tests/migration-cli.test.ts | 14 +- .../tests/migrations.test.ts | 136 +- .../tests/runtime.conformance.test.ts | 126 +- .../0001_tanstack_ai_interrupt_batches.sql | 68 + .../drizzle/meta/0001_snapshot.json | 446 +++++ .../drizzle/meta/_journal.json | 7 + .../0001_tanstack_ai_interrupt_batches.sql | 68 + packages/ai-persistence-drizzle/src/index.ts | 66 +- .../ai-persistence-drizzle/src/migrations.ts | 6 + packages/ai-persistence-drizzle/src/schema.ts | 82 +- packages/ai-persistence-drizzle/src/sqlite.ts | 42 +- packages/ai-persistence-drizzle/src/stores.ts | 606 ++++++- .../tests/api-types.test-d.ts | 38 +- .../tests/drizzle.conformance.test.ts | 128 +- .../tests/migration-cli.test.ts | 13 +- .../tests/migrations.test.ts | 98 +- .../prisma/tanstack-ai.prisma | 21 + .../src/assets/tanstack-ai.prisma | 21 + packages/ai-persistence-prisma/src/stores.ts | 740 +++++++- .../tests/api-types.test-d.ts | 9 + .../tests/models-cli.test.ts | 2 + .../tests/models.test.ts | 10 + .../tests/prisma.conformance.test.ts | 296 +++- packages/ai-persistence/src/capabilities.ts | 38 +- packages/ai-persistence/src/index.ts | 33 +- packages/ai-persistence/src/interrupts.ts | 118 +- packages/ai-persistence/src/memory.ts | 369 +++- packages/ai-persistence/src/middleware.ts | 1006 +++++++++-- packages/ai-persistence/src/recovery.ts | 179 ++ .../ai-persistence/src/testkit/conformance.ts | 268 ++- packages/ai-persistence/src/types.ts | 51 +- .../ai-persistence/tests/interrupts.test.ts | 1039 ++++++------ .../tests/memory.conformance.test.ts | 55 +- .../tests/persistence-fixtures.ts | 16 + .../tests/persistence-validation.test.ts | 25 + .../tests/with-persistence.test.ts | 275 ++- packages/ai-preact/src/types.ts | 21 +- packages/ai-preact/src/use-chat.ts | 116 +- packages/ai-preact/tests/test-utils.ts | 75 + .../ai-preact/tests/use-chat-types.test.ts | 141 +- packages/ai-preact/tests/use-chat.test.ts | 85 +- packages/ai-react/src/types.ts | 21 +- packages/ai-react/src/use-chat.ts | 66 +- packages/ai-react/tests/test-utils.ts | 74 + .../ai-react/tests/use-chat-types.test.ts | 138 ++ packages/ai-react/tests/use-chat.test.ts | 85 +- packages/ai-solid/src/types.ts | 21 +- packages/ai-solid/src/use-chat.ts | 65 +- packages/ai-solid/tests/test-utils.ts | 75 + .../ai-solid/tests/use-chat-types.test.ts | 140 ++ packages/ai-solid/tests/use-chat.test.ts | 84 +- packages/ai-svelte/src/create-chat.svelte.ts | 63 +- packages/ai-svelte/src/types.ts | 21 +- .../ai-svelte/tests/create-chat-types.test.ts | 139 +- packages/ai-svelte/tests/test-utils.ts | 76 + packages/ai-svelte/tests/use-chat.test.ts | 84 +- packages/ai-vue/src/types.ts | 23 +- packages/ai-vue/src/use-chat.ts | 60 +- packages/ai-vue/tests/test-utils.ts | 84 +- packages/ai-vue/tests/use-chat-types.test.ts | 139 +- packages/ai-vue/tests/use-chat.test.ts | 84 +- packages/ai/package.json | 3 + packages/ai/src/activities/chat/index.ts | 573 ++++++- .../ai/src/activities/chat/mcp/manager.ts | 8 +- packages/ai/src/activities/chat/mcp/types.ts | 4 +- .../src/activities/chat/middleware/index.ts | 1 + .../src/activities/chat/middleware/types.ts | 14 +- .../activities/chat/tools/approval-schema.ts | 205 +++ .../chat/tools/json-schema-validator.ts | 138 ++ .../src/activities/chat/tools/tool-calls.ts | 73 +- .../activities/chat/tools/tool-definition.ts | 242 ++- packages/ai/src/client.ts | 28 + packages/ai/src/index.ts | 27 + packages/ai/src/interrupt-serialization.ts | 68 + packages/ai/src/interrupts.ts | 200 +++ packages/ai/src/types.ts | 27 +- packages/ai/tests/chat-mcp-manager.test.ts | 12 +- packages/ai/tests/chat.test.ts | 464 +++-- packages/ai/tests/interrupts-types.test-d.ts | 74 + packages/ai/tests/interrupts.test.ts | 234 +++ packages/ai/tests/stream-processor.test.ts | 40 +- pnpm-lock.yaml | 17 +- testing/e2e/src/lib/interrupts-v2-fixture.ts | 393 +++++ testing/e2e/src/routeTree.gen.ts | 73 + .../src/routes/api.interrupts-v2.recovery.ts | 24 + testing/e2e/src/routes/api.interrupts-v2.ts | 235 +++ testing/e2e/src/routes/interrupts-v2.tsx | 552 ++++++ testing/e2e/tests/interrupts.spec.ts | 270 +++ 138 files changed, 18894 insertions(+), 1513 deletions(-) create mode 100644 .changeset/ag-ui-interrupts.md create mode 100644 docs/chat/interrupts.md create mode 100644 docs/migration/interrupts.md create mode 100644 packages/ai-client/src/interrupt-manager.ts create mode 100644 packages/ai-client/tests/chat-client-interrupts.test.ts create mode 100644 packages/ai-client/tests/interrupt-recovery-adapters.test.ts create mode 100644 packages/ai-client/tests/interrupts-types.test-d.ts create mode 100644 packages/ai-persistence-cloudflare/migrations/0001_tanstack_ai_interrupt_batches.sql create mode 100644 packages/ai-persistence-cloudflare/src/assets/0001_tanstack_ai_interrupt_batches.sql create mode 100644 packages/ai-persistence-drizzle/drizzle/0001_tanstack_ai_interrupt_batches.sql create mode 100644 packages/ai-persistence-drizzle/drizzle/meta/0001_snapshot.json create mode 100644 packages/ai-persistence-drizzle/src/assets/0001_tanstack_ai_interrupt_batches.sql create mode 100644 packages/ai-persistence/src/recovery.ts create mode 100644 packages/ai/src/activities/chat/tools/approval-schema.ts create mode 100644 packages/ai/src/activities/chat/tools/json-schema-validator.ts create mode 100644 packages/ai/src/interrupt-serialization.ts create mode 100644 packages/ai/src/interrupts.ts create mode 100644 packages/ai/tests/interrupts-types.test-d.ts create mode 100644 testing/e2e/src/lib/interrupts-v2-fixture.ts create mode 100644 testing/e2e/src/routes/api.interrupts-v2.recovery.ts create mode 100644 testing/e2e/src/routes/api.interrupts-v2.ts create mode 100644 testing/e2e/src/routes/interrupts-v2.tsx create mode 100644 testing/e2e/tests/interrupts.spec.ts diff --git a/.changeset/ag-ui-interrupts.md b/.changeset/ag-ui-interrupts.md new file mode 100644 index 000000000..af939a21e --- /dev/null +++ b/.changeset/ag-ui-interrupts.md @@ -0,0 +1,24 @@ +--- +'@tanstack/ai': minor +'@tanstack/ai-client': minor +'@tanstack/ai-react': minor +'@tanstack/ai-preact': minor +'@tanstack/ai-solid': minor +'@tanstack/ai-vue': minor +'@tanstack/ai-svelte': minor +'@tanstack/ai-angular': minor +'@tanstack/ai-persistence': patch +'@tanstack/ai-persistence-drizzle': patch +'@tanstack/ai-persistence-prisma': patch +'@tanstack/ai-persistence-cloudflare': patch +--- + +Adopt the AG-UI interrupt lifecycle for tool approvals, generic responses, and +client-tool execution, with typed bound resolvers, atomic batches, structured +errors, persistence-backed recovery, and exact continuation replay. + +This changes native approval and client-tool streams from legacy custom events +to snapshot-plus-`RUN_FINISHED` interrupt outcomes. Deprecated +`pendingInterrupts`, `addToolApprovalResponse`, raw `resumeInterrupts`, and +legacy event readers remain as limited compatibility surfaces for migration; +`addToolResult` remains supported. diff --git a/docs/architecture/approval-flow-processing.md b/docs/architecture/approval-flow-processing.md index 7838c9b14..7ec204e36 100644 --- a/docs/architecture/approval-flow-processing.md +++ b/docs/architecture/approval-flow-processing.md @@ -16,7 +16,7 @@ Tool approval is an interrupt-and-resume protocol. A run that needs user input ends with one canonical event: ```ts -{ +const interruptTerminal = { type: 'RUN_FINISHED', runId: 'run-1', threadId: 'thread-1', @@ -26,17 +26,28 @@ ends with one canonical event: interrupts: [ { id: 'approval-1', - reason: 'approval_required', - metadata: { kind: 'approval' }, + reason: 'tool_call', + toolCallId: 'call-1', + responseSchema: { + oneOf: [ + { type: 'object', properties: { approved: { const: true } } }, + { type: 'object', properties: { approved: { const: false } } }, + ], + }, }, ], }, } ``` -The canonical event stream is the only approval event stream. Server state -persistence stores the interrupt record; SSE delivery durability separately -assigns opaque resume offsets to delivered events. +The canonical event stream is the only native approval event stream. Server +state persistence stores the complete descriptor/binding batch before the +terminal is exposed; SSE delivery durability separately assigns opaque resume +offsets to delivered events. Native paths do not emit `approval-requested` or +`tool-input-available` custom events. + +See [Interrupts](../chat/interrupts) for the public server/client guide and +[Migrate to AG-UI interrupts](../migration/interrupts) for deprecated readers. ## Responsibilities @@ -44,11 +55,39 @@ assigns opaque resume offsets to delivered events. | --- | --- | | Tool definition | Declares `needsApproval: true` for a sensitive operation. | | Chat engine | Stops before tool execution and emits the interrupt outcome. | -| Chat persistence middleware | Creates pending interrupt records, marks the run interrupted, and snapshots messages. | -| Chat client | Exposes `pendingInterrupts`, preserves `resumeState`, and sends typed resume entries. | -| Application UI | Explains the operation and resolves or cancels each interrupt. | +| Chat persistence middleware | Atomically opens the descriptor/binding batch, marks the run interrupted, and snapshots messages. | +| Chat client | Binds descriptors to typed methods, stages drafts, and submits one exact resume batch. | +| Application UI | Explains the operation and uses `resolveInterrupt`, `cancel`, or root batch controls. | | Delivery adapter | Optionally replays SSE events by opaque adapter-owned offsets. | +## Descriptor to continuation pipeline + +The invariant is **descriptor → validate all → compare-and-swap → +continuation → history**: + +1. The engine builds public descriptors and protected bindings, then requires + an atomic persistence capability. +2. Persistence opens the entire batch and assigns its generation before output + includes `MESSAGES_SNAPSHOT`, optional `STATE_SNAPSHOT`, and the interrupt + `RUN_FINISHED` terminal. +3. The client binds only descriptors whose reason, tool identity, call ID, + schema hashes, interrupted run, and generation match its tool registry. + Anything untrusted degrades to `generic` rather than gaining a typed tool + resolver. +4. Item methods validate and stage local drafts. The submit boundary contains + every pending interrupt ID exactly once. +5. The server validates **all** payloads, edited inputs, outputs, expiry, hashes, + and correlation before mutating state. +6. One transaction compares the current interrupted run and generation, stores + the canonical resolution fingerprint, creates a fresh continuation whose + `parentRunId` is the interrupted run, and records the receipt. +7. Resumed tool calls emit results only; they do not replay synthetic tool-call + start/argument events. Successful history belongs to the continuation run. + +An exact retry returns the recorded continuation. A stale or different +submission returns authoritative recovery state. Neither path re-executes an +approved tool. + ## Server setup Define the tool and add state persistence before handling requests: @@ -58,13 +97,15 @@ Define the tool and add state persistence before handling requests: import { toolDefinition } from '@tanstack/ai' import { z } from 'zod' -export const deleteProject = toolDefinition({ +export const deleteProjectDefinition = toolDefinition({ name: 'delete_project', description: 'Delete a project permanently', inputSchema: z.object({ projectId: z.string() }), outputSchema: z.object({ deleted: z.boolean() }), needsApproval: true, -}).server(async ({ projectId }) => { +}) + +export const deleteProject = deleteProjectDefinition.server(async ({ projectId }) => { await deleteProjectFromDatabase(projectId) return { deleted: true } }) @@ -117,60 +158,51 @@ A single approval follows this sequence: 1. The model emits a tool call. 2. The client tool-call part reaches `approval-requested`. 3. The run ends with `RUN_FINISHED.outcome.type === 'interrupt'`. -4. `useChat` exposes the descriptor in `pendingInterrupts` and retains - `{ threadId, runId }` in `resumeState`. -5. The UI calls `resumeInterrupts(...)` with a resolution for every pending - interrupt. -6. The next request carries those values in `RunAgentInput.resume`. -7. Persistence validates and resolves the stored records before the engine - continues the tool call. +4. `useChat` exposes a bound item in `interrupts`. +5. The UI calls `resolveInterrupt(...)` or `cancel()`; a singleton + submits immediately, while a multi-item batch waits for every valid draft. +6. The next request carries a fresh `runId`, the interrupted `parentRunId`, and + the exact AG-UI `resume` array. +7. Persistence validates the full set and commits it atomically before the + engine continues the tool call. Normal input is rejected at step 4. This prevents a second branch from being created while the existing run still waits for a decision. -## Real React approval UI +## React approval UI -Use the values returned by `useChat`. Rendering the actual -`pendingInterrupts` array keeps ids and types connected to the hook that owns -the run. +Use the bound values returned by `useChat`. Rendering `interrupts` keeps IDs, +tool types, drafts, and errors connected to the hook that owns the run. -```tsx +```tsx group=approval-ui +import type { ItemInterruptError } from '@tanstack/ai' import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' +import { deleteProjectDefinition } from './tools' export function ApprovalQueue() { const chat = useChat({ id: 'project-chat', threadId: 'project-thread', connection: fetchServerSentEvents('/api/chat'), + tools: [deleteProjectDefinition] as const, }) return (
- {chat.pendingInterrupts.map((interrupt) => ( + {chat.interrupts.map((interrupt) => (

Approval required: {interrupt.reason}

- - + {interrupt.kind === 'tool-approval' ? ( + + ) : null} + + {interrupt.errors.map((error: ItemInterruptError) => ( +

+ {error.message} +

+ ))}
))}
@@ -178,38 +210,30 @@ export function ApprovalQueue() { } ``` -For a batch, send one entry per pending interrupt in one call: - -```tsx -import type { RunAgentResumeItem } from '@tanstack/ai/client' - -function resolveInterrupt( - interruptId: string, - approved: boolean, -): RunAgentResumeItem { - return approved - ? { - interruptId, - status: 'resolved', - payload: { approved: true }, - } - : { interruptId, status: 'cancelled' } -} +For a batch, stage every resolution in one synchronous root callback: +```tsx group=approval-ui function ResolveAll({ approved }: { approved: boolean }) { const chat = useChat({ threadId: 'project-thread', connection: fetchServerSentEvents('/api/chat'), + tools: [deleteProjectDefinition] as const, }) return ( + + + + {transfer.errors.map((error: ItemInterruptError) => ( +

+ {error.message} +

+ ))} + + ) : null} + + {interruptErrors.map((error) => ( +

{error.message}

+ ))} + + + + + + ) +} +``` + +For approval, `true` without `editedArgs` executes the original input. When +provided, `editedArgs` is a validated **full replacement**, not a partial merge. +Only approval can edit arguments. Rejection options never accept `editedArgs`. + +Branch data always belongs under `payload`: + +```ts ignore +interrupt.resolveInterrupt(true, { + editedArgs: { amount: 12, recipient: 'Ada' }, + payload: { note: 'Reviewed' }, +}) + +interrupt.resolveInterrupt(false, { + payload: { reason: 'Policy limit' }, +}) +``` + +The selected `approve` or `reject` schema validates only its nested `payload`. +The server validates the complete envelope again before executing anything. + +### Rejection is not cancellation + +`resolveInterrupt(false, ...)` is a resolved denial. The continuation receives +the denial and its optional rejection payload, so the model can respond to the +decision. `cancelInterrupt()` produces AG-UI `status: 'cancelled'`, carries no +payload, and does not select the reject branch. + +Use denial when the user has answered "no." Use cancellation when the workflow +is being abandoned without an answer. + +## Singleton and batch submission + +Item methods return `void` because they stage local state. A valid resolution +for a singleton batch submits automatically. In a multi-item native batch, +each `resolveInterrupt` stages its item and the last valid item automatically +submits the entire batch. The server accepts all resolutions atomically or none +of them. + +For one synchronous transaction, use root `resolveInterrupts` and resolve every +item inside its callback: + +```ts ignore +await resolveInterrupts((interrupt) => { + if ( + interrupt.kind === 'tool-approval' && + interrupt.toolName === 'transfer' + ) { + interrupt.resolveInterrupt(true, { + editedArgs: interrupt.originalArgs, + payload: { note: 'Approved in batch review' }, + }) + return + } + + interrupt.cancel() +}) +``` + +The callback must synchronously produce a valid resolution or cancellation for +every item. It cannot be async. If it throws or leaves an item unresolved, no +partial submission occurs. + +The boolean shorthand `resolveInterrupts(true)` or +`resolveInterrupts(false)` is only valid when every item is a tool approval and +the selected branches require no payload or edits. It is rejected for generic +interrupts, client-tool outputs, mixed batches, or required branch payloads. +`cancelInterrupts()` is the payloadless all-items operation. + +`clearResolution()` removes an item's staged draft and validation errors. +`retryInterrupts()` retries the current complete batch after a retryable root +transport or server error; it does not invent missing item responses. + +## Validate a generic AG-UI response schema + +A generic AG-UI descriptor may carry a Draft 2020-12 `responseSchema`. Because +the schema arrives over the wire, its application payload remains `unknown` +statically. Convert the received schema to a runtime validator before calling +the bound resolver. + +This example uses the real Zod 4 JSON Schema converter and never asserts a +wire-derived static type: + +```ts +import { z } from 'zod' +import type { GenericAGUIInterrupt } from '@tanstack/ai-client' + +export function resolveGenericEditorValue( + interrupt: GenericAGUIInterrupt, + editorValue: string, +): readonly string[] { + if (!interrupt.responseSchema) { + return ['This interrupt has no response schema.'] + } + + let candidateResponse: unknown + try { + candidateResponse = JSON.parse(editorValue) + } catch { + return ['Enter valid JSON.'] + } + + const responseValidator = z.fromJSONSchema(interrupt.responseSchema) + const parsed = responseValidator.safeParse(candidateResponse) + if (!parsed.success) { + return parsed.error.issues.map( + (issue) => `${issue.path.join('.')}: ${issue.message}`, + ) + } + + interrupt.resolveInterrupt(parsed.data) + return [] +} +``` + +The conversion produces a runtime schema, not a trustworthy static application +type. TanStack AI still performs canonical Draft 2020-12 validation in the +client, and server validation remains authoritative. An invalid or unsupported +wire schema surfaces an item or root `invalid-response-schema` error rather +than bypassing validation. + +Applications that emit generic descriptors must persist the complete batch +before exposing the interrupt terminal and must send JSON-compatible Draft +2020-12 schemas. Tool approvals usually prefer shared `approvalSchema` because +it adds compile-time branch inference as well as runtime validation. + +## Client-tool execution is a separate interrupt + +Approval and execution are independent axes. A client tool with +`needsApproval: true` first produces a `tool-approval` decision. If approved, +the server can later produce a `client-tool-execution` item for the browser +result. + +```ts ignore +const execution = interrupts.find( + (item) => + item.kind === 'client-tool-execution' && item.toolName === 'showReceipt', +) + +if ( + execution?.kind === 'client-tool-execution' && + execution.toolName === 'showReceipt' +) { + execution.resolveInterrupt({ displayed: true }) +} +``` + +The output is validated against the tool's output schema. The existing +`addToolResult` API remains supported: for a matching native item it delegates +to the same staged client-tool resolution, while legacy streams retain their +historical result path. + +## Errors and retry + +Each item exposes `errors` for payload, edited-argument, output, expiry, and +binding failures. Root `interruptErrors` reports batch, transport, persistence, +stale-generation, conflict, and recovery failures. `canResolve` is false when +an item is expired, submitting, or quarantined for recovery. + +A failed candidate does not erase the last valid staged response. Correct the +form and call `resolveInterrupt` again, call `clearResolution()` to start over, +or call `retryInterrupts()` after a retryable submission failure. Non-retryable +stale and conflict errors require authoritative recovery. + +## Persistence, drafts, recovery, and tab conflicts + +Server state and browser drafts have different authority: + +- `withChatPersistence` atomically opens the descriptor/binding batch and + commits the exact resolution set with compare-and-swap semantics. +- Browser `ChatResumeSnapshotV2` stores raw JSON-safe drafts only. It never + stores bound methods, validators, tool implementations, or hydrated errors. +- On reload, authoritative recovery runs before descriptors are rebound and + before a draft can be restored. +- Submission fingerprints make an exact retry idempotent. The same accepted + batch replays or joins the winning continuation instead of executing tools + again. +- A second tab with a stale generation loses the compare-and-swap. It replaces + local state from recovery or joins the already-committed continuation; it + never starts a second tool execution. + +Recovery routes are opt-in and explicit. TanStack AI never infers a recovery +or continuation URL from the chat URL. Configure the connection with the +application-owned endpoints: + +```ts +import { + createInterruptContinuationLoader, + createInterruptStateFetcher, + fetchServerSentEvents, +} from '@tanstack/ai-client' + +const connection = fetchServerSentEvents('/api/chat', { + interruptStateFetcher: createInterruptStateFetcher( + '/api/interrupts/recovery', + ), + continuationLoader: createInterruptContinuationLoader( + '/api/interrupts/continuation', + ), +}) +``` + +The recovery endpoint must authenticate the caller, authorize the requested +thread, and return the application store's authoritative recovery DTO. The +continuation endpoint is a read-only GET of the winning run. See +[Chat persistence](../persistence/chat-persistence), +[Custom stores](../persistence/custom-stores), and +[Delivery durability](../persistence/delivery-durability). + +`resumeInterruptsUnsafe` remains a distinct low-level escape hatch for already +validated raw AG-UI resume entries and explicit recovery tooling. It is not the +replacement for normal approval UI; prefer bound item methods and root batch +controls. + +## Next steps + +- [Tool approval flow](../tools/tool-approval) focuses on approval forms. +- [Client tools](../tools/client-tools) explains browser execution. +- [Migrate to AG-UI interrupts](../migration/interrupts) maps deprecated APIs. +- [Approval flow architecture](../architecture/approval-flow-processing) + details validate-all, CAS, continuation, and replay processing. diff --git a/docs/chat/persistence.md b/docs/chat/persistence.md index dfa146104..d86828799 100644 --- a/docs/chat/persistence.md +++ b/docs/chat/persistence.md @@ -18,14 +18,15 @@ keywords: `useChat` can persist two browser-side values independently: - `persistence.client` stores the rendered `UIMessage[]` under the chat `id`. -- `persistence.server` stores the latest `ChatResumeSnapshot` under the - `threadId`. The snapshot contains the run identity and any pending - interrupts needed after a reload. +- `persistence.server` stores the latest `ChatResumeSnapshotV2` under the + `threadId`. The snapshot contains run identity, authoritative recovery + correlation, and raw JSON-safe interrupt drafts. This is client hydration, not server state durability. Persist authoritative messages, runs, and interrupts with `withChatPersistence(...)`; make an in-flight response replayable with SSE delivery durability. See -[Chat Persistence](../persistence/chat-persistence) for the complete flow. +[Chat Persistence](../persistence/chat-persistence) for the complete flow and +[Interrupts](./interrupts) for resolution and recovery semantics. ## Use IndexedDB for chat messages @@ -121,10 +122,12 @@ const messages = localStoragePersistence>({ serialize: serializeJson, deserialize(value) { const stored: Array = JSON.parse(value) - return stored.map((message) => ({ - ...message, - ...(message.createdAt - ? { createdAt: new Date(message.createdAt) } + return stored.map(({ id, role, parts, createdAt }) => ({ + id, + role, + parts, + ...(createdAt + ? { createdAt: new Date(createdAt) } : {}), })) }, @@ -153,57 +156,85 @@ export function PersistentChat() { `localStorage` survives browser restarts. `sessionStorage` is scoped to the current tab. Both adapters default to the key prefix `tanstack-ai:`. -## Resume pending interrupts after reload +## Recover interrupts before rebinding drafts -The hook exposes the persisted interrupt state directly. Render -`pendingInterrupts` and call the returned `resumeInterrupts` function; do not -manufacture a hook return type outside a component. +The browser stores raw V2 drafts, not hydrated bound items. It never serializes +`resolveInterrupt`, validators, configured tools, discriminated `kind`, or +hydrated error objects. On reload, authoritative server recovery must confirm +the thread, interrupted run, generation, exact IDs, schema hashes, expiry, and +commit state before the client rebinds descriptors or restores a draft. + +Render `interrupts` and use their bound methods. `pendingInterrupts` and raw +`resumeInterrupts` remain deprecated compatibility surfaces. ```tsx +import { toolDefinition } from '@tanstack/ai' import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' +const approvalTool = toolDefinition({ + name: 'sensitive_action', + description: 'Perform a sensitive action', + needsApproval: true, +}) + export function Interrupts() { - const { pendingInterrupts, resumeInterrupts } = useChat({ + const { interrupts, interruptErrors, retryInterrupts } = useChat({ threadId: 'thread-1', connection: fetchServerSentEvents('/api/chat'), + tools: [approvalTool] as const, }) return (
    - {pendingInterrupts.map((interrupt) => ( + {interrupts.map((interrupt) => (
  • {interrupt.reason} - - + {interrupt.kind === 'tool-approval' ? ( + + ) : null} +
  • ))} + {interruptErrors.map((error) => ( +
  • {error.message}
  • + ))} +
  • + +
) } ``` -The server still validates that every resume entry matches a pending -interrupt. Normal new input is rejected while interrupts remain unresolved. +The server still validates every entry and commits the exact batch atomically. +Normal new input is rejected while interrupts remain unresolved. + +Recovery is opt-in. Configure explicit application-owned recovery and winning +continuation URLs; connection adapters never infer them from the chat URL: + +```ts +import { + createInterruptContinuationLoader, + createInterruptStateFetcher, + fetchServerSentEvents, +} from '@tanstack/ai-client' + +export const connection = fetchServerSentEvents('/api/chat', { + interruptStateFetcher: createInterruptStateFetcher( + '/api/interrupts/recovery', + ), + continuationLoader: createInterruptContinuationLoader( + '/api/interrupts/continuation', + ), +}) +``` + +Exact retries use the stored fingerprint and join the winning continuation. +Stale tabs recover the authoritative generation instead of executing a second +continuation. See [Migrate to AG-UI interrupts](../migration/interrupts) for V1 +compatibility and rollout steps. ## SSR and custom adapters diff --git a/docs/config.json b/docs/config.json index a91934a2f..6eeb8d240 100644 --- a/docs/config.json +++ b/docs/config.json @@ -102,13 +102,13 @@ "label": "Client Tools", "to": "tools/client-tools", "addedAt": "2026-04-15", - "updatedAt": "2026-07-08" + "updatedAt": "2026-07-14" }, { "label": "Tool Approval Flow", "to": "tools/tool-approval", "addedAt": "2026-04-15", - "updatedAt": "2026-07-09" + "updatedAt": "2026-07-14" }, { "label": "Lazy Tool Discovery", @@ -174,11 +174,16 @@ "to": "chat/thinking-content", "addedAt": "2026-04-15" }, + { + "label": "Interrupts", + "to": "chat/interrupts", + "addedAt": "2026-07-14" + }, { "label": "Persistence", "to": "chat/persistence", "addedAt": "2026-06-02", - "updatedAt": "2026-07-10" + "updatedAt": "2026-07-14" } ] }, @@ -417,7 +422,7 @@ "label": "Delivery Durability", "to": "persistence/delivery-durability", "addedAt": "2026-07-09", - "updatedAt": "2026-07-10" + "updatedAt": "2026-07-14" }, { "label": "Persistence Controls", @@ -429,7 +434,7 @@ "label": "Migrations", "to": "persistence/migrations", "addedAt": "2026-07-08", - "updatedAt": "2026-07-10" + "updatedAt": "2026-07-14" }, { "label": "Chat Persistence", @@ -483,7 +488,7 @@ "label": "Custom Stores", "to": "persistence/custom-stores", "addedAt": "2026-07-03", - "updatedAt": "2026-07-10" + "updatedAt": "2026-07-14" }, { "label": "Persistence Internals", @@ -536,6 +541,12 @@ "label": "Typed Pre-Configured Options", "to": "advanced/typed-options", "addedAt": "2026-05-25" + }, + { + "label": "Approval Flow Processing", + "to": "architecture/approval-flow-processing", + "addedAt": "2026-04-15", + "updatedAt": "2026-07-14" } ] }, @@ -559,6 +570,11 @@ "addedAt": "2026-05-16", "updatedAt": "2026-07-08" }, + { + "label": "AG-UI Interrupts", + "to": "migration/interrupts", + "addedAt": "2026-07-14" + }, { "label": "Sampling \u2192 modelOptions", "to": "migration/sampling-options-to-model-options", @@ -901,7 +917,8 @@ { "label": "toolDefinition", "to": "reference/functions/toolDefinition", - "addedAt": "2026-04-15" + "addedAt": "2026-04-15", + "updatedAt": "2026-07-14" }, { "label": "uiMessageToModelMessages", diff --git a/docs/migration/interrupts.md b/docs/migration/interrupts.md new file mode 100644 index 000000000..73a6f2057 --- /dev/null +++ b/docs/migration/interrupts.md @@ -0,0 +1,323 @@ +--- +title: Migrate to AG-UI Interrupts +id: migrate-interrupts +order: 8 +description: "Migrate TanStack AI approval and resume code from legacy custom events and raw resume APIs to typed, atomic AG-UI interrupts." +keywords: + - tanstack ai migration + - ag-ui interrupts + - addToolApprovalResponse + - pendingInterrupts + - resumeInterrupts +--- + +# Migrate to AG-UI Interrupts + +TanStack AI now represents approvals, generic pauses, and client-tool execution +with standard AG-UI interrupt descriptors and results. Native runs end with +`RUN_FINISHED.outcome.type === 'interrupt'`; the continuation is a new run whose +`parentRunId` is the interrupted run. + +There is no codemod. Migrate the server lifecycle, client rendering, and +persistence schema together. The old readers remain temporarily available for +legacy streams, but they cannot provide the full native contract. + +For the complete API and runnable server/client setup, read +[Interrupts](../chat/interrupts). + +## API mapping + +| Deprecated or legacy surface | Current surface | +| --- | --- | +| `pendingInterrupts` | `interrupts` (`pendingInterrupts` is a deprecated alias of the same bound array) | +| `ChatClient.getPendingInterrupts()` | `ChatClient.getInterrupts()` | +| `addToolApprovalResponse({ id, approved })` | Find the bound `tool-approval` item and call `interrupt.resolveInterrupt(approved)` | +| Raw `resumeInterrupts(entries, state)` | Bound item methods or root `resolveInterrupts(...)`; reserve `resumeInterruptsUnsafe` for validated low-level recovery tooling | +| Native `approval-requested` custom event | `RUN_FINISHED` interrupt descriptor with reason `tool_call` | +| Native `tool-input-available` custom event | `RUN_FINISHED` interrupt descriptor with reason `tanstack:client_tool_execution` | +| Boolean denial treated like cancellation | `resolveInterrupt(false)` for denial; `cancel()` for payloadless cancellation | + +`addToolResult` is **not** removed. It remains supported for client-tool results +and delegates to a matching native client-tool interrupt when one exists. +`needsApproval` also remains the tool-definition switch for approvals. + +## Migrate a single approval + +Before: + +```ts ignore +await addToolApprovalResponse({ + id: approval.id, + approved: true, +}) +``` + +After: + +```ts ignore +const interrupt = interrupts.find( + (item) => item.kind === 'tool-approval' && item.toolName === 'transfer', +) + +if ( + interrupt?.kind === 'tool-approval' && + interrupt.toolName === 'transfer' +) { + interrupt.resolveInterrupt(true) +} +``` + +The bound method validates locally and the server validates again. A valid +singleton submits automatically. + +## Add branch payloads and optional edits + +Legacy boolean approvals could not carry typed branch-specific application +data. Add `approvalSchema` to the shared tool definition: + +```ts +import { toolDefinition } from '@tanstack/ai' +import { z } from 'zod' + +export const transferTool = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema: z.object({ + amount: z.number().positive(), + recipient: z.string().min(1), + }), + approvalSchema: { + approve: z.object({ note: z.string().min(1) }), + reject: z.object({ reason: z.string().min(1) }), + }, +}) +``` + +Resolve the selected branch with data nested under `payload`: + +```ts ignore +interrupt.resolveInterrupt(true, { + editedArgs: { amount: 12, recipient: 'Ada' }, + payload: { note: 'Reviewed' }, +}) + +interrupt.resolveInterrupt(false, { + payload: { reason: 'Policy limit' }, +}) +``` + +`editedArgs` is optional and approval-only. It is a full replacement validated +against the original input schema. Omitting it keeps the original tool input. +Rejection never accepts edits. Custom branch fields at the top level are +invalid; keep them grouped under `payload`. + +If `approvalSchema` is one schema instead of `{ approve, reject }`, that schema +applies to the selected decision. If no approval schema exists, the normal +boolean shorthand remains valid. + +## Separate denial from cancellation + +These are intentionally different results: + +```ts ignore +// Canonical denial: a resolved answer, optionally with a reject payload. +interrupt.resolveInterrupt(false, { + payload: { reason: 'Not authorized' }, +}) + +// Cancellation: no decision and no payload. +interrupt.cancel() +``` + +A denial continues the model with an explicit rejected decision. Cancellation +emits AG-UI `status: 'cancelled'` and never validates or selects the reject +branch. Deprecated `addToolApprovalResponse({ approved: false })` maps to a +denial for compatibility; it does not map to cancellation. + +## Migrate a batch + +Previously, applications often looped over approval IDs and sent multiple +requests. Native interrupt batches are all-or-nothing. Stage each item; the +last valid item auto-submits the complete batch: + +```ts ignore +for (const interrupt of interrupts) { + if (interrupt.kind === 'tool-approval') { + interrupt.resolveInterrupt(true) + } else { + interrupt.cancel() + } +} +``` + +For a single synchronous transaction, use the root callback: + +```ts ignore +await resolveInterrupts((interrupt) => { + if ( + interrupt.kind === 'tool-approval' && + interrupt.toolName === 'transfer' + ) { + interrupt.resolveInterrupt(true, { + editedArgs: interrupt.originalArgs, + payload: { note: 'Approved during batch review' }, + }) + return + } + + interrupt.cancel() +}) +``` + +The callback must synchronously resolve or cancel every item. It cannot return a +promise. An invalid item, thrown error, or incomplete set prevents the entire +submission. + +`resolveInterrupts(true)` and `resolveInterrupts(false)` are convenience +operations only for all-tool-approval batches whose selected branches require +no payload or edits. They are rejected for generic items, client-tool outputs, +mixed batches, and required branch payloads. Use `cancelInterrupts()` for a +payloadless all-items cancellation. + +Use `interrupt.clearResolution()` to remove one staged draft and its errors. +Use `retryInterrupts()` only after every item still has a valid staged response +and the root error is retryable. + +## Migrate generic responses + +Generic AG-UI descriptors carry a wire `responseSchema`, but their application +payload remains `unknown`. Do not derive a static TypeScript type from received +JSON. Parse the editor value as `unknown`, convert the Draft 2020-12 schema to a +runtime validator, then resolve the validated value: + +```ts +import { z } from 'zod' +import type { GenericAGUIInterrupt } from '@tanstack/ai-client' + +function submitGenericValue( + interrupt: GenericAGUIInterrupt, + editorValue: string, +): readonly string[] { + if (!interrupt.responseSchema) return ['Missing response schema.'] + + let value: unknown + try { + value = JSON.parse(editorValue) + } catch { + return ['Enter valid JSON.'] + } + + const validator = z.fromJSONSchema(interrupt.responseSchema) + const result = validator.safeParse(value) + if (!result.success) { + return result.error.issues.map((issue) => issue.message) + } + + interrupt.resolveInterrupt(result.data) + return [] +} +``` + +The client performs canonical validation too, and server validation remains +authoritative. + +## Migrate server events and persistence + +A native server must emit snapshots before the interrupt terminal: + +1. `MESSAGES_SNAPSHOT` +2. optional `STATE_SNAPSHOT` +3. `RUN_FINISHED` with a nonempty interrupt outcome + +It must persist the complete descriptor/binding batch before exposing that +terminal. Without the atomic persistence capability, TanStack AI emits one +`RUN_ERROR` with `persistence-required` instead. + +Continuations use a fresh `runId`, the same `threadId`, and the interrupted run +as `parentRunId`. The resume request contains every pending ID exactly once. +Persistence validates every payload, edited input, schema hash, expiry, and +generation before one compare-and-swap commit. Exact retries attach to the +winning continuation rather than executing tools again. + +Upgrade persistence before deploying the new runtime: + +- add the interrupt batch, binding, generation, submission fingerprint, + continuation, and accepted-tombstone storage required by your adapter; +- add the parent/current run correlation used by compare-and-swap; +- generate and deploy your application's native Prisma migration after copying + the updated model fragment; +- deploy Drizzle or Cloudflare migration assets through the application's normal + migration process. + +See [Persistence migrations](../persistence/migrations) and +[Custom stores](../persistence/custom-stores). + +## Opt in to native recovery explicitly + +Browser persistence now writes a V2 envelope containing raw JSON-safe drafts. +Bound methods, validators, tool implementations, kinds, and hydrated errors are +not stored. After reload, the client must fetch authoritative recovery state +before rebinding descriptors or restoring a draft. + +Configure application-owned URLs explicitly: + +```ts +import { + createInterruptContinuationLoader, + createInterruptStateFetcher, + fetchServerSentEvents, +} from '@tanstack/ai-client' + +const connection = fetchServerSentEvents('/api/chat', { + interruptStateFetcher: createInterruptStateFetcher( + '/api/interrupts/recovery', + ), + continuationLoader: createInterruptContinuationLoader( + '/api/interrupts/continuation', + ), +}) +``` + +No recovery route is inferred from `/api/chat`. Your recovery endpoint must +authenticate the caller and authorize the thread. The continuation endpoint +must read the winning run without scheduling a new model invocation. + +`resumeInterruptsUnsafe` is available only for low-level recovery integrations +that already possess validated raw AG-UI resume entries and correlation state. +It is not the normal migration target for approval UI. + +## Legacy compatibility limits + +Deprecated readers recognize well-formed historical `approval-requested` and +`tool-input-available` events. They do not create a second native writer. A +legacy batch is converted into one cloned-history follow-up only after every +item is covered. + +Legacy descriptors do not support: + +- edited arguments or custom approval payloads; +- generic AG-UI responses; +- native payloadless cancellation semantics; +- expiry and schema-hash reconciliation; +- generation conflicts or authoritative native recovery. + +Those operations fail with `legacy-unsupported`. Native and legacy items cannot +be mixed in one batch. If legacy transport submission fails, staged decisions +remain available and the root reports `legacy-submit-failed`. + +## Migration checklist + +1. Deploy persistence schema changes for every configured backend. +2. Add `withChatPersistence(...)` to interrupt-producing chat routes. +3. Replace native custom-event writers with the AG-UI interrupt terminal. +4. Replace `pendingInterrupts` rendering with bound `interrupts`. +5. Replace boolean approval helpers with `resolveInterrupt` and explicit + denial/cancellation behavior. +6. Replace approval loops with atomic batch staging or root + `resolveInterrupts(...)`. +7. Add explicit recovery and continuation endpoints before enabling V2 draft + restoration. +8. Keep `addToolResult` for client-tool results where it is still useful. +9. Test reloads, exact retries, two-tab conflicts, expired items, and failed + transport recovery before removing legacy support. diff --git a/docs/persistence/chat-persistence.md b/docs/persistence/chat-persistence.md index bffcf6061..865a286bf 100644 --- a/docs/persistence/chat-persistence.md +++ b/docs/persistence/chat-persistence.md @@ -155,7 +155,7 @@ handled. This prevents accidental conversation forks. To reconnect to an in-flight SSE response, add a `StreamDurability` adapter to the response: -```ts +```ts ignore import { durableStream } from '@tanstack/ai-durable-stream' declare function getDurableStreamsToken(): Promise diff --git a/docs/persistence/cloudflare.md b/docs/persistence/cloudflare.md index 02564e7d8..68fa9471a 100644 --- a/docs/persistence/cloudflare.md +++ b/docs/persistence/cloudflare.md @@ -66,13 +66,14 @@ export { CloudflareLockDurableObject } from '@tanstack/ai-persistence-cloudflare ## Create persistence -```ts +```ts group=cloudflare-persistence import { cloudflarePersistence } from '@tanstack/ai-persistence-cloudflare' +import type { CloudflarePersistenceOptions } from '@tanstack/ai-persistence-cloudflare' interface Env { - AI_STATE: D1Database - AI_MEDIA: R2Bucket - AI_LOCKS: DurableObjectNamespace + AI_STATE: NonNullable + AI_MEDIA: NonNullable + AI_LOCKS: NonNullable } export function createPersistence(env: Env) { @@ -96,7 +97,7 @@ serializes owners and uses leases/alarms for recovery. ## Use it with chat -```ts +```ts group=cloudflare-persistence import { chat, chatParamsFromRequest, @@ -157,12 +158,13 @@ for (const migration of d1Migrations) { Use Cloudflare as the base and replace only application-owned stores: -```ts +```ts group=cloudflare-persistence import { composePersistence } from '@tanstack/ai-persistence' import type { InterruptStore, RunStore } from '@tanstack/ai-persistence' declare const customInterrupts: InterruptStore declare const customRuns: RunStore +declare const env: Env const persistence = composePersistence(createPersistence(env), { overrides: { diff --git a/docs/persistence/controls.md b/docs/persistence/controls.md index 201acb625..68f5a79e2 100644 --- a/docs/persistence/controls.md +++ b/docs/persistence/controls.md @@ -32,21 +32,22 @@ the complete capability selection mechanism. `composePersistence` takes the base backend first and a configuration object second: -```ts +```ts group=persistence-controls import { composePersistence, defineAIPersistence, } from '@tanstack/ai-persistence' import { cloudflarePersistence } from '@tanstack/ai-persistence-cloudflare' +import type { CloudflarePersistenceOptions } from '@tanstack/ai-persistence-cloudflare' import type { InterruptStore, RunStore, } from '@tanstack/ai-persistence' declare const env: { - AI_STATE: D1Database - AI_MEDIA: R2Bucket - AI_LOCKS: DurableObjectNamespace + AI_STATE: NonNullable + AI_MEDIA: NonNullable + AI_LOCKS: NonNullable } declare const interrupts: InterruptStore declare const runs: RunStore @@ -76,7 +77,7 @@ Each override is independent: | a store object | Replace that store only. | | `false` | Remove that store. | -```ts +```ts group=persistence-controls const withoutGeneratedMedia = composePersistence(base, { overrides: { artifacts: false, @@ -101,7 +102,7 @@ Some capabilities require related stores: Known-invalid static compositions fail to type-check at the middleware call. Runtime validation covers dynamically typed inputs. -```ts +```ts group=persistence-controls import { withGenerationPersistence } from '@tanstack/ai-persistence' // Valid: both stores remain present. diff --git a/docs/persistence/custom-stores.md b/docs/persistence/custom-stores.md index 22889ce0d..f49ef4f93 100644 --- a/docs/persistence/custom-stores.md +++ b/docs/persistence/custom-stores.md @@ -14,15 +14,17 @@ database. ```ts import { defineAIPersistence } from '@tanstack/ai-persistence' import type { + InterruptStore, MessageStore, RunStore, } from '@tanstack/ai-persistence' declare const messages: MessageStore declare const runs: RunStore +declare const interrupts: InterruptStore export const persistence = defineAIPersistence({ - stores: { messages, runs }, + stores: { messages, runs, interrupts }, }) ``` @@ -74,22 +76,73 @@ Implement `createOrResume` idempotently. Retries may repeat the same run id. ### Interrupts +Interrupt persistence is an atomic capability, not a collection of sequential +`create`, `resolve`, and `cancel` calls. Opening a batch must store all +descriptors and protected bindings together. Committing a response must compare +the expected generation and exact ID set, then resolve every item and create +the continuation receipt in the same transaction. + ```ts -import type { InterruptRecord } from '@tanstack/ai-persistence' +import type { + CommitInterruptResolutionsInput, + InterruptCommitResult, + InterruptRecoveryQuery, + InterruptRecoveryStateV1, + OpenInterruptBatchInput, +} from '@tanstack/ai' interface InterruptStore { - create(record: Omit): Promise - resolve(interruptId: string, response?: unknown): Promise - cancel(interruptId: string): Promise - get(interruptId: string): Promise - list(threadId: string): Promise> - listPending(threadId: string): Promise> - listByRun(runId: string): Promise> - listPendingByRun(runId: string): Promise> + openInterruptBatch(input: OpenInterruptBatchInput): Promise<{ + generation: number + descriptors: OpenInterruptBatchInput['descriptors'] + }> + commitInterruptResolutions( + input: CommitInterruptResolutionsInput, + ): Promise + getInterruptRecoveryState( + input: InterruptRecoveryQuery, + ): Promise } ``` +The commit result is the durable receipt: + +- `committed` returns the newly accepted `continuationRunId`; +- `replayed` returns the same winning continuation for an exact idempotent + retry; +- `conflict` returns authoritative recovery state for a stale generation or a + different submission. + +Compute and compare the canonical submission fingerprint inside the same +transaction as the current-run/parent-run compare-and-swap. Never rerun tools +for `replayed`. Do not accept a subset or superset of the opened IDs. + An `interrupts` store requires a `runs` store when used with chat persistence. +`defineAIPersistence` rejects an adapter that advertises interrupts without the +atomic methods. See [Interrupts](../chat/interrupts) for client semantics and +[Delivery durability](./delivery-durability) for exact replay behavior. + +### Explicit recovery handler + +Expose recovery only on an application-owned route. The handler does not infer +a URL and must authorize the caller and thread before reading state: + +```ts +// app/api/interrupts/recovery/route.ts +import { createInterruptRecoveryHandler } from '@tanstack/ai-persistence' +import { persistence } from '../../../../persistence' +import { authorizeInterruptRecovery } from '../../../../authorization' + +export const POST = createInterruptRecoveryHandler({ + gateway: persistence.stores.interrupts, + authorize: authorizeInterruptRecovery, +}) +``` + +Treat the concrete authorization object as application policy: authenticate the +request, authorize the requested thread, and redact protected bindings or +responses that the caller must not receive. Configure the matching client URL +explicitly with `createInterruptStateFetcher`. ### Metadata @@ -159,7 +212,7 @@ transactions, conditional writes, or stable idempotency keys. ## Override selected packaged stores -```ts +```ts ignore /// import { composePersistence } from '@tanstack/ai-persistence' diff --git a/docs/persistence/delivery-durability.md b/docs/persistence/delivery-durability.md index 556b1c166..405dec5c2 100644 --- a/docs/persistence/delivery-durability.md +++ b/docs/persistence/delivery-durability.md @@ -9,6 +9,12 @@ Delivery durability records an ordered SSE stream so a dropped connection can replay chunks without invoking the provider again. It is separate from state persistence for messages, runs, interrupts, metadata, and artifacts. +Interrupt submission has a second exact-replay rule. The persistence store +canonicalizes the complete resolution set into an idempotency fingerprint and +commits it with the continuation run ID. Repeating that exact set returns the +same winning continuation; it never executes approved tools twice. A different +set or stale generation is a conflict and returns authoritative recovery state. + Two adapters ship: - `memoryStream(request)` stores a process-local log for development and tests. @@ -115,13 +121,44 @@ declare const runId: string const connection = fetchServerSentEvents('/api/chat') for await (const chunk of connection.joinRun(runId)) { - console.log(chunk.type) + if ('type' in chunk) { + console.log(chunk.type) + } } ``` `joinRun` performs a GET with `offset=-1` and `runId`. The endpoint must accept that GET, as shown above. +For interrupt recovery, configure this winning-run loader explicitly with +`createInterruptContinuationLoader`. TanStack AI never guesses a continuation +or recovery route from the chat URL. See [Interrupts](../chat/interrupts) for +the client setup. + +## Accepted tombstones and exact replay + +After a native batch is accepted, the browser first persists a V2 +`phase: 'accepted'` tombstone containing the winning continuation ID and no +drafts, then attempts to remove the resume record. If removal fails, reload does +not show stale pending UI. Authoritative committed recovery confirms the winner +and retries cleanup. + +The layers cooperate without sharing identifiers: + +1. Interrupt persistence compares the generation, exact ID set, and canonical + resolution fingerprint. +2. A first valid submission atomically stores the resolutions and winning + continuation receipt. +3. An exact retry returns `replayed` with that continuation ID. +4. The client joins the winning run through SSE durability or an explicit + continuation loader. +5. SSE delivery resumes from the adapter-owned opaque offset and de-duplicates + the replayed prefix. + +The fingerprint is not an SSE offset. The accepted tombstone is not the +authoritative server commit. Keep both state persistence and delivery +durability when the workflow must survive request retries and connection loss. + ## Offset ownership `StreamDurability` owns its offset format. Core only passes returned @@ -172,3 +209,7 @@ Delivery logs replay chunks. They are not the queryable source of truth for thread messages, pending interrupts, generation artifacts, or retention policies. Add [Chat Persistence](./chat-persistence) or [Generation Persistence](./generation-persistence) for that state. + +For atomic interrupt store requirements and conflict receipts, see +[Custom stores](./custom-stores). For the legacy-to-native transition, see +[Migrate to AG-UI interrupts](../migration/interrupts). diff --git a/docs/persistence/drizzle.md b/docs/persistence/drizzle.md index 9baed12e0..662712cc2 100644 --- a/docs/persistence/drizzle.md +++ b/docs/persistence/drizzle.md @@ -18,7 +18,7 @@ backend. ## Node SQLite -```ts +```ts group=drizzle-node import { sqlitePersistence } from '@tanstack/ai-persistence-drizzle/sqlite' export const persistence = sqlitePersistence({ @@ -34,13 +34,16 @@ deployment-time migrations in production. ## Bring your own SQLite Drizzle database ```ts -import { drizzle } from 'drizzle-orm/d1' +import type { CloudflarePersistenceOptions } from '@tanstack/ai-persistence-cloudflare' import { drizzlePersistence, schema, } from '@tanstack/ai-persistence-drizzle' +import { drizzle } from 'drizzle-orm/d1' -declare const env: { AI_STATE: D1Database } +declare const env: { + AI_STATE: NonNullable +} const db = drizzle(env.AI_STATE, { schema }) export const persistence = drizzlePersistence(db) @@ -52,7 +55,7 @@ lifecycle and migration timing. ## Use the middleware -```ts +```ts group=drizzle-node import { chat, chatParamsFromRequest, diff --git a/docs/persistence/generation-persistence.md b/docs/persistence/generation-persistence.md index 3f3651ed6..cade7cfb8 100644 --- a/docs/persistence/generation-persistence.md +++ b/docs/persistence/generation-persistence.md @@ -13,7 +13,7 @@ Generated bytes never belong in browser resume state. ## Create the server endpoint -```ts +```ts group=generation-persistence import { generateImage, generationParamsFromRequest, @@ -122,7 +122,7 @@ restart or reconnect to provider work. Artifact metadata and blob bodies are separate. The default generated blob key is `artifacts//`. -```ts +```ts group=generation-persistence export async function GET( _request: Request, context: { params: Promise<{ artifactId: string }> }, @@ -154,7 +154,7 @@ Built-in extraction handles input media, generated image/audio/video output, transcription input audio, and structured transcription results. Supply `extractArtifacts` to replace the built-in extraction for a run: -```ts +```ts group=generation-persistence const generationPersistence = withGenerationPersistence(persistence, { extractArtifacts({ result }) { return [ diff --git a/docs/persistence/migrations.md b/docs/persistence/migrations.md index bf5ca1285..ed3f33a92 100644 --- a/docs/persistence/migrations.md +++ b/docs/persistence/migrations.md @@ -8,6 +8,18 @@ id: migrations Migration ownership depends on the backend. Apply schema changes before deploying code that reads or writes the corresponding stores. +AG-UI interrupts add atomic batch state. Every backend must persist the +descriptor and protected binding set, generation, exact ID set, canonical +submission fingerprint, committed resolutions, continuation run ID, and +accepted/commit timestamps. Runs also need current/parent correlation for the +compare-and-swap that creates a continuation. Add uniqueness for one generation +of an interrupted run and indexes for pending recovery lookups. + +Do not approximate this with sequential interrupt rows. A partial write can +expose an interrupt that cannot be resumed safely. See +[Custom stores](./custom-stores) for the atomic capability and +[Interrupts](../chat/interrupts) for the lifecycle. + ## Drizzle SQLite The Drizzle package bundles ordered canonical SQLite migrations. Copy them into @@ -27,6 +39,11 @@ The CLI preserves existing identical files and refuses to overwrite divergent files without `--force`. Apply the copied SQL using your normal SQLite, Drizzle, or D1 deployment process. +Refresh the manifest to pick up the interrupt batch table/columns, generation +and pending-recovery indexes, submission fingerprint uniqueness, continuation +receipt, and parent-run correlation. Review the generated SQL before applying +it to every SQLite environment. + For local Node development, the `/sqlite` factory can apply the same manifest: ```ts @@ -53,6 +70,11 @@ It also exports `d1Migrations` for programmatic tooling. R2 and Durable Object locks do not use the D1 table migration set; configure their bindings and Durable Object migration tags in Wrangler. +The D1 migration set carries the same atomic interrupt fields and indexes as +Drizzle SQLite. Apply it before deploying Workers that can emit native +interrupts. A Durable Object lock can serialize workers, but it does not replace +the D1 compare-and-swap or accepted receipt. + ## Prisma Prisma ships a provider-neutral models fragment, then delegates SQL generation @@ -69,12 +91,33 @@ The copied fragment contains no datasource or generator. Keep those in your application schema and commit the native migration Prisma creates for your provider. +The updated fragment adds the models and fields needed for atomic interrupt +batches, bindings, generations, fingerprints, accepted receipts, continuation +runs, and parent/current-run correlation. The package does **not** ship SQL for +your provider. Generate and deploy your own Prisma migration before using the +new package version: + +```bash +pnpm exec tanstack-ai-prisma-models --out prisma/schema +pnpm prisma migrate dev --name add-ag-ui-interrupts +pnpm prisma generate +pnpm prisma migrate deploy +``` + +Inspect and commit the migration that Prisma generates. Do not deploy the +runtime first or rely on `prisma db push` as a production migration. + ## Custom stores Custom `AIPersistence` adapters own their schema and migrations entirely. Maintain compatibility with the public store records and method semantics; TanStack AI does not inspect your table layout. +For this upgrade, add one transaction boundary spanning exact-set validation, +generation compare-and-swap, resolution persistence, continuation receipt, and +run-parent linkage. Preserve committed fingerprints long enough for exact retry +replay and recovery after a browser accepted-tombstone cleanup failure. + ## Upgrade discipline 1. Read the package release notes for schema changes. diff --git a/docs/persistence/overview.md b/docs/persistence/overview.md index 12d44f98f..57a9acc91 100644 --- a/docs/persistence/overview.md +++ b/docs/persistence/overview.md @@ -79,12 +79,13 @@ the second argument: import { composePersistence } from '@tanstack/ai-persistence' import { cloudflarePersistence } from '@tanstack/ai-persistence-cloudflare' import { defineAIPersistence } from '@tanstack/ai-persistence' +import type { CloudflarePersistenceOptions } from '@tanstack/ai-persistence-cloudflare' import type { AIPersistenceStores } from '@tanstack/ai-persistence' declare const env: { - AI_STATE: D1Database - AI_MEDIA: R2Bucket - AI_LOCKS: DurableObjectNamespace + AI_STATE: NonNullable + AI_MEDIA: NonNullable + AI_LOCKS: NonNullable } declare const myInterruptStore: NonNullable< AIPersistenceStores['interrupts'] diff --git a/docs/persistence/sandbox-persistence.md b/docs/persistence/sandbox-persistence.md index a74d34cae..997ffca17 100644 --- a/docs/persistence/sandbox-persistence.md +++ b/docs/persistence/sandbox-persistence.md @@ -129,14 +129,13 @@ product contract. Cloudflare can back all workspace checkpoint stores: ```ts -/// - import { cloudflarePersistence } from '@tanstack/ai-persistence-cloudflare' +import type { CloudflarePersistenceOptions } from '@tanstack/ai-persistence-cloudflare' declare const env: { - AI_STATE: D1Database - AI_MEDIA: R2Bucket - AI_LOCKS: DurableObjectNamespace + AI_STATE: NonNullable + AI_MEDIA: NonNullable + AI_LOCKS: NonNullable } const persistence = cloudflarePersistence({ diff --git a/docs/persistence/sql-backends.md b/docs/persistence/sql-backends.md index fbf85471f..875632c22 100644 --- a/docs/persistence/sql-backends.md +++ b/docs/persistence/sql-backends.md @@ -31,13 +31,14 @@ export const persistence = sqlitePersistence({ ## Existing SQLite or D1 Drizzle database ```ts +import type { CloudflarePersistenceOptions } from '@tanstack/ai-persistence-cloudflare' import { drizzlePersistence, schema, } from '@tanstack/ai-persistence-drizzle' import { drizzle } from 'drizzle-orm/d1' -declare const stateDatabase: D1Database +declare const stateDatabase: NonNullable const db = drizzle(stateDatabase, { schema }) export const persistence = drizzlePersistence(db) @@ -47,7 +48,7 @@ The package root is edge-safe; `/sqlite` is Node-only. ## Prisma -```ts +```ts group=sql-backends import { PrismaClient } from '@prisma/client' import { prismaPersistence } from '@tanstack/ai-persistence-prisma' @@ -64,7 +65,7 @@ Both adapters provide messages, runs, interrupts, metadata, artifacts, blobs, and an in-process lock store. Replace `locks` with a distributed store when multiple processes can mutate the same run: -```ts +```ts group=sql-backends import { composePersistence } from '@tanstack/ai-persistence' import type { LockStore } from '@tanstack/ai' diff --git a/docs/reference/functions/toolDefinition.md b/docs/reference/functions/toolDefinition.md index 66ce2837e..f769550f2 100644 --- a/docs/reference/functions/toolDefinition.md +++ b/docs/reference/functions/toolDefinition.md @@ -6,7 +6,19 @@ title: toolDefinition # Function: toolDefinition() ```ts -function toolDefinition(config): ToolDefinition; +function toolDefinition< + TInput, + TOutput, + TName, + TNeedsApproval, + TApprovalSchema, +>(config): ToolDefinition< + TInput, + TOutput, + TName, + TNeedsApproval, + TApprovalSchema +> ``` Defined in: [packages/ai/src/activities/chat/tools/tool-definition.ts:191](https://github.com/TanStack/ai/blob/main/packages/ai/src/activities/chat/tools/tool-definition.ts#L191) @@ -21,6 +33,34 @@ The definition contains all tool metadata (name, description, schemas) and can b Supports any Standard JSON Schema compliant library (Zod v4+, ArkType, Valibot, etc.) or plain JSON Schema objects. +## Conditional approval schema + +`approvalSchema` is available only when `needsApproval: true`. It accepts either +one Standard Schema/JSON Schema for both decisions or a nonempty branch map: + +```ts +type ApprovalSchemaConfig = + | SchemaInput + | { approve: SchemaInput; reject?: SchemaInput } + | { approve?: SchemaInput; reject: SchemaInput } +``` + +The schema generic is preserved by `.server()` and `.client()`. Client +`tool-approval` interrupts infer the selected branch payload, require it when +the schema requires it, and place it under `payload`. Approval may also carry an +optional, fully validated `editedArgs` replacement when the tool has an input +schema. Rejection never accepts edited arguments. + +Plain JSON Schema remains runtime-only and therefore produces `unknown` payload +data. Standard Schema inputs such as Zod infer both runtime validation and the +bound resolver overloads. + +At runtime, defining `approvalSchema` without `needsApproval: true` throws. +TanStack AI converts the input, output, and selected approval branches to +canonical JSON Schema, embeds their hashes in the protected interrupt binding, +and validates again on resume. See [Interrupts](../../chat/interrupts) for the +full lifecycle. + ## Type Parameters ### TInput @@ -35,6 +75,16 @@ or plain JSON Schema objects. `TName` *extends* `string` = `string` +### TNeedsApproval + +`TNeedsApproval` *extends* `boolean` = `false`. The literal `true` enables the +approval capability in mapped client interrupt types. + +### TApprovalSchema + +`TApprovalSchema` *extends* `ApprovalSchemaConfig | undefined` = `undefined`. +This generic is conditionally permitted only when `TNeedsApproval` is `true`. + ## Parameters ### config @@ -65,6 +115,10 @@ const addToCartTool = toolDefinition({ cartId: z.string(), totalItems: z.number(), }), + approvalSchema: { + approve: z.object({ note: z.string() }), + reject: z.object({ reason: z.string() }), + }, }); // Use directly in chat (server-side, no execute function) @@ -89,3 +143,17 @@ const addToCartClient = addToCartTool.client(async (args) => { return { success: true, cartId: 'local', totalItems: 1 }; }); ``` + +With `tools: [addToCartTool] as const`, the corresponding bound approval has +branch-specific overloads: + +```ts +interrupt.resolveInterrupt(true, { + editedArgs: { guitarId: 'guitar-2', quantity: 2 }, + payload: { note: 'Reviewed' }, +}) + +interrupt.resolveInterrupt(false, { + payload: { reason: 'Budget limit' }, +}) +``` diff --git a/docs/tools/client-tools.md b/docs/tools/client-tools.md index a6232d6f3..488b1c4c6 100644 --- a/docs/tools/client-tools.md +++ b/docs/tools/client-tools.md @@ -27,7 +27,7 @@ sequenceDiagram Note over Server: No execute function
= client tool - Server->>Browser: Forward tool-input-available
chunk via SSE/HTTP + Server->>Browser: RUN_FINISHED client-tool
interrupt via SSE/HTTP Browser->>Browser: Find registered
client tool Browser->>ClientTool: execute(args) ClientTool->>ClientTool: Update UI,
localStorage, etc. @@ -53,12 +53,57 @@ sequenceDiagram 1. **Tool Call from LLM**: LLM decides to call a client tool 2. **Server Detection**: Server sees the tool has no `execute` function -3. **Client Notification**: Server sends a `tool-input-available` chunk to the browser +3. **Client Notification**: Server emits a `client-tool-execution` AG-UI interrupt 4. **Client Execution**: The browser finds the registered `.client()` implementation by tool name and runs it with the parsed input 5. **Result Return**: Client executes the tool and returns the result 6. **Server Update**: Result is sent back to the server and added to the conversation 7. **LLM Continuation**: LLM receives the result and continues the conversation +Native client-tool execution uses the same atomic interrupt lifecycle as other +waits. See [Interrupts](../chat/interrupts) for persistence, batches, recovery, +and migration from the historical `tool-input-available` custom event. + +## Approval and execution are separate axes + +A client tool can require approval, but approval is not the browser result. A +tool with `needsApproval: true` first produces a `tool-approval` interrupt. Once +approved, its browser work is represented separately by a +`client-tool-execution` interrupt. + +```ts ignore +const approval = interrupts.find( + (interrupt) => + interrupt.kind === 'tool-approval' && + interrupt.toolName === 'delete_local_data', +) + +if ( + approval?.kind === 'tool-approval' && + approval.toolName === 'delete_local_data' +) { + approval.resolveInterrupt(true) +} + +const execution = interrupts.find( + (interrupt) => + interrupt.kind === 'client-tool-execution' && + interrupt.toolName === 'delete_local_data', +) + +if ( + execution?.kind === 'client-tool-execution' && + execution.toolName === 'delete_local_data' +) { + execution.resolveInterrupt({ deleted: true }) +} +``` + +The result is validated against the tool's output schema. The existing +`addToolResult` API remains supported and delegates to the same staged native +item when one matches; it also preserves the historical path for legacy +streams. See [Tool approval flow](./tool-approval) for approval forms and +[Migrate to AG-UI interrupts](../migration/interrupts) for compatibility limits. + ## Defining Client Tools Client tools use the same `toolDefinition()` API but with the `.client()` method: @@ -232,7 +277,7 @@ function MessageComponent({ message }: { message: ChatMessages[number] }) { Client tools are **automatically executed** when the model calls them. The flow is: 1. LLM calls a client tool -2. Server sends `tool-input-available` chunk to browser +2. Server sends a `client-tool-execution` interrupt to the browser 3. Client automatically executes the matching tool implementation 4. Result is sent back to server 5. Conversation continues with the result diff --git a/docs/tools/tool-approval.md b/docs/tools/tool-approval.md index 2ab21975d..cc3287c48 100644 --- a/docs/tools/tool-approval.md +++ b/docs/tools/tool-approval.md @@ -15,6 +15,11 @@ keywords: The tool approval flow allows you to require user approval before executing sensitive tools, giving users control over actions like sending emails, making purchases, or deleting data. A tool call moves through the `ToolCallState` lifecycle: +The current client API exposes approvals as bound AG-UI interrupts. For the +complete server/client lifecycle, atomic batch controls, generic interrupts, +and recovery, see [Interrupts](../chat/interrupts). For deprecated API mapping, +see [Migrate to AG-UI interrupts](../migration/interrupts). + 1. **`awaiting-input`** — Tool call started, no arguments yet 2. **`input-streaming`** — Arguments arriving incrementally 3. **`input-complete`** — All arguments received @@ -36,6 +41,69 @@ When a tool requires approval, the typical flow is: 4. Tool executes (if approved) or is cancelled (if denied) 5. Conversation continues with the result +## Resolve an approval interrupt + +Without an `approvalSchema`, use the boolean shorthand. Approval uses the +original tool input by default: + +```ts ignore +const approval = interrupts.find( + (interrupt) => interrupt.kind === 'tool-approval', +) + +if (approval?.kind === 'tool-approval') { + approval.resolveInterrupt(true) +} +``` + +An `approvalSchema` can define separate application payloads for approval and +rejection: + +```ts +import { toolDefinition } from '@tanstack/ai' +import { z } from 'zod' + +const transferDefinition = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema: z.object({ + amount: z.number().positive(), + recipient: z.string(), + }), + approvalSchema: { + approve: z.object({ note: z.string() }), + reject: z.object({ reason: z.string() }), + }, +}) +``` + +Keep branch data under `payload`. Approved arguments can optionally be replaced +in full with `editedArgs`; rejection never accepts edits: + +```ts ignore +approval.resolveInterrupt(true, { + editedArgs: { amount: 12, recipient: 'Ada' }, + payload: { note: 'Reviewed' }, +}) + +approval.resolveInterrupt(false, { + payload: { reason: 'Policy limit' }, +}) +``` + +Denial and cancellation are different. `resolveInterrupt(false, ...)` records a +resolved rejection for the continuation. `cancel()` is payloadless and +does not select the reject schema: + +```ts ignore +approval.cancel() +``` + +A singleton submits after its valid resolution. Multiple items stage until all +are valid, then submit atomically. Use root `resolveInterrupts(...)` for one +synchronous batch transaction. See [Interrupts](../chat/interrupts#singleton-and-batch-submission). + ## Enabling Approval Tools can be marked as requiring approval by setting `needsApproval: true` in the definition: @@ -91,9 +159,11 @@ export async function POST(request: Request) { } ``` -## Client-Side Approval Handling +## Deprecated approval response compatibility -The client receives approval requests and can respond: +`addToolApprovalResponse` remains temporarily available for old approval UIs, +but new code should use the bound interrupt API above. A compatibility +`approved: false` is a denial, not a cancellation. ```tsx import { useChat, fetchServerSentEvents } from '@tanstack/ai-react' @@ -186,7 +256,7 @@ const listData = toolDefinition({ name: 'list_data', description: 'List available keys', inputSchema: z.object({}), -}).client(async () => ({ keys: [] as Array })) +}).client(async () => ({ keys: new Array() })) function ApprovalHandler() { const { messages, addToolApprovalResponse } = useChat({ diff --git a/examples/ts-code-mode-web/src/routes/_database-demo/api.database-demo.ts b/examples/ts-code-mode-web/src/routes/_database-demo/api.database-demo.ts index 8bcd7e15a..8d9fd0908 100644 --- a/examples/ts-code-mode-web/src/routes/_database-demo/api.database-demo.ts +++ b/examples/ts-code-mode-web/src/routes/_database-demo/api.database-demo.ts @@ -10,10 +10,10 @@ import { skillsToTools, } from '@tanstack/ai-code-mode-skills' import { createFileSkillStorage } from '@tanstack/ai-code-mode-skills/storage' -import { anthropicText } from '@tanstack/ai-anthropic' -import { openaiText } from '@tanstack/ai-openai' -import { geminiText } from '@tanstack/ai-gemini' -import type { AnyTextAdapter, ServerTool, StreamChunk } from '@tanstack/ai' +import { ANTHROPIC_MODELS, anthropicText } from '@tanstack/ai-anthropic' +import { OPENAI_CHAT_MODELS, openaiText } from '@tanstack/ai-openai' +import { GEMINI_MODELS, geminiText } from '@tanstack/ai-gemini' +import type { AnyServerTool, AnyTextAdapter, StreamChunk } from '@tanstack/ai' import type { IsolateDriver } from '@tanstack/ai-code-mode' import { databaseTools, getSchemaInfoTool } from '@/lib/tools/database-tools' @@ -21,15 +21,25 @@ import { maxTokensModelOptions } from '@/lib/max-tokens-model-options' type Provider = 'anthropic' | 'openai' | 'gemini' +function selectModel( + requested: string | undefined, + fallback: TModel, + available: ReadonlyArray, +): TModel { + return available.find((candidate) => candidate === requested) ?? fallback +} + function getAdapter(provider: Provider, model?: string): AnyTextAdapter { switch (provider) { case 'openai': - return openaiText((model || 'gpt-4o') as 'gpt-4o') + return openaiText(selectModel(model, 'gpt-5.5', OPENAI_CHAT_MODELS)) case 'gemini': - return geminiText((model || 'gemini-2.5-flash') as 'gemini-2.5-flash') + return geminiText(selectModel(model, 'gemini-2.5-flash', GEMINI_MODELS)) case 'anthropic': default: - return anthropicText((model || 'claude-haiku-4-5') as 'claude-haiku-4-5') + return anthropicText( + selectModel(model, 'claude-haiku-4-5', ANTHROPIC_MODELS), + ) } } @@ -115,7 +125,7 @@ Rules: This is not optional — skill registration is a core part of your workflow.` async function getSkillToolsAndPrompt(driver: IsolateDriver): Promise<{ - skillTools: Array> + skillTools: Array skillsPrompt: string }> { const allSkills = await skillStorage.loadAll() @@ -250,7 +260,7 @@ export const Route = createFileRoute( const { adapter: instrumentedAdapter } = instrumentAdapter(rawAdapter) try { - let tools: Array> + let tools: Array let systemPrompts: Array if (useCodeMode) { diff --git a/packages/ai-acp/tests/compatible.test.ts b/packages/ai-acp/tests/compatible.test.ts index 3a47958bb..8b21528b7 100644 --- a/packages/ai-acp/tests/compatible.test.ts +++ b/packages/ai-acp/tests/compatible.test.ts @@ -68,7 +68,7 @@ const FAKE_ACP_AGENT = fakeAcpAgent() const baseDir = path.join(os.tmpdir(), `tanstack-ai-acp-test-${Date.now()}`) // No removeOnDestroy: destroying a sandbox right after killing its agent races // the OS releasing the dir (EBUSY on Windows). Clean the tree once at the end. -const provider = localProcessSandbox({ baseDir }) +const provider = localProcessSandbox({ baseDir, removeOnDestroy: false }) afterAll(async () => { await fsp.rm(baseDir, { diff --git a/packages/ai-acp/tests/sandbox-provisioning.test.ts b/packages/ai-acp/tests/sandbox-provisioning.test.ts index dac4f7768..4f13b0de0 100644 --- a/packages/ai-acp/tests/sandbox-provisioning.test.ts +++ b/packages/ai-acp/tests/sandbox-provisioning.test.ts @@ -112,7 +112,7 @@ const baseDir = path.join(os.tmpdir(), `tanstack-ai-acp-prov-${Date.now()}`) // No removeOnDestroy: destroying a sandbox right after killing its agent races // the OS releasing the dir (EBUSY on Windows). Clean the whole tree once at the // end instead, with retries for any lingering handle. -const provider = localProcessSandbox({ baseDir }) +const provider = localProcessSandbox({ baseDir, removeOnDestroy: false }) afterAll(async () => { await fsp.rm(baseDir, { diff --git a/packages/ai-angular/src/inject-chat.ts b/packages/ai-angular/src/inject-chat.ts index 186d5dee9..d67aa2332 100644 --- a/packages/ai-angular/src/inject-chat.ts +++ b/packages/ai-angular/src/inject-chat.ts @@ -15,11 +15,15 @@ import type { AnyClientTool, InferSchemaType, ModelMessage, + RunAgentResumeItem, SchemaInput, StreamChunk, } from '@tanstack/ai' import type { ChatClientState, + ChatInterrupt, + ChatInterruptState, + ChatResumeState, ConnectionStatus, InferredClientContext, StructuredOutputPart, @@ -32,6 +36,9 @@ import type { UIMessage, } from './types' +const EMPTY_INTERRUPTS = Object.freeze([]) +const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) + let nextId = 0 export function injectChat< @@ -63,6 +70,12 @@ export function injectChat< const isSubscribed = signal(false) const connectionStatus = signal('disconnected') const sessionGenerating = signal(false) + const interruptState = signal>({ + interrupts: EMPTY_INTERRUPTS, + pendingInterrupts: EMPTY_INTERRUPTS, + interruptErrors: EMPTY_INTERRUPT_ERRORS, + resuming: false, + }) // Reactive option sources. Plain values become constant computeds. const bodySource = @@ -112,6 +125,10 @@ export function injectChat< onResumeStateChange: (resumeState, pendingInterrupts) => { options.onResumeStateChange?.(resumeState, pendingInterrupts) }, + onInterruptStateChange: (nextInterruptState) => { + interruptState.set(nextInterruptState) + options.onInterruptStateChange?.(nextInterruptState) + }, tools: options.tools, onCustomEvent: (eventType, data, context) => options.onCustomEvent?.(eventType, data, context), @@ -128,6 +145,7 @@ export function injectChat< }) messages.set(client.getMessages()) + interruptState.set(client.getInterruptState()) // Sync reactive body / forwardedProps / context to the client. if (bodySource || forwardedPropsSource || contextSource) { @@ -245,6 +263,25 @@ export function injectChat< }) => { await client.addToolApprovalResponse(response) } + const interrupts = computed(() => interruptState().interrupts) + const pendingInterrupts = computed(() => interruptState().interrupts) + const interruptErrors = computed(() => interruptState().interruptErrors) + const resuming = computed(() => interruptState().resuming) + const resolveInterrupts = ( + resolution: boolean | ((interrupt: ChatInterrupt) => undefined), + ) => { + if (typeof resolution === 'boolean') { + client.resolveInterrupts(resolution) + } else { + client.resolveInterrupts(resolution) + } + } + const cancelInterrupts = () => client.cancelInterrupts() + const retryInterrupts = () => client.retryInterrupts() + const resumeInterruptsUnsafe = ( + resumeItems: Array, + state?: ChatResumeState, + ) => client.resumeInterruptsUnsafe(resumeItems, state) // eslint-disable-next-line no-restricted-syntax -- return shape diverges from conditional InjectChatResult; TS can't structurally narrow the TSchema-gated partial/final signals return { @@ -263,6 +300,14 @@ export function injectChat< clear, addToolResult, addToolApprovalResponse, + interrupts, + pendingInterrupts, + interruptErrors, + resuming, + resolveInterrupts, + cancelInterrupts, + retryInterrupts, + resumeInterruptsUnsafe, partial, final, } as unknown as InjectChatResult diff --git a/packages/ai-angular/src/types.ts b/packages/ai-angular/src/types.ts index d6c369263..3f6ee4a54 100644 --- a/packages/ai-angular/src/types.ts +++ b/packages/ai-angular/src/types.ts @@ -2,13 +2,18 @@ import type { AnyClientTool, InferSchemaType, ModelMessage, + RunAgentResumeItem, SchemaInput, } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, + BoundInterrupts, ChatClientOptions, ChatClientState, + ChatInterrupt, + ChatInterruptState, ChatRequestBody, + ChatResumeState, ClientContextOptionFromTools, ConnectionStatus, DistributedOmit, @@ -118,6 +123,24 @@ interface BaseInjectChatResult< id: string approved: boolean }) => Promise + /** Immutable bound interrupts for the current interrupted run. */ + interrupts: Signal> + /** @deprecated Use `interrupts`. */ + pendingInterrupts: Signal> + /** Batch-level interrupt errors. */ + interruptErrors: Signal['interruptErrors']> + /** Whether the client is submitting an interrupt batch. */ + resuming: Signal + resolveInterrupts: { + (approved: boolean): void + (resolver: (interrupt: ChatInterrupt) => undefined): void + } + cancelInterrupts: () => void + retryInterrupts: () => void + resumeInterruptsUnsafe: ( + resume: Array, + state?: ChatResumeState, + ) => Promise /** Reload the last assistant message. */ reload: () => Promise /** Stop the current response generation. */ diff --git a/packages/ai-angular/tests/inject-chat-types.test.ts b/packages/ai-angular/tests/inject-chat-types.test.ts index d56de87a9..205f1448b 100644 --- a/packages/ai-angular/tests/inject-chat-types.test.ts +++ b/packages/ai-angular/tests/inject-chat-types.test.ts @@ -7,7 +7,8 @@ import { describe, expectTypeOf, it } from 'vitest' import { z } from 'zod' import type { Signal } from '@angular/core' -import type { AnyClientTool } from '@tanstack/ai' +import { toolDefinition, type AnyClientTool } from '@tanstack/ai' +import { clientTools } from '@tanstack/ai-client' import { injectChat } from '../src/inject-chat' import type { DeepPartial, InjectChatResult } from '../src/types' @@ -75,3 +76,126 @@ describe('injectChat() return type (angular)', () => { }) }) }) + +describe('injectChat() interrupt types', () => { + it('preserves approval, generic, and client-tool inference', () => { + const inputSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { cents: 0 }, + output: { cents: 0 }, + }, + validate: (value: unknown) => ({ + value: + value !== null && + typeof value === 'object' && + 'cents' in value && + typeof value.cents === 'number' + ? { cents: value.cents } + : { cents: 0 }, + }), + }, + } + const approveSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { note: '' }, + output: { note: '' }, + }, + validate: () => ({ value: { note: '' } }), + }, + } + const rejectSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { reason: '' }, + output: { reason: '' }, + }, + validate: () => ({ value: { reason: '' } }), + }, + } + const outputSchema = z.object({ accountId: z.string() }) + const transfer = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema, + approvalSchema: { + approve: approveSchema, + reject: rejectSchema, + }, + }).client() + const confirm = toolDefinition({ + name: 'confirm', + description: 'Confirm without schemas', + needsApproval: true, + }).client() + const lookup = toolDefinition({ + name: 'lookup', + description: 'Lookup account', + outputSchema, + }).client(() => ({ accountId: 'account-1' })) + const tools = clientTools(transfer, confirm, lookup) + type Interrupt = ReturnType< + InjectChatResult['interrupts'] + >[number] + type Transfer = Extract< + Interrupt, + { kind: 'tool-approval'; toolName: 'transfer' } + > + type Confirm = Extract< + Interrupt, + { kind: 'tool-approval'; toolName: 'confirm' } + > + type Lookup = Extract< + Interrupt, + { kind: 'client-tool-execution'; toolName: 'lookup' } + > + type Generic = Extract + + const check = ( + transferInterrupt: Transfer, + confirmInterrupt: Confirm, + lookupInterrupt: Lookup, + genericInterrupt: Generic, + ) => { + transferInterrupt.resolveInterrupt(true, { + editedArgs: { cents: 100 }, + payload: { note: 'approved' }, + }) + transferInterrupt.resolveInterrupt(false, { + payload: { reason: 'declined' }, + }) + // @ts-expect-error rejected approvals cannot edit tool input + transferInterrupt.resolveInterrupt(false, { editedArgs: { cents: 1 } }) + // @ts-expect-error approve payload uses the approve branch + transferInterrupt.resolveInterrupt(true, { + payload: { reason: 'wrong branch' }, + }) + + confirmInterrupt.resolveInterrupt(true) + confirmInterrupt.resolveInterrupt(false) + // @ts-expect-error omitted input schema forbids edited input + confirmInterrupt.resolveInterrupt(true, { editedArgs: { cents: 1 } }) + // @ts-expect-error omitted approval branches forbid payloads + confirmInterrupt.resolveInterrupt(false, { + payload: { reason: 'no branch' }, + }) + + lookupInterrupt.resolveInterrupt({ accountId: 'account-1' }) + expectTypeOf(lookupInterrupt.resolveInterrupt) + .parameter(0) + .toEqualTypeOf<{ accountId: string }>() + expectTypeOf(genericInterrupt.resolveInterrupt) + .parameter(0) + .toEqualTypeOf() + } + void check + }) +}) diff --git a/packages/ai-angular/tests/inject-chat.test.ts b/packages/ai-angular/tests/inject-chat.test.ts index 6855074af..39f930637 100644 --- a/packages/ai-angular/tests/inject-chat.test.ts +++ b/packages/ai-angular/tests/inject-chat.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it, vi } from 'vitest' +import { afterEach, describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { EventType } from '@tanstack/ai/client' import { Component, signal } from '@angular/core' @@ -7,13 +7,93 @@ import { ChatClient } from '@tanstack/ai-client' import { injectChat } from '../src/inject-chat' import { createMockConnectionAdapter, + createInterruptResumeSnapshot, createTextChunks, renderInjectChat, } from './test-utils' const tick = () => new Promise((r) => setTimeout(r, 0)) +afterEach(() => { + vi.restoreAllMocks() +}) + describe('injectChat', () => { + describe('interrupt state', () => { + it('projects one immutable reactive snapshot with the deprecated pending alias', () => { + const onInterruptStateChange = vi.fn() + const { result, flush } = renderInjectChat({ + connection: createMockConnectionAdapter(), + initialResumeSnapshot: createInterruptResumeSnapshot(), + onInterruptStateChange, + }) + + expect(Object.isFrozen(result.interrupts())).toBe(true) + expect(result.pendingInterrupts()).toBe(result.interrupts()) + expect(result.interrupts()[0]).toMatchObject({ + id: 'staged-interrupt', + status: 'staged', + }) + expect(result.interrupts()[1]).toMatchObject({ + id: 'invalid-interrupt', + status: 'error', + error: { code: 'invalid-payload' }, + }) + expect(result.interruptErrors()).toEqual([]) + expect(result.resuming()).toBe(false) + expect(result.interrupts()[0]).toEqual( + expect.objectContaining({ + resolveInterrupt: expect.any(Function), + cancel: expect.any(Function), + clearResolution: expect.any(Function), + }), + ) + + result.resolveInterrupts(false) + TestBed.flushEffects() + flush() + expect(result.interruptErrors()[0]?.code).toBe( + 'unsupported-bulk-operation', + ) + expect(onInterruptStateChange).toHaveBeenLastCalledWith( + expect.objectContaining({ + interrupts: result.interrupts(), + interruptErrors: result.interruptErrors(), + }), + ) + }) + + it('delegates every root interrupt control to ChatClient', async () => { + const resolve = vi + .spyOn(ChatClient.prototype, 'resolveInterrupts') + .mockImplementation(() => {}) + const cancel = vi + .spyOn(ChatClient.prototype, 'cancelInterrupts') + .mockImplementation(() => {}) + const retry = vi + .spyOn(ChatClient.prototype, 'retryInterrupts') + .mockImplementation(() => {}) + const unsafe = vi + .spyOn(ChatClient.prototype, 'resumeInterruptsUnsafe') + .mockResolvedValue(true) + const { result } = renderInjectChat({ + connection: createMockConnectionAdapter(), + }) + const resolver = () => undefined + const resume = [{ interruptId: 'one', status: 'cancelled' as const }] + + result.resolveInterrupts(resolver) + result.cancelInterrupts() + result.retryInterrupts() + await expect(result.resumeInterruptsUnsafe(resume)).resolves.toBe(true) + + expect(resolve).toHaveBeenCalledWith(resolver) + expect(cancel).toHaveBeenCalledOnce() + expect(retry).toHaveBeenCalledOnce() + expect(unsafe).toHaveBeenCalledWith(resume, undefined) + }) + }) + it('initializes with default state', () => { const adapter = createMockConnectionAdapter() const { result } = renderInjectChat({ connection: adapter }) diff --git a/packages/ai-angular/tests/test-utils.ts b/packages/ai-angular/tests/test-utils.ts index 9e217a942..b27b72497 100644 --- a/packages/ai-angular/tests/test-utils.ts +++ b/packages/ai-angular/tests/test-utils.ts @@ -7,6 +7,7 @@ import { import { injectChat } from '../src/inject-chat' import type { InjectChatOptions } from '../src/types' import type { InjectChatResult } from '../src/types' +import type { ChatResumeSnapshotV2 } from '@tanstack/ai-client' export { createMockConnectionAdapter, @@ -14,6 +15,80 @@ export { createToolCallChunks, } from '../../ai-client/tests/test-utils' +export function createInterruptResumeSnapshot(): ChatResumeSnapshotV2 { + const pendingInterrupts = [ + { + id: 'staged-interrupt', + reason: 'confirmation', + metadata: { + 'tanstack:interruptBinding': { + kind: 'generic' as const, + interruptId: 'staged-interrupt', + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'none', + }, + }, + }, + { + id: 'invalid-interrupt', + reason: 'confirmation', + metadata: { + 'tanstack:interruptBinding': { + kind: 'generic' as const, + interruptId: 'invalid-interrupt', + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'none', + }, + }, + }, + ] + + return { + schemaVersion: 2, + resumeState: { threadId: 'thread-1', runId: 'run-1' }, + pendingInterrupts, + interruptState: { + recoveryState: { + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + pendingInterrupts, + }, + drafts: [ + { + interruptId: 'staged-interrupt', + response: { + interruptId: 'staged-interrupt', + status: 'resolved', + payload: { answer: 42 }, + }, + status: 'staged', + }, + { + interruptId: 'invalid-interrupt', + response: null, + status: 'error', + error: { + scope: 'item', + interruptId: 'invalid-interrupt', + code: 'invalid-payload', + message: 'Invalid persisted response', + source: 'client', + retryable: false, + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + }, + }, + ], + }, + } +} + // Ensure TestBed is initialized in this module's scope, regardless of whether // the setup file's initialization was in a different module context (possible // when the Angular plugin creates separate ESM module instances for compiled diff --git a/packages/ai-anthropic/tests/anthropic-adapter.test.ts b/packages/ai-anthropic/tests/anthropic-adapter.test.ts index 43d3ded61..dd8ae2a3a 100644 --- a/packages/ai-anthropic/tests/anthropic-adapter.test.ts +++ b/packages/ai-anthropic/tests/anthropic-adapter.test.ts @@ -1,8 +1,12 @@ import { describe, it, expect, beforeEach, vi } from 'vitest' import { + InterruptPersistenceCapability, chat, + defineChatMiddleware, + provideInterruptPersistence, StreamProcessor, type Tool, + type InterruptPersistenceGateway, type StreamChunk, type UIMessage, } from '@tanstack/ai' @@ -55,6 +59,43 @@ const weatherTool: Tool = { }), } +const testInterruptGateway: InterruptPersistenceGateway = { + openInterruptBatch: async (input) => ({ + generation: 1, + descriptors: input.descriptors, + }), + commitInterruptResolutions: async (input) => ({ + status: 'committed', + continuationRunId: input.continuationRunId, + }), + getInterruptRecoveryState: async (input) => ({ + schemaVersion: 1, + state: 'missing', + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + generation: input.knownGeneration, + pendingInterrupts: [], + }), +} + +const testInterruptPersistence = defineChatMiddleware({ + name: 'anthropic-test-interrupt-persistence', + provides: [InterruptPersistenceCapability], + setup(ctx) { + provideInterruptPersistence(ctx, testInterruptGateway) + }, + async onChunk(_ctx, chunk) { + if (chunk.type === 'RUN_FINISHED' && chunk.outcome?.type === 'interrupt') { + await testInterruptGateway.openInterruptBatch({ + threadId: chunk.threadId, + interruptedRunId: chunk.runId, + descriptors: chunk.outcome.interrupts, + bindings: [], + }) + } + }, +}) + function createTextStream(text: string) { return (async function* () { yield { @@ -1715,6 +1756,7 @@ describe('Anthropic stream processing', () => { adapter, messages: [{ role: 'user', content: 'Weather in Berlin?' }], tools: [weatherTool], + middleware: [testInterruptPersistence], })) { chunks.push(chunk) } diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index fc3a078cf..cbf20a717 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -11,12 +11,18 @@ import { fetcherToConnectionAdapter, getChunkRunId, normalizeConnectionAdapter, + parseInterruptRecoveryState, } from './connection-adapters' import { ChatPersistenceController } from './chat-persistence-controller' import { ClearedStreamTracker } from './cleared-stream-tracker' +import { InterruptManager } from './interrupt-manager' import type { AnyClientTool, ContentPart, + InterruptCommitResult, + InterruptRecoveryQuery, + InterruptRecoveryStateV1, + InterruptSubmissionError, ModelMessage, RunAgentResumeItem, StreamChunk, @@ -35,18 +41,24 @@ import type { ChatDevtoolsBridgeOptions, } from './devtools' import type { + BoundInterrupts, ChatClientOptions, ChatClientState, + ChatContinuationLoader, ChatFetcher, + ChatInterrupt, + ChatInterruptState, ChatPendingInterrupt, ChatResumeSnapshot, ChatResumeState, ConnectionStatus, MessagePart, MultimodalContent, + PersistedInterruptDraft, ToolCallPart, UIMessage, } from './types' +import type { InterruptManagerSubmission } from './interrupt-manager' type ChatClientUpdateOptionsWithoutContext< TTools extends ReadonlyArray, @@ -66,8 +78,9 @@ type ChatClientUpdateOptionsWithoutContext< onSessionGeneratingChange?: (isGenerating: boolean) => void onResumeStateChange?: ( resumeState: ChatResumeState | null, - pendingInterrupts: Array, + pendingInterrupts: BoundInterrupts, ) => void + onInterruptStateChange?: (state: ChatInterruptState) => void onCustomEvent?: ( eventType: string, data: unknown, @@ -98,6 +111,110 @@ function resolveTransport(transport: { throw new Error('ChatClient: either `connection` or `fetcher` is required.') } +function replayContinuationRunId(chunk: StreamChunk): string | undefined { + if (chunk.type !== 'RUN_FINISHED' || !('result' in chunk)) return undefined + const result: unknown = chunk.result + if (result === null || typeof result !== 'object' || Array.isArray(result)) { + return undefined + } + if (!('replayed' in result) || result.replayed !== true) return undefined + return 'continuationRunId' in result && + typeof result.continuationRunId === 'string' + ? result.continuationRunId + : undefined +} + +interface ParsedPersistedInterruptState { + recoveryState: InterruptRecoveryStateV1 + drafts: ReadonlyArray +} + +interface PersistedRecoveryRequest { + query: InterruptRecoveryQuery + drafts: ReadonlyArray + expectedGeneration?: number + expectedInterruptIds: ReadonlyArray +} + +function readResumeState( + snapshot: ChatResumeSnapshot, +): ChatResumeState | undefined { + const value: unknown = snapshot + if ( + value === null || + typeof value !== 'object' || + !('resumeState' in value) + ) { + return undefined + } + const resumeState = value.resumeState + if ( + resumeState === null || + typeof resumeState !== 'object' || + !('threadId' in resumeState) || + typeof resumeState.threadId !== 'string' || + resumeState.threadId.length === 0 || + !('runId' in resumeState) || + typeof resumeState.runId !== 'string' || + resumeState.runId.length === 0 + ) { + return undefined + } + return { threadId: resumeState.threadId, runId: resumeState.runId } +} + +function readPersistedInterruptState( + snapshot: ChatResumeSnapshot, +): ParsedPersistedInterruptState | undefined { + if (snapshot.schemaVersion !== 2) { + return undefined + } + const interruptStateValue: unknown = snapshot.interruptState + if ( + interruptStateValue === undefined || + interruptStateValue === null || + typeof interruptStateValue !== 'object' || + Array.isArray(interruptStateValue) + ) { + return undefined + } + const interruptState = interruptStateValue as Record + let recoveryState: InterruptRecoveryStateV1 + try { + recoveryState = parseInterruptRecoveryState(interruptState.recoveryState) + } catch { + return undefined + } + const drafts = interruptState.drafts + if ( + recoveryState.threadId !== snapshot.resumeState.threadId || + recoveryState.interruptedRunId !== snapshot.resumeState.runId || + !Array.isArray(drafts) + ) { + return undefined + } + if ( + !drafts.every( + (draft) => + draft !== null && + typeof draft === 'object' && + typeof draft.interruptId === 'string' && + ['pending', 'validating', 'staged', 'submitting', 'error'].includes( + draft.status, + ) && + 'response' in draft && + (draft.error === undefined || + (draft.error !== null && typeof draft.error === 'object')), + ) + ) { + return undefined + } + return { + recoveryState, + drafts: drafts as ReadonlyArray, + } +} + export class ChatClient< TTools extends ReadonlyArray = any, TContext = unknown, @@ -118,15 +235,23 @@ export class ChatClient< // run, so approvals/client-tool results can be sent back. Cleared when the // run terminates. This is STATE (interrupt) resume, not delivery/cursor. private lastResume: ChatResumeState | null = null - private pendingInterrupts: Array = [] - private pendingInterruptRunId: string | null = null - private readonly pendingInterruptResumeItems = new Map< - string, - RunAgentResumeItem - >() + private readonly interruptManager: InterruptManager + private readonly continuationLoader: ChatContinuationLoader | undefined + private activeInterruptSubmission = false + private pendingInterruptReplayRunId: string | undefined + private interruptSubmissionFailure: + | { + errors: ReadonlyArray + recovery?: InterruptRecoveryStateV1 + } + | undefined + private readonly joinedRunWaiters = new Map void>() + private readonly recoveryReady: boolean + private pendingInitialRecovery: InterruptRecoveryStateV1 | undefined + private pendingPersistedRecovery: PersistedRecoveryRequest | undefined // When set, the next streamResponse() continues this interrupted run instead // of starting a fresh run (consumed once). - private pendingResumeRunId: string | null = null + private pendingResumeParentRunId: string | null = null private pendingResumeThreadId: string | null = null private pendingResumeItems: Array | null = null private activeResumeThreadId: string | null = null @@ -190,8 +315,9 @@ export class ChatClient< onSessionGeneratingChange: (isGenerating: boolean) => void onResumeStateChange: ( resumeState: ChatResumeState | null, - pendingInterrupts: Array, + pendingInterrupts: BoundInterrupts, ) => void + onInterruptStateChange: (state: ChatInterruptState) => void onCustomEvent: ( eventType: string, data: unknown, @@ -243,10 +369,19 @@ export class ChatClient< onSessionGeneratingChange: options.onSessionGeneratingChange || (() => {}), onResumeStateChange: options.onResumeStateChange || (() => {}), + onInterruptStateChange: options.onInterruptStateChange || (() => {}), onCustomEvent: options.onCustomEvent || (() => {}), }, } + this.interruptManager = new InterruptManager({ + ...(options.tools !== undefined ? { tools: options.tools } : {}), + submit: (submission) => this.submitInterruptBatch(submission), + recover: (state) => this.recoverInterrupts(state), + onChange: () => this.notifyResumeStateChange(), + }) + this.continuationLoader = options.continuationLoader + this.persistenceController = new ChatPersistenceController({ chatId: this.uniqueId, threadId: this.threadId, @@ -486,6 +621,17 @@ export class ChatClient< if (!options.initialResumeSnapshot) { this.persistenceController.hydrateResumeSnapshot(storedResumeSnapshot) } + this.recoveryReady = true + if (this.pendingInitialRecovery !== undefined) { + const recovery = this.pendingInitialRecovery + this.pendingInitialRecovery = undefined + this.startInitialRecovery(recovery) + } + if (this.pendingPersistedRecovery !== undefined) { + const recovery = this.pendingPersistedRecovery + this.pendingPersistedRecovery = undefined + this.startPersistedRecovery(recovery) + } } /** @@ -508,10 +654,80 @@ export class ChatClient< } private applyResumeSnapshot(snapshot: ChatResumeSnapshot): void { - this.lastResume = { ...snapshot.resumeState } - this.pendingInterrupts = [...(snapshot.pendingInterrupts ?? [])] - this.pendingInterruptRunId = - this.pendingInterrupts.length > 0 ? this.lastResume.runId : null + const resumeState = readResumeState(snapshot) + if (resumeState === undefined) { + this.interruptManager.reset() + return + } + this.lastResume = resumeState + const persistedState = readPersistedInterruptState(snapshot) + if (persistedState !== undefined) { + const { recoveryState, drafts } = persistedState + if (recoveryState.state === 'pending') { + this.interruptManager.hydrate({ + threadId: recoveryState.threadId, + interruptedRunId: recoveryState.interruptedRunId, + generation: recoveryState.generation, + interrupts: recoveryState.pendingInterrupts, + }) + this.interruptManager.restorePersistedDrafts(drafts) + this.schedulePersistedRecovery({ + query: { + threadId: recoveryState.threadId, + interruptedRunId: recoveryState.interruptedRunId, + knownGeneration: recoveryState.generation, + }, + drafts, + expectedGeneration: recoveryState.generation, + expectedInterruptIds: recoveryState.pendingInterrupts.map( + (interrupt) => interrupt.id, + ), + }) + return + } + this.interruptManager.reset() + if (this.recoveryReady) { + this.startInitialRecovery(recoveryState) + } else { + this.pendingInitialRecovery = recoveryState + } + return + } + const pendingInterrupts = Array.isArray(snapshot.pendingInterrupts) + ? snapshot.pendingInterrupts + : [] + if (pendingInterrupts.length === 0) { + this.interruptManager.reset() + return + } + const generation = this.interruptGeneration(pendingInterrupts) + this.interruptManager.hydrate({ + threadId: resumeState.threadId, + interruptedRunId: resumeState.runId, + generation, + interrupts: pendingInterrupts, + }) + this.schedulePersistedRecovery({ + query: { + threadId: resumeState.threadId, + interruptedRunId: resumeState.runId, + knownGeneration: generation, + }, + drafts: [], + expectedGeneration: generation, + expectedInterruptIds: pendingInterrupts.map((interrupt) => interrupt.id), + }) + } + + private schedulePersistedRecovery(request: PersistedRecoveryRequest): void { + if (this.connection.loadInterruptState === undefined) { + return + } + if (this.recoveryReady) { + this.startPersistedRecovery(request) + } else { + this.pendingPersistedRecovery = request + } } mountDevtools(): void { @@ -585,6 +801,15 @@ export class ChatClient< if (chunk.type !== 'RUN_FINISHED' && chunk.type !== 'RUN_ERROR') { return } + + if (this.activeInterruptSubmission) { + const continuationRunId = replayContinuationRunId(chunk) + if (continuationRunId !== undefined) { + this.pendingInterruptReplayRunId = continuationRunId + return + } + if (chunk.type === 'RUN_ERROR') return + } const runId = getChunkRunId(chunk) const threadId = 'threadId' in chunk && typeof chunk.threadId === 'string' @@ -600,10 +825,12 @@ export class ChatClient< threadId: threadId ?? this.threadId, runId: interruptedRunId, } - this.pendingInterruptRunId = interruptedRunId - this.pendingInterrupts = [...chunk.outcome.interrupts] - this.pendingInterruptResumeItems.clear() - this.notifyResumeStateChange() + this.interruptManager.hydrate({ + threadId: this.lastResume.threadId, + interruptedRunId, + generation: this.interruptGeneration(chunk.outcome.interrupts), + interrupts: chunk.outcome.interrupts, + }) return } @@ -611,9 +838,6 @@ export class ChatClient< const isTrackedRunTerminal = Boolean( runId && this.lastResume?.runId === runId, ) - const isPendingInterruptRunTerminal = Boolean( - runId && this.pendingInterruptRunId === runId, - ) const isCurrentRunTerminal = Boolean( (runId && this.currentRunId === runId) || (this.currentRunId && this.lastResume?.runId === this.currentRunId), @@ -623,14 +847,12 @@ export class ChatClient< if ( isRunlessSessionError || isTrackedRunTerminal || - isPendingInterruptRunTerminal || isCurrentRunTerminal || isCurrentStreamTerminal ) { this.lastResume = null - this.pendingInterrupts = [] - this.pendingInterruptRunId = null - this.pendingInterruptResumeItems.clear() + this.interruptManager.reset() + return } this.notifyResumeStateChange() } @@ -644,22 +866,343 @@ export class ChatClient< return this.lastResume ? { ...this.lastResume } : null } - getPendingInterrupts(): Array { - return [...this.pendingInterrupts] + getInterruptState(): ChatInterruptState { + return this.interruptManager.getState() } - resumeInterrupts( + getInterrupts(): BoundInterrupts { + return this.interruptManager.getInterrupts() + } + + /** @deprecated Use getInterrupts(). */ + getPendingInterrupts(): BoundInterrupts { + return this.interruptManager.getInterrupts() + } + + resolveInterrupts(approved: boolean): void + resolveInterrupts( + resolver: (interrupt: ChatInterrupt) => undefined, + ): void + resolveInterrupts( + resolution: boolean | ((interrupt: ChatInterrupt) => undefined), + ): void { + if (typeof resolution === 'boolean') { + this.interruptManager.resolve(resolution) + } else { + this.interruptManager.resolve(resolution) + } + } + + cancelInterrupts(): void { + this.interruptManager.cancel() + } + + retryInterrupts(): void { + this.interruptManager.retry() + } + + /** Unsafe low-level resume escape hatch. Prefer bound interrupt methods. */ + resumeInterruptsUnsafe( resume: Array, state?: ChatResumeState, ): Promise { const target = state ?? this.lastResume if (!target || this.isLoading) return Promise.resolve(false) this.pendingResumeThreadId = target.threadId - this.pendingResumeRunId = target.runId + this.pendingResumeParentRunId = target.runId this.pendingResumeItems = [...resume] return this.streamResponse() } + /** @deprecated Use bound interrupt methods or resumeInterruptsUnsafe(). */ + resumeInterrupts( + resume: Array, + state?: ChatResumeState, + ): Promise { + return this.resumeInterruptsUnsafe(resume, state) + } + + private async submitInterruptBatch( + submission: InterruptManagerSubmission, + ): Promise { + this.activeInterruptSubmission = true + this.pendingInterruptReplayRunId = undefined + this.interruptSubmissionFailure = undefined + const resumed = await this.resumeInterruptsUnsafe( + [...submission.resolutions], + { + threadId: submission.threadId, + runId: submission.interruptedRunId, + }, + ).finally(() => { + this.activeInterruptSubmission = false + }) + const replayRunId = this.takeInterruptReplayRunId() + const failure = this.takeInterruptSubmissionFailure() + if (failure?.recovery !== undefined) { + return { status: 'conflict', authoritativeState: failure.recovery } + } + if (failure !== undefined) { + throw { errors: failure.errors } + } + if (replayRunId !== undefined) { + await this.joinContinuationRun(replayRunId, { + schemaVersion: 1, + state: 'committed', + threadId: submission.threadId, + interruptedRunId: submission.interruptedRunId, + generation: submission.generation, + pendingInterrupts: [], + committed: { + fingerprint: submission.fingerprint, + continuationRunId: replayRunId, + committedAt: new Date().toISOString(), + }, + }) + return { status: 'replayed', continuationRunId: replayRunId } + } + if (!resumed) { + throw new Error('Interrupt continuation could not be started.') + } + } + + private takeInterruptSubmissionFailure(): + | { + errors: ReadonlyArray + recovery?: InterruptRecoveryStateV1 + } + | undefined { + const failure = this.interruptSubmissionFailure + this.interruptSubmissionFailure = undefined + return failure + } + + private takeInterruptReplayRunId(): string | undefined { + const runId = this.pendingInterruptReplayRunId + this.pendingInterruptReplayRunId = undefined + return runId + } + + private async recoverInterrupts( + authoritativeState?: InterruptRecoveryStateV1, + ): Promise { + let state = authoritativeState + if (state === undefined) { + const current = this.interruptManager.getRecoveryState() + if ( + current === undefined || + this.connection.loadInterruptState === undefined + ) { + this.interruptManager.reportRecoveryError( + 'recovery-unavailable', + 'Authoritative interrupt recovery is unavailable for this connection.', + current, + ) + return + } + try { + state = await this.connection.loadInterruptState({ + threadId: current.threadId, + interruptedRunId: current.interruptedRunId, + knownGeneration: current.generation, + }) + } catch { + this.interruptManager.reportRecoveryError( + 'recovery-unavailable', + 'Authoritative interrupt recovery failed.', + current, + ) + return + } + } + + if (state.state === 'pending') { + this.lastResume = { + threadId: state.threadId, + runId: state.interruptedRunId, + } + this.interruptManager.hydrate({ + threadId: state.threadId, + interruptedRunId: state.interruptedRunId, + generation: state.generation, + interrupts: state.pendingInterrupts, + }) + return + } + + if (state.state === 'committed' || state.state === 'legacy-committed') { + const continuationRunId = state.committed?.continuationRunId + if (continuationRunId === undefined) { + this.interruptManager.reportRecoveryError( + 'recovery-unavailable', + 'The committed interrupt state has no continuation run to join.', + state, + ) + return + } + await this.joinContinuationRun(continuationRunId, state, { + preserveRootErrors: this.interruptManager + .getInterruptErrors() + .some((error) => error.code === 'conflict'), + }) + return + } + + this.interruptManager.reportRecoveryError( + state.state === 'expired' ? 'expired' : 'stale', + state.state === 'expired' + ? 'The interrupt batch has expired.' + : 'The interrupt batch no longer exists.', + state, + ) + } + + private startInitialRecovery(recoveryState: InterruptRecoveryStateV1): void { + // Constructors cannot await replay. The structured manager error keeps an + // unexpected async failure observable without turning hydration into an + // unhandled promise rejection. + void this.recoverInterrupts(recoveryState).catch(() => { + this.interruptManager.reportRecoveryError( + 'recovery-unavailable', + 'Initial interrupt recovery failed.', + recoveryState, + ) + }) + } + + private startPersistedRecovery(request: PersistedRecoveryRequest): void { + void this.recoverPersistedInterrupts(request).catch(() => { + this.interruptManager.reportRecoveryError( + 'recovery-unavailable', + 'Persisted interrupt recovery failed.', + this.persistedRecoveryErrorState(request.query), + ) + }) + } + + private async recoverPersistedInterrupts( + request: PersistedRecoveryRequest, + ): Promise { + if (this.connection.loadInterruptState === undefined) { + this.interruptManager.reportRecoveryError( + 'recovery-unavailable', + 'Authoritative interrupt recovery is unavailable for this connection.', + this.persistedRecoveryErrorState(request.query), + ) + return + } + + let state: InterruptRecoveryStateV1 + try { + state = parseInterruptRecoveryState( + await this.connection.loadInterruptState(request.query), + ) + if ( + state.threadId !== request.query.threadId || + state.interruptedRunId !== request.query.interruptedRunId + ) { + throw new TypeError('Interrupt recovery correlation did not match.') + } + } catch { + this.interruptManager.reportRecoveryError( + 'recovery-unavailable', + 'Authoritative interrupt recovery failed.', + this.persistedRecoveryErrorState(request.query), + ) + return + } + + await this.recoverInterrupts(state) + if ( + state.state === 'pending' && + request.expectedGeneration === state.generation && + request.expectedInterruptIds.length === state.pendingInterrupts.length && + request.expectedInterruptIds.every( + (id, index) => state.pendingInterrupts[index]?.id === id, + ) + ) { + this.interruptManager.restorePersistedDrafts(request.drafts) + } + } + + private persistedRecoveryErrorState( + query: InterruptRecoveryQuery, + ): InterruptRecoveryStateV1 { + return { + schemaVersion: 1, + state: 'missing', + threadId: query.threadId, + interruptedRunId: query.interruptedRunId, + generation: query.knownGeneration, + pendingInterrupts: [], + } + } + + private async joinContinuationRun( + continuationRunId: string, + recoveryState: InterruptRecoveryStateV1, + options?: { preserveRootErrors?: boolean }, + ): Promise { + try { + if (this.connection.joinRun !== undefined) { + this.ensureSubscription() + const processed = new Promise((resolve) => { + this.joinedRunWaiters.set(continuationRunId, resolve) + }) + await this.connection.joinRun(continuationRunId) + await processed + } else if (this.continuationLoader !== undefined) { + for await (const chunk of this.continuationLoader(continuationRunId)) { + await this.processIncomingChunk(chunk) + } + } else { + this.interruptManager.reportRecoveryError( + 'recovery-unavailable', + 'A continuation loader or joinRun connection operation is required.', + recoveryState, + ) + return + } + } catch { + this.joinedRunWaiters.delete(continuationRunId) + this.interruptManager.reportRecoveryError( + 'recovery-unavailable', + 'The committed continuation run could not be replayed.', + recoveryState, + ) + return + } + this.lastResume = null + this.interruptManager.reset({ + preserveRootErrors: options?.preserveRootErrors === true, + }) + } + + private interruptGeneration( + interrupts: ReadonlyArray, + ): number { + let generation: number | undefined + for (const interrupt of interrupts) { + const candidate: unknown = + interrupt.metadata?.['tanstack:interruptBinding'] + if ( + candidate === null || + typeof candidate !== 'object' || + !('generation' in candidate) || + typeof candidate.generation !== 'number' || + !Number.isInteger(candidate.generation) || + candidate.generation < 0 + ) { + return 0 + } + if (generation !== undefined && generation !== candidate.generation) { + return 0 + } + generation = candidate.generation + } + return generation ?? 0 + } + private generateUniqueId(prefix: string): string { return `${prefix}-${Date.now()}-${Math.random().toString(36).substring(7)}` } @@ -697,13 +1240,31 @@ export class ChatClient< private notifyResumeStateChange(): void { const resumeState = this.getResumeState() - const pendingInterrupts = this.getPendingInterrupts() + const pendingInterrupts = [...this.interruptManager.getDescriptors()] + const recoveryState = this.interruptManager.getRecoveryState() this.persistenceController.persistResumeSnapshot( - resumeState ? { resumeState, pendingInterrupts } : null, + resumeState + ? { + schemaVersion: 2, + resumeState, + pendingInterrupts, + ...(recoveryState !== undefined + ? { + interruptState: { + recoveryState, + drafts: this.interruptManager.getPersistedDrafts(), + }, + } + : {}), + } + : null, ) this.callbacksRef.current.onResumeStateChange( resumeState, - pendingInterrupts, + this.interruptManager.getInterrupts(), + ) + this.callbacksRef.current.onInterruptStateChange( + this.interruptManager.getState(), ) } @@ -859,33 +1420,54 @@ export class ChatClient< const stream = this.connection.subscribe(signal) for await (const chunk of stream) { if (signal.aborted) break - if (this.connectionStatus === 'connecting') { - this.setConnectionStatus('connected') + await this.processIncomingChunk(chunk) + } + } + + private async processIncomingChunk(chunk: StreamChunk): Promise { + if ( + chunk.type === 'RUN_ERROR' && + this.activeInterruptSubmission && + (chunk['tanstack:interruptErrors']?.length ?? 0) > 0 + ) { + this.interruptSubmissionFailure = { + errors: chunk['tanstack:interruptErrors'] ?? [], + ...(chunk['tanstack:interruptRecovery'] !== undefined + ? { recovery: chunk['tanstack:interruptRecovery'] } + : {}), } - const shouldIgnore = this.clearedStreamTracker.shouldIgnoreChunk(chunk) - if (shouldIgnore) { - if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') { - if (getChunkRunId(chunk)) { - this.updateRunLifecycle(chunk, { resolveProcessing: false }) - } - this.retireIgnoredClearedTerminalChunk(chunk) + } + if (this.connectionStatus === 'connecting') { + this.setConnectionStatus('connected') + } + const shouldIgnore = this.clearedStreamTracker.shouldIgnoreChunk(chunk) + if (shouldIgnore) { + if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') { + if (getChunkRunId(chunk)) { + this.updateRunLifecycle(chunk, { resolveProcessing: false }) } - continue + this.retireIgnoredClearedTerminalChunk(chunk) + this.resolveJoinedRun(chunk) } - this.callbacksRef.current.onChunk(chunk) - this.devtoolsBridge.observeChunk(chunk) - this.processor.processChunk(chunk) - // Run lifecycle (active-run tracking, session-generating state, and - // processing resolution for RUN_FINISHED / RUN_ERROR) is handled in a - // single place so the ignored-chunk path above and this path can't - // diverge. RUN_ERROR carries its runId via the AG-UI passthrough so a - // per-run error only clears that run, while a runId-less RUN_ERROR is - // treated as a session-level error that clears every active run. - this.updateRunLifecycle(chunk) - this.observeInterruptState(chunk) - // Yield control back to event loop for UI updates - await new Promise((resolve) => setTimeout(resolve, 0)) + return } + this.callbacksRef.current.onChunk(chunk) + this.devtoolsBridge.observeChunk(chunk) + this.processor.processChunk(chunk) + this.updateRunLifecycle(chunk) + this.observeInterruptState(chunk) + await new Promise((resolve) => setTimeout(resolve, 0)) + this.resolveJoinedRun(chunk) + } + + private resolveJoinedRun(chunk: StreamChunk): void { + if (chunk.type !== 'RUN_FINISHED' && chunk.type !== 'RUN_ERROR') return + const runId = getChunkRunId(chunk) + if (runId === undefined) return + const resolve = this.joinedRunWaiters.get(runId) + if (resolve === undefined) return + this.joinedRunWaiters.delete(runId) + resolve() } /** @@ -964,7 +1546,7 @@ export class ChatClient< if (emptyMessage || this.isLoading) { return } - if (this.pendingInterrupts.length > 0 && this.lastResume) { + if (this.interruptManager.getInterrupts().length > 0 && this.lastResume) { throw new Error( 'ChatClient: cannot send normal input while pending interrupts exist. Use resumeInterrupts() instead.', ) @@ -1004,7 +1586,7 @@ export class ChatClient< */ async append(message: UIMessage | ModelMessage): Promise { this.mountDevtools() - if (this.pendingInterrupts.length > 0 && this.lastResume) { + if (this.interruptManager.getInterrupts().length > 0 && this.lastResume) { throw new Error( 'ChatClient: cannot append normal input while pending interrupts exist. Use resumeInterrupts() instead.', ) @@ -1051,20 +1633,18 @@ export class ChatClient< // Track generation so a superseded stream's cleanup doesn't clobber the new one const generation = ++this.streamGeneration - // Resuming an interrupt reuses the original runId so the server continues - // that run with the resume decisions. + // Native interrupt continuation is a fresh child run. The interrupted run + // is carried as parentRunId and the complete resolution batch as resume. const resumeThreadId = this.pendingResumeThreadId - const resumeRunId = this.pendingResumeRunId + const resumeParentRunId = this.pendingResumeParentRunId const resumeItems = this.pendingResumeItems const isResumeRequest = Boolean( - resumeThreadId || resumeRunId || resumeItems, + resumeThreadId || resumeParentRunId || resumeItems, ) this.pendingResumeThreadId = null - this.pendingResumeRunId = null + this.pendingResumeParentRunId = null this.pendingResumeItems = null - const runId = - resumeRunId ?? - `run-${Date.now()}-${Math.random().toString(36).slice(2, 8)}` + const runId = `run-${Date.now()}-${Math.random().toString(36).slice(2, 8)}` this.currentRunId = runId this.activeResumeThreadId = resumeThreadId ?? this.threadId this.activeResumeRunId = runId @@ -1149,6 +1729,9 @@ export class ChatClient< const runContext = { threadId: resumeThreadId ?? this.threadId, runId, + ...(resumeParentRunId !== null + ? { parentRunId: resumeParentRunId } + : {}), clientTools: Array.from(clientTools.values()).map((t) => ({ name: t.name, description: t.description, @@ -1400,13 +1983,10 @@ export class ChatClient< this.processor.clearMessages() this.persistenceController.removeMessages() this.lastResume = null - this.pendingInterrupts = [] - this.pendingInterruptRunId = null - this.pendingInterruptResumeItems.clear() + this.interruptManager.reset() this.pendingResumeThreadId = null - this.pendingResumeRunId = null + this.pendingResumeParentRunId = null this.pendingResumeItems = null - this.notifyResumeStateChange() this.setError(undefined) this.events.messagesCleared() } @@ -1424,6 +2004,17 @@ export class ChatClient< clientTool: AnyClientTool | undefined, context?: ChatClientRunEventContext, ): Promise { + if ( + this.interruptManager.resolveClientToolOutput( + result.toolCallId, + result.state === 'output-error' + ? { error: result.errorText || 'Tool execution failed' } + : result.output, + ) + ) { + return + } + if (clientTool && result.state !== 'output-error') { try { result = { @@ -1459,35 +2050,6 @@ export class ChatClient< : undefined, ) - const pendingInterrupt = this.pendingInterrupts.find( - (interrupt) => - interrupt.toolCallId === result.toolCallId || - interrupt.id === result.toolCallId, - ) - if (pendingInterrupt && this.lastResume) { - const resumeItem: RunAgentResumeItem = { - interruptId: pendingInterrupt.id, - status: 'resolved', - payload: - result.state === 'output-error' - ? { error: result.errorText || 'Tool execution failed' } - : result.output, - } - const resumeItems = this.collectPendingInterruptResumeItems(resumeItem) - if (!resumeItems) { - return - } - if (this.isLoading) { - this.queuePostStreamAction(async () => { - await this.resumeInterrupts(resumeItems) - }) - return - } - - await this.resumeInterrupts(resumeItems) - return - } - // If stream is in progress, queue continuation check for after it ends if (this.isLoading) { this.queuePostStreamAction(() => this.checkForContinuation()) @@ -1515,9 +2077,14 @@ export class ChatClient< id: string // approval.id, not toolCallId approved: boolean }): Promise { - const pendingInterrupt = this.pendingInterrupts.find( - (interrupt) => interrupt.id === response.id, - ) + if ( + this.interruptManager.resolveToolApprovalDecision( + response.id, + response.approved, + ) + ) { + return + } // Find the tool call ID from the approval ID const messages = this.processor.getMessages() let foundToolCallId: string | undefined @@ -1545,25 +2112,6 @@ export class ChatClient< this.processor.addToolApprovalResponse(response.id, response.approved) this.devtoolsBridge.emitSnapshot() - if (pendingInterrupt && this.lastResume) { - const resumeItems = this.collectPendingInterruptResumeItems({ - interruptId: response.id, - status: response.approved ? 'resolved' : 'cancelled', - payload: { approved: response.approved }, - }) - if (!resumeItems) { - return - } - if (this.isLoading) { - this.queuePostStreamAction(async () => { - await this.resumeInterrupts(resumeItems) - }) - return - } - await this.resumeInterrupts(resumeItems) - return - } - // If stream is in progress, queue continuation check for after it ends if (this.isLoading) { this.queuePostStreamAction(() => this.checkForContinuation()) @@ -1573,23 +2121,6 @@ export class ChatClient< await this.checkForContinuation() } - private collectPendingInterruptResumeItems( - resumeItem: RunAgentResumeItem, - ): Array | null { - this.pendingInterruptResumeItems.set(resumeItem.interruptId, resumeItem) - const resumeItems: Array = [] - for (const interrupt of this.pendingInterrupts) { - const answeredInterrupt = this.pendingInterruptResumeItems.get( - interrupt.id, - ) - if (!answeredInterrupt) { - return null - } - resumeItems.push(answeredInterrupt) - } - return resumeItems - } - /** * Queue an action to be executed after the current stream ends */ @@ -1775,6 +2306,7 @@ export class ChatClient< this.context = options.context } if (options.tools !== undefined) { + this.interruptManager.updateTools(options.tools) this.clientToolsRef.current = new Map() for (const tool of options.tools) { this.clientToolsRef.current.set(tool.name, tool) @@ -1810,6 +2342,10 @@ export class ChatClient< this.callbacksRef.current.onResumeStateChange = options.onResumeStateChange } + if (options.onInterruptStateChange !== undefined) { + this.callbacksRef.current.onInterruptStateChange = + options.onInterruptStateChange + } if (options.onCustomEvent !== undefined) { this.callbacksRef.current.onCustomEvent = options.onCustomEvent } diff --git a/packages/ai-client/src/chat-persistence-controller.ts b/packages/ai-client/src/chat-persistence-controller.ts index 91eabb3fe..847a9b85f 100644 --- a/packages/ai-client/src/chat-persistence-controller.ts +++ b/packages/ai-client/src/chat-persistence-controller.ts @@ -173,6 +173,21 @@ export class ChatPersistenceController< this.resumeGeneration++ this.desiredResumeSnapshot = snapshot ? { + ...(snapshot.schemaVersion === 2 + ? { + schemaVersion: 2 as const, + ...(snapshot.interruptState !== undefined + ? { + interruptState: { + recoveryState: snapshot.interruptState.recoveryState, + drafts: [...snapshot.interruptState.drafts], + }, + } + : {}), + } + : snapshot.schemaVersion === 1 + ? { schemaVersion: 1 as const } + : {}), resumeState: { ...snapshot.resumeState }, pendingInterrupts: [...(snapshot.pendingInterrupts ?? [])], } diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 8f19d255b..9eac6713c 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -5,6 +5,9 @@ import { } from './response-stream' import { parseSseDataLine } from './sse-utils' import type { + Interrupt, + InterruptRecoveryQuery, + InterruptRecoveryStateV1, ModelMessage, RunAgentResumeItem, RunErrorEvent, @@ -376,6 +379,16 @@ export interface ConnectConnectionAdapter { abortSignal?: AbortSignal, runContext?: RunAgentInputContext, ) => AsyncIterable + /** Read-only replay of an existing continuation run. */ + joinRun?: ( + runId: string, + abortSignal?: AbortSignal, + ) => AsyncIterable + /** Explicit authoritative interrupt recovery operation. */ + loadInterruptState?: ( + query: InterruptRecoveryQuery, + abortSignal?: AbortSignal, + ) => Promise } /** @@ -408,6 +421,12 @@ export interface SubscribeConnectionAdapter { abortSignal?: AbortSignal, runContext?: RunAgentInputContext, ) => Promise + /** Pumps a read-only continuation replay through subscribe(). */ + joinRun?: (runId: string, abortSignal?: AbortSignal) => Promise + loadInterruptState?: ( + query: InterruptRecoveryQuery, + abortSignal?: AbortSignal, + ) => Promise } /** @@ -445,6 +464,14 @@ export function normalizeConnectionAdapter( return { subscribe: connection.subscribe.bind(connection), send: connection.send.bind(connection), + ...(connection.joinRun !== undefined + ? { joinRun: connection.joinRun.bind(connection) } + : {}), + ...(connection.loadInterruptState !== undefined + ? { + loadInterruptState: connection.loadInterruptState.bind(connection), + } + : {}), } } @@ -565,6 +592,23 @@ export function normalizeConnectionAdapter( throw err } }, + ...(connection.joinRun !== undefined + ? { + async joinRun(runId: string, abortSignal?: AbortSignal) { + for await (const chunk of connection.joinRun?.( + runId, + abortSignal, + ) ?? []) { + push(chunk, runId) + } + }, + } + : {}), + ...(connection.loadInterruptState !== undefined + ? { + loadInterruptState: connection.loadInterruptState.bind(connection), + } + : {}), } } @@ -577,6 +621,195 @@ export interface FetchConnectionOptions { signal?: AbortSignal body?: Record fetchClient?: typeof globalThis.fetch + /** Explicit authoritative interrupt recovery operation. */ + interruptStateFetcher?: NonNullable< + ConnectConnectionAdapter['loadInterruptState'] + > + /** Explicit read-only loader for a committed continuation run. */ + continuationLoader?: NonNullable +} + +export type InterruptStateFetcher = NonNullable< + ConnectConnectionAdapter['loadInterruptState'] +> + +export type InterruptContinuationLoader = NonNullable< + ConnectConnectionAdapter['joinRun'] +> + +export interface InterruptFetchOptions { + headers?: Record | Headers + credentials?: RequestCredentials + signal?: AbortSignal + fetchClient?: typeof globalThis.fetch +} + +function isUnknownRecord(value: unknown): value is Record { + return value !== null && typeof value === 'object' && !Array.isArray(value) +} + +function isInterruptDescriptor(value: unknown): value is Interrupt { + if (!isUnknownRecord(value)) return false + if (typeof value.id !== 'string' || typeof value.reason !== 'string') { + return false + } + if (value.message !== undefined && typeof value.message !== 'string') { + return false + } + if (value.toolCallId !== undefined && typeof value.toolCallId !== 'string') { + return false + } + if ( + value.responseSchema !== undefined && + !isUnknownRecord(value.responseSchema) + ) { + return false + } + if (value.expiresAt !== undefined && typeof value.expiresAt !== 'string') { + return false + } + return value.metadata === undefined || isUnknownRecord(value.metadata) +} + +function isResumeEntry(value: unknown): value is RunAgentResumeItem { + return ( + isUnknownRecord(value) && + typeof value.interruptId === 'string' && + (value.status === 'resolved' || value.status === 'cancelled') + ) +} + +function isInterruptRecoveryStateName( + value: unknown, +): value is InterruptRecoveryStateV1['state'] { + return ( + typeof value === 'string' && + ['pending', 'committed', 'expired', 'missing', 'legacy-committed'].includes( + value, + ) + ) +} + +/** @internal Runtime boundary parser shared by persistence and fetch recovery. */ +export function parseInterruptRecoveryState( + value: unknown, +): InterruptRecoveryStateV1 { + if ( + !isUnknownRecord(value) || + value.schemaVersion !== 1 || + !isInterruptRecoveryStateName(value.state) || + typeof value.threadId !== 'string' || + value.threadId.length === 0 || + typeof value.interruptedRunId !== 'string' || + value.interruptedRunId.length === 0 || + typeof value.generation !== 'number' || + !Number.isSafeInteger(value.generation) || + value.generation < 0 || + !Array.isArray(value.pendingInterrupts) || + !value.pendingInterrupts.every(isInterruptDescriptor) + ) { + throw new TypeError('Invalid interrupt recovery response.') + } + let committed: InterruptRecoveryStateV1['committed'] + if (value.committed !== undefined) { + if ( + !isUnknownRecord(value.committed) || + typeof value.committed.fingerprint !== 'string' || + typeof value.committed.committedAt !== 'string' || + (value.committed.continuationRunId !== undefined && + typeof value.committed.continuationRunId !== 'string') || + (value.committed.resolutions !== undefined && + (!Array.isArray(value.committed.resolutions) || + !value.committed.resolutions.every(isResumeEntry))) + ) { + throw new TypeError('Invalid interrupt recovery response.') + } + committed = { + fingerprint: value.committed.fingerprint, + committedAt: value.committed.committedAt, + ...(value.committed.continuationRunId === undefined + ? {} + : { continuationRunId: value.committed.continuationRunId }), + ...(value.committed.resolutions === undefined + ? {} + : { resolutions: value.committed.resolutions }), + } + } + return { + schemaVersion: 1, + state: value.state, + threadId: value.threadId, + interruptedRunId: value.interruptedRunId, + generation: value.generation, + pendingInterrupts: value.pendingInterrupts, + ...(typeof value.submissionId === 'string' + ? { submissionId: value.submissionId } + : {}), + ...(typeof value.continuationRunId === 'string' + ? { continuationRunId: value.continuationRunId } + : {}), + ...(committed === undefined ? {} : { committed }), + } +} + +/** Create an authoritative interrupt-state fetcher for an explicit endpoint. */ +export function createInterruptStateFetcher( + url: string | (() => string), + options: InterruptFetchOptions = {}, +): InterruptStateFetcher { + return async (query, abortSignal) => { + const resolvedUrl = typeof url === 'function' ? url() : url + const requestUrl = withSearchParams(resolvedUrl, { + threadId: query.threadId, + interruptedRunId: query.interruptedRunId, + knownGeneration: String(query.knownGeneration), + }) + const response = await (options.fetchClient ?? fetch)(requestUrl, { + method: 'POST', + headers: mergeHeaders(options.headers), + credentials: options.credentials ?? 'same-origin', + ...((abortSignal ?? options.signal) === undefined + ? {} + : { signal: abortSignal ?? options.signal }), + }) + if (!response.ok) { + throw new Error( + `Interrupt recovery request failed with status ${response.status}.`, + ) + } + const state = parseInterruptRecoveryState(await response.json()) + if ( + state.threadId !== query.threadId || + state.interruptedRunId !== query.interruptedRunId + ) { + throw new TypeError('Invalid interrupt recovery response correlation.') + } + return state + } +} + +/** Create a read-only continuation loader for an explicit endpoint. */ +export function createInterruptContinuationLoader( + url: string | (() => string), + options: InterruptFetchOptions = {}, +): InterruptContinuationLoader { + return async function* (runId, abortSignal) { + const resolvedUrl = typeof url === 'function' ? url() : url + const requestUrl = withSearchParams(resolvedUrl, { + offset: '-1', + runId, + }) + yield* resumableServerSentEvents( + options.fetchClient ?? fetch, + requestUrl, + { + method: 'GET', + headers: mergeHeaders(options.headers), + credentials: options.credentials ?? 'same-origin', + }, + abortSignal ?? options.signal, + ) + } } /** @@ -726,6 +959,14 @@ export function fetchServerSentEvents( const resolvedOptions = typeof options === 'function' ? await options() : options + if (resolvedOptions.continuationLoader !== undefined) { + yield* resolvedOptions.continuationLoader( + runId, + abortSignal ?? resolvedOptions.signal, + ) + return + } + const joinUrl = withSearchParams(resolvedUrl, { offset: '-1', runId, @@ -748,6 +989,19 @@ export function fetchServerSentEvents( signal, ) }, + async loadInterruptState(query, abortSignal) { + const resolvedOptions = + typeof options === 'function' ? await options() : options + if (resolvedOptions.interruptStateFetcher === undefined) { + throw new Error( + 'Interrupt recovery requires an explicit interruptStateFetcher.', + ) + } + return resolvedOptions.interruptStateFetcher( + query, + abortSignal ?? resolvedOptions.signal, + ) + }, } } @@ -1197,6 +1451,9 @@ export function fetcherToConnectionAdapter( data, threadId: runContext.threadId, runId: runContext.runId, + ...(runContext.parentRunId !== undefined && { + parentRunId: runContext.parentRunId, + }), ...(runContext.resume !== undefined && { resume: runContext.resume }), }, { signal: abortSignal }, diff --git a/packages/ai-client/src/index.ts b/packages/ai-client/src/index.ts index bb8b7f7ae..d2dec69d6 100644 --- a/packages/ai-client/src/index.ts +++ b/packages/ai-client/src/index.ts @@ -6,6 +6,12 @@ export type { InferAudioRecordingOutput, } from './audio-recorder' export { ChatClient } from './chat-client' +export { InterruptManager } from './interrupt-manager' +export type { + InterruptManagerHydration, + InterruptManagerOptions, + InterruptManagerSubmission, +} from './interrupt-manager' export { createMcpAppBridge } from './mcp-app-bridge' export type { McpAppBridge, CreateMcpAppBridgeOptions } from './mcp-app-bridge' export { RealtimeClient } from './realtime-client' @@ -27,9 +33,21 @@ export type { ChatPersistenceOptions, ChatClientOptions, ChatPendingInterrupt, + BoundInterruptBase, + BoundInterrupts, + ChatInterrupt, + ChatInterruptState, + ClientToolExecutionInterrupt, + GenericAGUIInterrupt, + InterruptItemStatus, + ToolApprovalInterrupt, ClientContextOptionFromTools, ChatResumeState, ChatResumeSnapshot, + ChatResumeSnapshotV1, + ChatResumeSnapshotV2, + ChatContinuationLoader, + PersistedInterruptDraft, ChatRequestBody, InferChatMessages, InferredClientContext, @@ -120,6 +138,8 @@ export type { RealtimeStateChangeCallback, } from './realtime-types' export { + createInterruptContinuationLoader, + createInterruptStateFetcher, fetchServerSentEvents, fetchHttpStream, xhrServerSentEvents, @@ -131,6 +151,9 @@ export { type ConnectConnectionAdapter, type ConnectionAdapter, type FetchConnectionOptions, + type InterruptContinuationLoader, + type InterruptFetchOptions, + type InterruptStateFetcher, type ResumableConnectConnectionAdapter, type RunAgentInputContext, type SubscribeConnectionAdapter, diff --git a/packages/ai-client/src/interrupt-manager.ts b/packages/ai-client/src/interrupt-manager.ts new file mode 100644 index 000000000..b0ad355da --- /dev/null +++ b/packages/ai-client/src/interrupt-manager.ts @@ -0,0 +1,1396 @@ +import { + canonicalInterruptJson, + canonicalizeInterruptResolutions, + cloneAndDeepFreezeJson, + compileJsonSchema202012, + digestInterruptJson, + hashSchemaInput, + isStandardSchema, + normalizeApprovalSchema, +} from '@tanstack/ai/client' +import type { + AnyClientTool, + BatchInterruptError, + Interrupt, + InterruptBinding, + InterruptCommitResult, + InterruptRecoveryStateV1, + InterruptSubmissionError, + ItemInterruptError, + RunAgentResumeItem, +} from '@tanstack/ai/client' +import type { + BoundInterruptBase, + BoundInterrupts, + ChatInterrupt, + ChatInterruptState, + GenericAGUIInterrupt, + InterruptItemStatus, + PersistedInterruptDraft, +} from './types' + +export interface InterruptManagerHydration { + threadId: string + interruptedRunId: string + generation: number + interrupts: ReadonlyArray +} + +export interface InterruptManagerSubmission { + threadId: string + interruptedRunId: string + generation: number + resolutions: ReadonlyArray + canonicalResolutions: string + fingerprint: string +} + +export interface InterruptManagerOptions< + TTools extends ReadonlyArray, +> { + tools?: TTools + submit: ( + submission: InterruptManagerSubmission, + ) => Promise + recover?: (state?: InterruptRecoveryStateV1) => void | Promise + onChange?: () => void +} + +type UnknownObject = { [key: string]: unknown } + +type RuntimeKind = 'generic' | 'tool-approval' | 'client-tool-execution' + +interface RuntimeInterrupt { + descriptor: Interrupt + binding: InterruptBinding + kind: RuntimeKind + status: InterruptItemStatus + canResolve: boolean + error?: ItemInterruptError + resolution?: RunAgentResumeItem + tool?: AnyClientTool + validationGeneration: number +} + +interface ValidationFailure { + code: ItemInterruptError['code'] + message: string + path?: ReadonlyArray +} + +type ValidationResult = { valid: true; payload: unknown } | ValidationFailure + +interface TransactionToken { + active: boolean +} + +interface RuntimeInterruptCheckpoint { + status: InterruptItemStatus + resolution?: RunAgentResumeItem + error?: ItemInterruptError + validationGeneration: number +} + +const itemErrorCodes = new Set([ + 'invalid-payload', + 'invalid-edited-args', + 'invalid-tool-output', + 'invalid-response-schema', + 'unknown-interrupt', + 'expired', + 'stale', + 'conflict', + 'legacy-unsupported', +]) + +const batchErrorCodes = new Set([ + 'incomplete-batch', + 'item-validation-failed', + 'unsupported-bulk-operation', + 'async-resolver', + 'inactive-transaction', + 'mixed-provenance', + 'transport', + 'server', + 'protocol', + 'invalid-response-schema', + 'expired', + 'stale', + 'conflict', + 'persistence-required', + 'atomic-commit-unsupported', + 'recovery-unavailable', + 'legacy-submit-failed', +]) + +function isUnknownObject(value: unknown): value is UnknownObject { + return value !== null && typeof value === 'object' && !Array.isArray(value) +} + +function isLegacyApprovalMetadata(value: unknown): boolean { + return ( + isUnknownObject(value) && + value['kind'] === 'approval' && + typeof value['toolName'] === 'string' && + 'input' in value + ) +} + +function isLegacyClientToolMetadata(value: unknown): boolean { + return ( + isUnknownObject(value) && + value['kind'] === 'client_tool' && + typeof value['toolName'] === 'string' && + 'input' in value + ) +} + +function isBindingBase(value: UnknownObject): boolean { + return ( + typeof value['kind'] === 'string' && + typeof value['interruptId'] === 'string' && + typeof value['interruptedRunId'] === 'string' && + typeof value['generation'] === 'number' && + Number.isInteger(value['generation']) && + value['generation'] >= 0 && + typeof value['responseSchemaHash'] === 'string' && + (value['expiresAt'] === undefined || + (typeof value['expiresAt'] === 'string' && + Number.isFinite(Date.parse(value['expiresAt'])))) + ) +} + +function readBinding(value: unknown): InterruptBinding | undefined { + if (!isUnknownObject(value) || !isBindingBase(value)) return undefined + const expiresAt = + typeof value['expiresAt'] === 'string' ? value['expiresAt'] : undefined + if (value['kind'] === 'generic') { + return { + kind: 'generic', + interruptId: String(value['interruptId']), + interruptedRunId: String(value['interruptedRunId']), + generation: Number(value['generation']), + responseSchemaHash: String(value['responseSchemaHash']), + ...(expiresAt !== undefined ? { expiresAt } : {}), + } + } + if ( + value['kind'] === 'client-tool-execution' && + typeof value['toolName'] === 'string' && + typeof value['toolCallId'] === 'string' && + typeof value['outputSchemaHash'] === 'string' + ) { + return { + kind: 'client-tool-execution', + interruptId: String(value['interruptId']), + interruptedRunId: String(value['interruptedRunId']), + generation: Number(value['generation']), + toolName: value['toolName'], + toolCallId: value['toolCallId'], + outputSchemaHash: value['outputSchemaHash'], + responseSchemaHash: String(value['responseSchemaHash']), + ...(expiresAt !== undefined ? { expiresAt } : {}), + } + } + if ( + value['kind'] === 'tool-approval' && + typeof value['toolName'] === 'string' && + typeof value['toolCallId'] === 'string' && + typeof value['inputSchemaHash'] === 'string' && + typeof value['approvalSchemaHash'] === 'string' && + 'originalArgs' in value + ) { + return { + kind: 'tool-approval', + interruptId: String(value['interruptId']), + interruptedRunId: String(value['interruptedRunId']), + generation: Number(value['generation']), + toolName: value['toolName'], + toolCallId: value['toolCallId'], + originalArgs: value['originalArgs'], + inputSchemaHash: value['inputSchemaHash'], + approvalSchemaHash: value['approvalSchemaHash'], + responseSchemaHash: String(value['responseSchemaHash']), + ...(expiresAt !== undefined ? { expiresAt } : {}), + } + } + return undefined +} + +function getDescriptorBinding( + interrupt: Interrupt, +): InterruptBinding | undefined { + const candidate: unknown = interrupt.metadata?.['tanstack:interruptBinding'] + return readBinding(candidate) +} + +function isToolApprovalReason(reason: string): boolean { + return reason === 'tool_call' || reason === 'approval_required' +} + +function isClientToolExecutionReason(reason: string): boolean { + return ( + reason === 'tanstack:client_tool_execution' || + reason === 'client_tool_input' + ) +} + +function responseSchemaHash(interrupt: Interrupt): string | undefined { + if (interrupt.responseSchema === undefined) return undefined + try { + return digestInterruptJson(canonicalInterruptJson(interrupt.responseSchema)) + } catch { + return undefined + } +} + +function canValidateGeneric(interrupt: Interrupt): boolean { + if (interrupt.responseSchema === undefined) return true + try { + compileJsonSchema202012(interrupt.responseSchema) + return true + } catch { + return false + } +} + +function isPromiseLike(value: unknown): value is PromiseLike { + return ( + value !== null && + (typeof value === 'object' || typeof value === 'function') && + 'then' in value && + typeof value.then === 'function' + ) +} + +function validateWithSchema( + schema: unknown, + value: unknown, + code: ItemInterruptError['code'], +): ValidationResult | Promise { + if (schema === undefined) return { valid: true, payload: value } + if (isStandardSchema(schema)) { + const result = schema['~standard'].validate(value) + const normalize = ( + validation: Awaited, + ): ValidationResult => { + if (!validation.issues) { + return { valid: true, payload: validation.value } + } + return { + code, + message: validation.issues[0]?.message ?? 'Schema validation failed.', + } + } + return isPromiseLike(result) + ? Promise.resolve(result).then(normalize) + : normalize(result) + } + try { + const issues = compileJsonSchema202012(schema)(value) + const issue = issues[0] + return issue + ? { code, message: issue.message, path: issue.path } + : { valid: true, payload: value } + } catch { + return { + code: 'invalid-response-schema', + message: 'The interrupt response schema is invalid.', + } + } +} + +function isItemErrorCode(value: string): value is ItemInterruptError['code'] { + for (const code of itemErrorCodes) if (code === value) return true + return false +} + +function isBatchErrorCode(value: string): value is BatchInterruptError['code'] { + for (const code of batchErrorCodes) if (code === value) return true + return false +} + +function isSubmissionError(value: unknown): value is InterruptSubmissionError { + if (!isUnknownObject(value)) return false + const scope = value['scope'] + const code = value['code'] + const base = + typeof code === 'string' && + typeof value['message'] === 'string' && + typeof value['retryable'] === 'boolean' && + typeof value['threadId'] === 'string' && + typeof value['interruptedRunId'] === 'string' && + typeof value['generation'] === 'number' + if (!base) return false + if (scope === 'item') { + return ( + isItemErrorCode(code) && + typeof value['interruptId'] === 'string' && + (value['source'] === 'client' || value['source'] === 'server') + ) + } + return ( + scope === 'batch' && + isBatchErrorCode(code) && + Array.isArray(value['interruptIds']) && + value['interruptIds'].every((id) => typeof id === 'string') && + (value['source'] === 'client' || + value['source'] === 'server' || + value['source'] === 'transport') + ) +} + +function readSubmissionErrors( + error: unknown, +): ReadonlyArray { + if (isSubmissionError(error)) return [error] + if (!isUnknownObject(error) || !Array.isArray(error['errors'])) return [] + return error['errors'].every(isSubmissionError) ? error['errors'] : [] +} + +function genericBinding( + interrupt: Interrupt, + hydration: InterruptManagerHydration, + candidate: InterruptBinding | undefined, +): InterruptBinding { + return cloneAndDeepFreezeJson({ + kind: 'generic', + interruptId: interrupt.id, + interruptedRunId: hydration.interruptedRunId, + generation: hydration.generation, + responseSchemaHash: + responseSchemaHash(interrupt) ?? + candidate?.responseSchemaHash ?? + 'invalid', + ...(interrupt.expiresAt !== undefined + ? { expiresAt: interrupt.expiresAt } + : {}), + }) +} + +function baseSnapshot( + item: RuntimeInterrupt, + hydration: InterruptManagerHydration, + cancel: () => void, + clearResolution: () => void, +): BoundInterruptBase { + const descriptor = cloneAndDeepFreezeJson(item.descriptor) + const binding = cloneAndDeepFreezeJson(item.binding) + const errors: ReadonlyArray = + item.error === undefined + ? Object.freeze([]) + : Object.freeze([cloneAndDeepFreezeJson(item.error)]) + const error = errors[0] + return { + id: descriptor.id, + interruptId: descriptor.id, + reason: descriptor.reason, + ...(descriptor.message !== undefined + ? { message: descriptor.message } + : {}), + ...(descriptor.responseSchema !== undefined + ? { responseSchema: descriptor.responseSchema } + : {}), + ...(descriptor.expiresAt !== undefined + ? { expiresAt: descriptor.expiresAt } + : {}), + ...(descriptor.metadata !== undefined + ? { metadata: descriptor.metadata } + : {}), + threadId: hydration.threadId, + interruptedRunId: hydration.interruptedRunId, + generation: hydration.generation, + status: item.status, + binding, + errors, + ...(error !== undefined ? { error } : {}), + canResolve: item.canResolve, + cancel, + clearResolution, + } +} + +export class InterruptManager< + TTools extends ReadonlyArray = ReadonlyArray, +> { + private hydration: InterruptManagerHydration | undefined + private items: Array = [] + private snapshot: ReadonlyArray> = Object.freeze([]) + private rootErrors: ReadonlyArray = Object.freeze([]) + private state: ChatInterruptState = Object.freeze({ + interrupts: this.snapshot, + pendingInterrupts: this.snapshot, + interruptErrors: this.rootErrors, + resuming: false, + }) + private activeTransaction: TransactionToken | undefined + private retrySubmission: InterruptManagerSubmission | undefined + private resuming = false + private tools: TTools | undefined + + constructor(private readonly options: InterruptManagerOptions) { + this.tools = options.tools + } + + updateTools(tools: TTools): void { + this.tools = tools + } + + hydrate(hydration: InterruptManagerHydration): void { + this.hydration = { + threadId: hydration.threadId, + interruptedRunId: hydration.interruptedRunId, + generation: hydration.generation, + interrupts: cloneAndDeepFreezeJson(hydration.interrupts), + } + this.items = hydration.interrupts.map((interrupt) => + this.hydrateInterrupt(interrupt, hydration), + ) + this.rootErrors = Object.freeze([]) + this.retrySubmission = undefined + this.resuming = false + this.publish() + } + + getInterrupts(): BoundInterrupts { + return this.snapshot + } + + getState(): ChatInterruptState { + return this.state + } + + getDescriptors(): ReadonlyArray { + return this.hydration?.interrupts ?? Object.freeze([]) + } + + getRecoveryState(): InterruptRecoveryStateV1 | undefined { + const hydration = this.hydration + if (!hydration) return undefined + return cloneAndDeepFreezeJson({ + schemaVersion: 1, + state: 'pending', + threadId: hydration.threadId, + interruptedRunId: hydration.interruptedRunId, + generation: hydration.generation, + pendingInterrupts: hydration.interrupts, + }) + } + + getPersistedDrafts(): ReadonlyArray { + return cloneAndDeepFreezeJson( + this.items.flatMap((item) => { + if (item.resolution === undefined && item.error === undefined) return [] + return [ + { + interruptId: item.descriptor.id, + response: item.resolution ?? null, + status: item.status, + ...(item.error !== undefined ? { error: item.error } : {}), + }, + ] + }), + ) + } + + reportRecoveryError( + code: 'expired' | 'stale' | 'conflict' | 'recovery-unavailable', + message: string, + recoveryState?: InterruptRecoveryStateV1, + ): void { + if (!this.hydration && recoveryState !== undefined) { + this.hydration = { + threadId: recoveryState.threadId, + interruptedRunId: recoveryState.interruptedRunId, + generation: recoveryState.generation, + interrupts: cloneAndDeepFreezeJson(recoveryState.pendingInterrupts), + } + this.items = recoveryState.pendingInterrupts.map((interrupt) => + this.hydrateInterrupt(interrupt, this.requireHydration()), + ) + } + for (const item of this.items) { + if (item.status === 'submitting') item.status = 'error' + } + this.addRootError(code, message, false, 'server') + this.publish() + } + + restorePersistedDrafts(drafts: ReadonlyArray): void { + for (const draft of drafts) { + const item = this.items.find( + (candidate) => candidate.descriptor.id === draft.interruptId, + ) + if (!item) continue + const response = draft.response + if ( + isUnknownObject(response) && + response['interruptId'] === item.descriptor.id && + response['status'] === 'cancelled' + ) { + item.resolution = Object.freeze({ + interruptId: item.descriptor.id, + status: 'cancelled', + }) + item.status = draft.error === undefined ? 'staged' : 'error' + item.error = draft.error + continue + } + if ( + isUnknownObject(response) && + response['interruptId'] === item.descriptor.id && + response['status'] === 'resolved' + ) { + const validationGeneration = ++item.validationGeneration + const validation = this.validateCandidate(item, response['payload']) + if (isPromiseLike(validation)) { + item.status = 'validating' + void Promise.resolve(validation) + .then((result) => { + if (validationGeneration !== item.validationGeneration) return + this.applyRestoredValidation(item, result, draft) + }) + .catch((error: unknown) => { + if (validationGeneration !== item.validationGeneration) return + this.applyRestoredValidation( + item, + { + code: this.validationCode(item), + message: + error instanceof Error ? error.message : String(error), + }, + draft, + ) + }) + } else { + this.applyRestoredValidation(item, validation, draft, false) + } + continue + } + if (draft.error !== undefined) { + item.status = 'error' + item.error = cloneAndDeepFreezeJson(draft.error) + } + } + this.publish() + } + + reset(options?: { preserveRootErrors?: boolean }): void { + this.hydration = undefined + this.items = [] + this.snapshot = Object.freeze([]) + if (options?.preserveRootErrors !== true) { + this.rootErrors = Object.freeze([]) + } + this.retrySubmission = undefined + this.resuming = false + this.state = Object.freeze({ + interrupts: this.snapshot, + pendingInterrupts: this.snapshot, + interruptErrors: this.rootErrors, + resuming: false, + }) + this.options.onChange?.() + } + + getInterruptErrors(): ReadonlyArray { + return this.rootErrors + } + + getResuming(): boolean { + return this.resuming + } + + resolve(approved: boolean): void + resolve(resolver: (interrupt: ChatInterrupt) => undefined): void + resolve( + resolution: boolean | ((interrupt: ChatInterrupt) => unknown), + ): void { + this.assertRootMutable() + if (typeof resolution === 'boolean') { + this.resolveBooleanBulk(resolution) + return + } + this.resolveTransaction(resolution) + } + + cancel(): void { + this.assertRootMutable() + this.invalidateRetry() + for (const item of this.items) { + item.validationGeneration++ + item.resolution = Object.freeze({ + interruptId: item.descriptor.id, + status: 'cancelled', + }) + item.status = 'staged' + item.error = undefined + } + this.publish() + this.maybeSubmit() + } + + retry(): void { + if (this.resuming) + throw new Error('Interrupt submission is already active.') + const submission = this.retrySubmission + if (!submission) { + this.addRootError( + 'transport', + 'There is no retryable interrupt submission.', + false, + ) + return + } + this.submitBatch(submission) + } + + resolveClientToolOutput(toolCallId: string, output: unknown): boolean { + const item = this.items.find( + (candidate) => + (candidate.kind === 'client-tool-execution' && + candidate.binding.kind === 'client-tool-execution' && + candidate.binding.toolCallId === toolCallId) || + (candidate.kind === 'generic' && + candidate.descriptor.reason === 'client_tool_input' && + candidate.descriptor.toolCallId === toolCallId && + isLegacyClientToolMetadata(candidate.descriptor.metadata)), + ) + if (!item) return false + this.resolveItem(item.descriptor.id, output) + return true + } + + resolveToolApprovalDecision(interruptId: string, approved: boolean): boolean { + const item = this.items.find( + (candidate) => + candidate.descriptor.id === interruptId && + (candidate.kind === 'tool-approval' || + (candidate.kind === 'generic' && + candidate.descriptor.reason === 'approval_required' && + isLegacyApprovalMetadata(candidate.descriptor.metadata))), + ) + if (!item) return false + this.resolveItem(item.descriptor.id, { approved }) + return true + } + + private hydrateInterrupt( + descriptor: Interrupt, + hydration: InterruptManagerHydration, + ): RuntimeInterrupt { + const interrupt = cloneAndDeepFreezeJson(descriptor) + const candidate = getDescriptorBinding(interrupt) + const correlated = + candidate !== undefined && + candidate.interruptId === interrupt.id && + candidate.interruptedRunId === hydration.interruptedRunId && + candidate.generation === hydration.generation && + candidate.responseSchemaHash === + (responseSchemaHash(interrupt) ?? candidate.responseSchemaHash) + + if (correlated && candidate.kind === 'tool-approval') { + const tool = this.tools?.find( + (configured) => configured.name === candidate.toolName, + ) + if ( + tool?.needsApproval === true && + isToolApprovalReason(interrupt.reason) && + interrupt.toolCallId === candidate.toolCallId + ) { + try { + const approval = normalizeApprovalSchema( + tool.approvalSchema, + tool.inputSchema, + ) + if ( + hashSchemaInput(tool.inputSchema) === candidate.inputSchemaHash && + approval.approvalSchemaHash === candidate.approvalSchemaHash && + approval.responseSchemaHash === candidate.responseSchemaHash + ) { + return { + descriptor: interrupt, + binding: cloneAndDeepFreezeJson(candidate), + kind: 'tool-approval', + status: 'pending', + canResolve: true, + tool, + validationGeneration: 0, + } + } + } catch { + // Invalid configured schemas cannot safely grant typed hydration. + } + } + } + + if (correlated && candidate.kind === 'client-tool-execution') { + const tool = this.tools?.find( + (configured) => configured.name === candidate.toolName, + ) + if ( + tool !== undefined && + isClientToolExecutionReason(interrupt.reason) && + interrupt.toolCallId === candidate.toolCallId && + hashSchemaInput(tool.outputSchema) === candidate.outputSchemaHash + ) { + return { + descriptor: interrupt, + binding: cloneAndDeepFreezeJson(candidate), + kind: 'client-tool-execution', + status: 'pending', + canResolve: true, + tool, + validationGeneration: 0, + } + } + } + + return { + descriptor: interrupt, + binding: genericBinding(interrupt, hydration, candidate), + kind: 'generic', + status: 'pending', + canResolve: canValidateGeneric(interrupt), + validationGeneration: 0, + } + } + + private buildSnapshot( + transaction?: TransactionToken, + ): BoundInterrupts { + const hydration = this.requireHydration() + const next = this.items.map((item) => { + const base = baseSnapshot( + item, + hydration, + () => this.cancelItem(item.descriptor.id, transaction), + () => this.clearItem(item.descriptor.id, transaction), + ) + if ( + item.kind === 'tool-approval' && + item.binding.kind === 'tool-approval' + ) { + const snapshot = { + ...base, + kind: 'tool-approval' as const, + toolName: item.binding.toolName, + toolCallId: item.binding.toolCallId, + originalArgs: cloneAndDeepFreezeJson(item.binding.originalArgs), + resolveInterrupt: (approved: boolean, options?: unknown) => { + const details = isUnknownObject(options) ? options : undefined + this.resolveItem( + item.descriptor.id, + { + approved, + ...(approved && details?.['editedArgs'] !== undefined + ? { editedArgs: details['editedArgs'] } + : {}), + ...(details?.['payload'] !== undefined + ? { payload: details['payload'] } + : {}), + }, + transaction, + ) + }, + } + return Object.freeze(snapshot) + } + if ( + item.kind === 'client-tool-execution' && + item.binding.kind === 'client-tool-execution' + ) { + const snapshot = { + ...base, + kind: 'client-tool-execution' as const, + toolName: item.binding.toolName, + toolCallId: item.binding.toolCallId, + resolveInterrupt: (output: unknown) => + this.resolveItem(item.descriptor.id, output, transaction), + } + return Object.freeze(snapshot) + } + const snapshot: GenericAGUIInterrupt = { + ...base, + kind: 'generic', + resolveInterrupt: (payload) => + this.resolveItem(item.descriptor.id, payload, transaction), + } + return Object.freeze(snapshot) + }) + + // The runtime items are created only from the exact configured TTools entry + // selected by name. TypeScript cannot preserve that per-element lookup + // through Array.map, so this generic return boundary restores the proven + // distributive public union. + return Object.freeze(next) as BoundInterrupts + } + + private publish(): void { + if (!this.hydration) { + this.snapshot = Object.freeze([]) + this.state = Object.freeze({ + interrupts: this.snapshot, + pendingInterrupts: this.snapshot, + interruptErrors: this.rootErrors, + resuming: this.resuming, + }) + this.options.onChange?.() + return + } + this.snapshot = this.buildSnapshot() + this.state = Object.freeze({ + interrupts: this.snapshot, + pendingInterrupts: this.snapshot, + interruptErrors: this.rootErrors, + resuming: this.resuming, + }) + this.options.onChange?.() + } + + private resolveItem( + interruptId: string, + payload: unknown, + transaction?: TransactionToken, + ): void { + this.assertItemMutable(transaction) + const item = this.findItem(interruptId) + this.invalidateRetry() + if (!item.canResolve) { + item.status = 'error' + item.error = this.itemError( + interruptId, + 'invalid-response-schema', + 'The interrupt response schema is invalid and cannot be resolved.', + ) + if (!transaction) this.publish() + return + } + const validationGeneration = ++item.validationGeneration + const validation = this.validateCandidate(item, payload) + if (isPromiseLike(validation)) { + item.status = 'validating' + item.error = undefined + if (!transaction) this.publish() + void Promise.resolve(validation) + .then((result) => { + if (validationGeneration !== item.validationGeneration) return + this.applyValidation(item, result, transaction) + }) + .catch((error: unknown) => { + if (validationGeneration !== item.validationGeneration) return + this.applyValidation( + item, + { + code: this.validationCode(item), + message: error instanceof Error ? error.message : String(error), + }, + transaction, + ) + }) + return + } + this.applyValidation(item, validation, transaction) + } + + private cancelItem( + interruptId: string, + transaction?: TransactionToken, + ): void { + this.assertItemMutable(transaction) + const item = this.findItem(interruptId) + this.invalidateRetry() + item.validationGeneration++ + item.resolution = Object.freeze({ interruptId, status: 'cancelled' }) + item.status = 'staged' + item.error = undefined + if (!transaction) { + this.publish() + this.maybeSubmit() + } + } + + private clearItem(interruptId: string, transaction?: TransactionToken): void { + this.assertItemMutable(transaction) + const item = this.findItem(interruptId) + this.invalidateRetry() + item.validationGeneration++ + item.resolution = undefined + item.error = undefined + item.status = 'pending' + if (!transaction) this.publish() + } + + private maybeSubmit(): void { + if ( + this.items.length === 0 || + this.items.some( + (item) => item.resolution === undefined || item.status !== 'staged', + ) + ) { + return + } + const hydration = this.requireHydration() + const canonical = canonicalizeInterruptResolutions( + this.items + .map((item) => item.resolution) + .filter((item) => item !== undefined), + ) + const submission = Object.freeze({ + threadId: hydration.threadId, + interruptedRunId: hydration.interruptedRunId, + generation: hydration.generation, + resolutions: canonical.resolutions, + canonicalResolutions: canonical.canonicalResolutions, + fingerprint: canonical.fingerprint, + }) + this.submitBatch(submission) + } + + private applyValidation( + item: RuntimeInterrupt, + result: ValidationResult, + transaction?: TransactionToken, + ): void { + if (!('valid' in result)) { + item.status = 'error' + item.error = this.itemError( + item.descriptor.id, + result.code, + result.message, + result.path, + ) + if (!transaction) this.publish() + return + } + item.resolution = cloneAndDeepFreezeJson({ + interruptId: item.descriptor.id, + status: 'resolved', + payload: result.payload, + }) + item.status = 'staged' + item.error = undefined + if (!transaction) { + this.publish() + this.maybeSubmit() + } + } + + private applyRestoredValidation( + item: RuntimeInterrupt, + result: ValidationResult, + draft: PersistedInterruptDraft, + publish = true, + ): void { + if (!('valid' in result)) { + item.status = 'error' + item.error = this.itemError( + item.descriptor.id, + result.code, + result.message, + result.path, + ) + } else { + item.resolution = cloneAndDeepFreezeJson({ + interruptId: item.descriptor.id, + status: 'resolved', + payload: result.payload, + }) + item.status = draft.error === undefined ? 'staged' : 'error' + item.error = + draft.error === undefined + ? undefined + : cloneAndDeepFreezeJson(draft.error) + } + if (publish) this.publish() + } + + private validateCandidate( + item: RuntimeInterrupt, + payload: unknown, + ): ValidationResult | Promise { + if (item.kind === 'generic') { + return validateWithSchema( + item.descriptor.responseSchema, + payload, + 'invalid-payload', + ) + } + if (item.kind === 'client-tool-execution') { + return validateWithSchema( + item.tool?.outputSchema, + payload, + 'invalid-tool-output', + ) + } + return this.validateApprovalCandidate(item, payload) + } + + private validateApprovalCandidate( + item: RuntimeInterrupt, + payload: unknown, + ): ValidationResult | Promise { + if (!isUnknownObject(payload) || typeof payload['approved'] !== 'boolean') { + return { + code: 'invalid-payload', + message: 'Tool approval resolutions require an approved boolean.', + } + } + const approved = payload['approved'] + const editedArgs = payload['editedArgs'] + if (!approved && editedArgs !== undefined) { + return { + code: 'invalid-edited-args', + message: 'Rejected tool approvals cannot edit tool arguments.', + } + } + if (approved && editedArgs !== undefined) { + const editedValidation = validateWithSchema( + item.tool?.inputSchema, + editedArgs, + 'invalid-edited-args', + ) + if (isPromiseLike(editedValidation)) { + return Promise.resolve(editedValidation).then((result) => + 'valid' in result + ? this.validateApprovalPayload(item, payload, result.payload) + : result, + ) + } + if (!('valid' in editedValidation)) return editedValidation + return this.validateApprovalPayload( + item, + payload, + editedValidation.payload, + ) + } + return this.validateApprovalPayload(item, payload, undefined) + } + + private validateApprovalPayload( + item: RuntimeInterrupt, + envelope: UnknownObject, + validatedEditedArgs: unknown, + ): ValidationResult | Promise { + const approved = envelope['approved'] === true + const schema = this.approvalBranchSchema(item.tool, approved) + const branchPayload = envelope['payload'] + if (schema === undefined && branchPayload !== undefined) { + return { + code: 'invalid-payload', + message: 'This approval branch does not accept a payload.', + } + } + if (schema !== undefined && branchPayload === undefined) { + return { + code: 'invalid-payload', + message: 'This approval branch requires a payload.', + } + } + const validation = validateWithSchema( + schema, + branchPayload, + 'invalid-payload', + ) + const buildEnvelope = (result: ValidationResult): ValidationResult => { + if (!('valid' in result)) return result + return { + valid: true, + payload: { + approved, + ...(validatedEditedArgs !== undefined + ? { editedArgs: validatedEditedArgs } + : {}), + ...(schema !== undefined ? { payload: result.payload } : {}), + }, + } + } + return isPromiseLike(validation) + ? Promise.resolve(validation).then(buildEnvelope) + : buildEnvelope(validation) + } + + private approvalBranchSchema( + tool: AnyClientTool | undefined, + approved: boolean, + ): unknown { + const approvalSchema: unknown = tool?.approvalSchema + if (!isUnknownObject(approvalSchema)) return approvalSchema + const hasBranches = + 'approve' in approvalSchema || 'reject' in approvalSchema + if (!hasBranches) return approvalSchema + return approved ? approvalSchema['approve'] : approvalSchema['reject'] + } + + private validationCode(item: RuntimeInterrupt): ItemInterruptError['code'] { + return item.kind === 'client-tool-execution' + ? 'invalid-tool-output' + : 'invalid-payload' + } + + private resolveBooleanBulk(approved: boolean): void { + const eligible = this.items.every( + (item) => + item.kind === 'tool-approval' && + this.approvalBranchSchema(item.tool, approved) === undefined, + ) + if (!eligible || this.items.length === 0) { + this.addRootError( + 'unsupported-bulk-operation', + 'Boolean bulk resolution requires payloadless tool approvals.', + false, + ) + return + } + this.invalidateRetry() + for (const item of this.items) { + item.validationGeneration++ + item.resolution = cloneAndDeepFreezeJson({ + interruptId: item.descriptor.id, + status: 'resolved', + payload: { approved }, + }) + item.status = 'staged' + item.error = undefined + } + this.publish() + this.maybeSubmit() + } + + private resolveTransaction( + resolver: (interrupt: ChatInterrupt) => unknown, + ): void { + const checkpoints = this.items.map((item) => ({ + status: item.status, + ...(item.resolution !== undefined ? { resolution: item.resolution } : {}), + ...(item.error !== undefined ? { error: item.error } : {}), + validationGeneration: item.validationGeneration, + })) + const token: TransactionToken = { active: true } + this.activeTransaction = token + const stable = this.buildSnapshot(token) + let failure: + | { code: BatchInterruptError['code']; message: string } + | undefined + try { + for (const interrupt of stable) { + const result = resolver(interrupt) + if (result !== undefined) { + failure = { + code: isPromiseLike(result) + ? 'async-resolver' + : 'inactive-transaction', + message: isPromiseLike(result) + ? 'Interrupt transaction resolvers must be synchronous.' + : 'Interrupt transaction resolvers must return literal undefined.', + } + break + } + } + if ( + failure === undefined && + this.items.some( + (item) => item.resolution === undefined || item.status !== 'staged', + ) + ) { + failure = { + code: 'incomplete-batch', + message: 'Interrupt transaction did not resolve every item.', + } + } + } catch (error) { + failure = { + code: 'item-validation-failed', + message: error instanceof Error ? error.message : String(error), + } + } finally { + token.active = false + this.activeTransaction = undefined + } + + if (failure) { + this.restoreCheckpoints(checkpoints) + this.addRootError(failure.code, failure.message, false) + return + } + this.publish() + this.maybeSubmit() + } + + private restoreCheckpoints( + checkpoints: ReadonlyArray, + ): void { + this.items.forEach((item, index) => { + const checkpoint = checkpoints[index] + if (!checkpoint) return + item.status = checkpoint.status + item.resolution = checkpoint.resolution + item.error = checkpoint.error + item.validationGeneration = checkpoint.validationGeneration + 1 + }) + this.publish() + } + + private assertItemMutable(transaction?: TransactionToken): void { + if (transaction && !transaction.active) { + throw new Error('Interrupt transaction is inactive.') + } + if (this.activeTransaction && transaction !== this.activeTransaction) { + throw new Error('Interrupt transaction is inactive.') + } + if (this.resuming) { + throw new Error('Interrupts cannot be mutated while submitting.') + } + } + + private assertRootMutable(): void { + if (this.activeTransaction) { + throw new Error('Interrupt transaction is already active.') + } + if (this.resuming) { + throw new Error('Interrupts cannot be mutated while submitting.') + } + } + + private invalidateRetry(): void { + this.retrySubmission = undefined + } + + private submitBatch(submission: InterruptManagerSubmission): void { + this.resuming = true + this.retrySubmission = undefined + for (const item of this.items) item.status = 'submitting' + this.publish() + void this.performSubmission(submission) + } + + private async performSubmission( + submission: InterruptManagerSubmission, + ): Promise { + try { + const result = await this.options.submit(submission) + if (result?.status === 'conflict') { + this.addRootError( + 'conflict', + 'Interrupt submission conflicted with authoritative state.', + false, + 'server', + ) + await this.options.recover?.(result.authoritativeState) + } + } catch (error) { + await this.handleSubmissionFailure(error, submission) + } finally { + this.resuming = false + this.publish() + } + } + + private async handleSubmissionFailure( + error: unknown, + submission: InterruptManagerSubmission, + ): Promise { + const errors = readSubmissionErrors(error) + if (errors.length === 0) { + const message = error instanceof Error ? error.message : String(error) + this.addRootError('transport', message, true, 'transport') + this.retrySubmission = submission + for (const item of this.items) item.status = 'error' + return + } + + let requiresRecovery = false + let retryable = false + for (const submissionError of errors) { + if ( + submissionError.code === 'stale' || + submissionError.code === 'expired' || + submissionError.code === 'conflict' + ) { + requiresRecovery = true + } + retryable ||= submissionError.retryable + if (submissionError.scope === 'item') { + const item = this.items.find( + (candidate) => + candidate.descriptor.id === submissionError.interruptId, + ) + if (item) { + item.status = 'error' + item.error = cloneAndDeepFreezeJson(submissionError) + } + } else { + this.rootErrors = Object.freeze([ + ...this.rootErrors, + cloneAndDeepFreezeJson(submissionError), + ]) + } + } + for (const item of this.items) { + if (item.status === 'submitting') item.status = 'error' + } + this.retrySubmission = + retryable && !requiresRecovery ? submission : undefined + if (requiresRecovery) await this.options.recover?.() + } + + private addRootError( + code: BatchInterruptError['code'], + message: string, + retryable: boolean, + source: BatchInterruptError['source'] = 'client', + ): void { + const hydration = this.requireHydration() + this.rootErrors = Object.freeze([ + ...this.rootErrors, + Object.freeze({ + scope: 'batch' as const, + code, + message, + source, + retryable, + interruptIds: Object.freeze( + this.items.map((item) => item.descriptor.id), + ), + threadId: hydration.threadId, + interruptedRunId: hydration.interruptedRunId, + generation: hydration.generation, + }), + ]) + this.publish() + } + + private findItem(interruptId: string): RuntimeInterrupt { + const item = this.items.find( + (candidate) => candidate.descriptor.id === interruptId, + ) + if (!item) throw new Error(`Unknown interrupt: ${interruptId}`) + return item + } + + private requireHydration(): InterruptManagerHydration { + if (!this.hydration) throw new Error('InterruptManager is not hydrated.') + return this.hydration + } + + private itemError( + interruptId: string, + code: ItemInterruptError['code'], + message: string, + path?: ReadonlyArray, + ): ItemInterruptError { + const hydration = this.requireHydration() + return Object.freeze({ + scope: 'item', + interruptId, + code, + message, + ...(path !== undefined ? { path: Object.freeze([...path]) } : {}), + source: 'client', + retryable: false, + threadId: hydration.threadId, + interruptedRunId: hydration.interruptedRunId, + generation: hydration.generation, + }) + } +} diff --git a/packages/ai-client/src/types.ts b/packages/ai-client/src/types.ts index 5ed1b0076..5564b829a 100644 --- a/packages/ai-client/src/types.ts +++ b/packages/ai-client/src/types.ts @@ -1,15 +1,25 @@ import type { AnyClientTool, + ApprovalCapabilityOf, + ApprovalSchemaOf, AudioPart, + BatchInterruptError, ChunkStrategy, ContentPart, DocumentPart, ImagePart, + InferSchemaType, InferToolInput, InferToolOutput, + InputSchemaOf, Interrupt, + InterruptBinding, + InterruptRecoveryStateV1, + ItemInterruptError, ModelMessage, + NoSchema, RunAgentResumeItem, + SchemaInput, StreamChunk, StructuredOutputPart, UIResourcePart, @@ -28,11 +38,173 @@ export interface ChatResumeState { export type ChatPendingInterrupt = Interrupt -export interface ChatResumeSnapshot { +export interface ChatResumeSnapshotV1 { + schemaVersion?: 1 resumeState: ChatResumeState pendingInterrupts?: Array } +/** JSON-safe client draft state. Authoritative interrupt state stays in recoveryState. */ +export interface PersistedInterruptDraft { + readonly interruptId: string + readonly response: unknown + readonly status: InterruptItemStatus + readonly error?: ItemInterruptError +} + +export interface ChatResumeSnapshotV2 { + schemaVersion: 2 + resumeState: ChatResumeState + pendingInterrupts?: Array + interruptState?: { + recoveryState: InterruptRecoveryStateV1 + drafts: ReadonlyArray + } +} + +export type ChatResumeSnapshot = ChatResumeSnapshotV1 | ChatResumeSnapshotV2 + +export type ChatContinuationLoader = ( + continuationRunId: string, + abortSignal?: AbortSignal, +) => AsyncIterable + +export type InterruptItemStatus = + | 'pending' + | 'validating' + | 'staged' + | 'submitting' + | 'error' + +export interface BoundInterruptBase { + readonly id: string + readonly interruptId: string + readonly reason: string + readonly message?: string + readonly responseSchema?: Readonly> + readonly expiresAt?: string + readonly metadata?: Readonly> + readonly threadId: string + readonly interruptedRunId: string + readonly generation: number + readonly status: InterruptItemStatus + readonly binding: Readonly + readonly errors: ReadonlyArray + /** @deprecated Use `errors[0]`. */ + readonly error?: ItemInterruptError + readonly canResolve: boolean + cancel: () => void + clearResolution: () => void +} + +export interface GenericAGUIInterrupt extends BoundInterruptBase { + readonly kind: 'generic' + resolveInterrupt: (payload: unknown) => void +} + +type ApprovalBranchSchema = + ApprovalSchemaOf extends infer TApproval + ? TApproval extends { approve?: SchemaInput; reject?: SchemaInput } + ? Exclude + : TApproval extends SchemaInput + ? TApproval + : never + : never + +type ApprovalEdits = + InputSchemaOf extends NoSchema + ? { editedArgs?: never } + : { editedArgs?: InferToolInput } + +type ApprovalPayload = [TSchema] extends [never] + ? { payload?: never } + : TSchema extends SchemaInput + ? { payload: InferSchemaType } + : { payload?: never } + +type ApproveArguments = [ + ApprovalBranchSchema, +] extends [never] + ? InputSchemaOf extends NoSchema + ? [options?: never] + : [options?: ApprovalEdits & { payload?: never }] + : [ + options: ApprovalEdits & + ApprovalPayload>, + ] + +type RejectArguments = [ApprovalBranchSchema] extends [ + never, +] + ? [options?: never] + : [ + options: { editedArgs?: never } & ApprovalPayload< + ApprovalBranchSchema + >, + ] + +export type ToolApprovalInterrupt = + TTool extends AnyClientTool + ? BoundInterruptBase & { + readonly kind: 'tool-approval' + readonly toolName: TTool['name'] + readonly toolCallId: string + readonly originalArgs: InferToolInput + resolveInterrupt: { + (approved: true, ...args: ApproveArguments): void + (approved: false, ...args: RejectArguments): void + } + } + : never + +export type ClientToolExecutionInterrupt< + TTool extends AnyClientTool = AnyClientTool, +> = TTool extends AnyClientTool + ? BoundInterruptBase & { + readonly kind: 'client-tool-execution' + readonly toolName: TTool['name'] + readonly toolCallId: string + resolveInterrupt: (output: InferToolOutput) => void + } + : never + +type ApprovalInterrupts> = + TTools[number] extends infer TTool + ? TTool extends AnyClientTool + ? ApprovalCapabilityOf extends true + ? ToolApprovalInterrupt + : never + : never + : never + +type ClientToolExecutionInterrupts< + TTools extends ReadonlyArray, +> = TTools[number] extends infer TTool + ? TTool extends AnyClientTool + ? ClientToolExecutionInterrupt + : never + : never + +export type ChatInterrupt< + TTools extends ReadonlyArray = ReadonlyArray, +> = + | GenericAGUIInterrupt + | ApprovalInterrupts + | ClientToolExecutionInterrupts + +export type BoundInterrupts< + TTools extends ReadonlyArray = ReadonlyArray, +> = ReadonlyArray> + +export interface ChatInterruptState< + TTools extends ReadonlyArray = ReadonlyArray, +> { + readonly interrupts: BoundInterrupts + readonly pendingInterrupts: BoundInterrupts + readonly interruptErrors: ReadonlyArray + readonly resuming: boolean +} + /** * `messages` is the full UIMessage history (not a delta). `data` is the * merged body — `ChatClientOptions.body` plus any per-call data passed to @@ -45,6 +217,7 @@ export interface ChatFetcherInput { data?: Record threadId: string runId: string + parentRunId?: string resume?: Array } @@ -435,6 +608,9 @@ export interface ChatClientBaseOptions< */ initialResumeSnapshot?: ChatResumeSnapshot + /** Explicit read-only loader for replaying a committed continuation run. */ + continuationLoader?: ChatContinuationLoader + /** * Arbitrary client-controlled JSON forwarded to the server in the * AG-UI `RunAgentInput.forwardedProps` field. Use this for per-session @@ -527,9 +703,12 @@ export interface ChatClientBaseOptions< */ onResumeStateChange?: ( resumeState: ChatResumeState | null, - pendingInterrupts: Array, + pendingInterrupts: BoundInterrupts, ) => void + /** Callback when the immutable interrupt state snapshot changes. */ + onInterruptStateChange?: (state: ChatInterruptState) => void + /** * Callback when a custom event is received from a server-side tool. * Custom events are emitted by tools using `context.emitCustomEvent()` during execution. diff --git a/packages/ai-client/tests/chat-client-interrupts.test.ts b/packages/ai-client/tests/chat-client-interrupts.test.ts new file mode 100644 index 000000000..4d69daada --- /dev/null +++ b/packages/ai-client/tests/chat-client-interrupts.test.ts @@ -0,0 +1,1486 @@ +import { describe, expect, it, vi } from 'vitest' +import { + InterruptPersistenceCapability, + chat, + defineChatMiddleware, + provideInterruptPersistence, +} from '@tanstack/ai' +import { + EventType, + canonicalInterruptJson, + convertSchemaToJsonSchema, + digestInterruptJson, + hashSchemaInput, + normalizeApprovalSchema, + toolDefinition, +} from '@tanstack/ai/client' +import { z } from 'zod' +import { InterruptManager } from '../src/interrupt-manager' +import { ChatClient } from '../src/chat-client' +import type { + AnyTextAdapter, + InterruptRecoveryStateV1, + InterruptPersistenceGateway, + StreamChunk, +} from '@tanstack/ai' +import type { Interrupt, InterruptBinding } from '@tanstack/ai/client' +import type { StandardSchemaV1 } from '@standard-schema/spec' +import type { InterruptManagerSubmission } from '../src/interrupt-manager' +import type { ChatInterrupt } from '../src/types' +import type { + ConnectConnectionAdapter, + RunAgentInputContext, +} from '../src/connection-adapters' +import type { UIMessage } from '../src/types' +import type { ChatResumeSnapshotV2, ChatServerPersistence } from '../src/types' + +const transferDefinition = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema: z.object({ cents: z.number() }), + outputSchema: z.object({ receipt: z.string() }), + approvalSchema: { + approve: z.object({ note: z.string() }), + reject: z.object({ reason: z.string() }), + }, +}) + +const lookupDefinition = toolDefinition({ + name: 'lookup', + description: 'Look up an account', + outputSchema: z.object({ accountId: z.string() }), +}) + +const tools = [transferDefinition.client(), lookupDefinition.client()] as const + +function descriptor( + binding: InterruptBinding, + overrides: Partial = {}, +): Interrupt { + return { + id: binding.interruptId, + reason: + binding.kind === 'tool-approval' + ? 'tool_call' + : binding.kind === 'client-tool-execution' + ? 'tanstack:client_tool_execution' + : 'confirmation', + ...(binding.kind !== 'generic' && { toolCallId: binding.toolCallId }), + metadata: { 'tanstack:interruptBinding': binding }, + ...overrides, + } +} + +function createManager() { + const submit = vi.fn(async (_submission: InterruptManagerSubmission) => ({ + status: 'committed' as const, + continuationRunId: 'continuation-1', + })) + const manager = new InterruptManager({ tools, submit }) + return { manager, submit } +} + +function genericDescriptor(id: string): Interrupt { + return descriptor({ + kind: 'generic', + interruptId: id, + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'none', + }) +} + +async function settle(): Promise { + await Promise.resolve() + await Promise.resolve() +} + +describe('InterruptManager hydration', () => { + it('hydrates correlated approval and client-tool bindings into frozen typed snapshots', () => { + const approval = normalizeApprovalSchema( + transferDefinition.approvalSchema, + transferDefinition.inputSchema, + ) + const approvalBinding: InterruptBinding = { + kind: 'tool-approval', + interruptId: 'approval-1', + interruptedRunId: 'run-1', + generation: 3, + toolName: 'transfer', + toolCallId: 'call-1', + originalArgs: { cents: 100 }, + inputSchemaHash: hashSchemaInput(transferDefinition.inputSchema), + approvalSchemaHash: approval.approvalSchemaHash, + responseSchemaHash: approval.responseSchemaHash, + } + const outputSchemaHash = hashSchemaInput(lookupDefinition.outputSchema) + const clientBinding: InterruptBinding = { + kind: 'client-tool-execution', + interruptId: 'client-1', + interruptedRunId: 'run-1', + generation: 3, + toolName: 'lookup', + toolCallId: 'call-2', + outputSchemaHash, + responseSchemaHash: outputSchemaHash, + } + const { manager } = createManager() + + manager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 3, + interrupts: [ + descriptor(approvalBinding, { + responseSchema: approval.responseSchema, + }), + descriptor(clientBinding), + ], + }) + + const snapshot = manager.getInterrupts() + expect(snapshot.map((item) => item.kind)).toEqual([ + 'tool-approval', + 'client-tool-execution', + ]) + expect(Object.isFrozen(snapshot)).toBe(true) + expect(Object.isFrozen(snapshot[0])).toBe(true) + expect(Object.isFrozen(snapshot[0]?.binding)).toBe(true) + }) + + it('hydrates a real core client-tool terminal with distinct schema identity hashes', async () => { + const coreChunks: Array = [ + { + type: EventType.RUN_STARTED, + runId: 'core-run', + threadId: 'core-thread', + timestamp: 1, + }, + { + type: EventType.TOOL_CALL_START, + toolCallId: 'core-call', + toolCallName: lookupDefinition.name, + toolName: lookupDefinition.name, + timestamp: 1, + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'core-call', + delta: '{}', + timestamp: 1, + }, + { + type: EventType.RUN_FINISHED, + runId: 'core-run', + threadId: 'core-thread', + finishReason: 'tool_calls', + timestamp: 1, + }, + ] + const adapter: AnyTextAdapter = { + kind: 'text', + name: 'core-interrupt-test', + model: 'test-model', + '~types': { + providerOptions: {}, + inputModalities: ['text'], + messageMetadataByModality: { + text: undefined, + image: undefined, + audio: undefined, + video: undefined, + document: undefined, + }, + toolCapabilities: [], + toolCallMetadata: undefined, + systemPromptMetadata: undefined, + }, + chatStream: () => + (async function* () { + await Promise.resolve() + for (const chunk of coreChunks) yield chunk + })(), + structuredOutput: () => Promise.resolve({ data: {}, rawText: '{}' }), + } + const gateway: InterruptPersistenceGateway = { + openInterruptBatch: (input) => + Promise.resolve({ generation: 1, descriptors: input.descriptors }), + commitInterruptResolutions: (input) => + Promise.resolve({ + status: 'committed', + continuationRunId: input.continuationRunId, + }), + getInterruptRecoveryState: (input) => + Promise.resolve({ + schemaVersion: 1, + state: 'missing', + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + generation: input.knownGeneration, + pendingInterrupts: [], + }), + } + const persistence = defineChatMiddleware({ + name: 'core-interrupt-test-persistence', + provides: [InterruptPersistenceCapability], + setup(ctx) { + provideInterruptPersistence(ctx, gateway) + }, + async onChunk(_ctx, chunk) { + if ( + chunk.type !== EventType.RUN_FINISHED || + chunk.outcome?.type !== 'interrupt' + ) { + return + } + const opened = await gateway.openInterruptBatch({ + threadId: chunk.threadId, + interruptedRunId: chunk.runId, + descriptors: chunk.outcome.interrupts, + bindings: [], + }) + return { + ...chunk, + outcome: { + ...chunk.outcome, + interrupts: chunk.outcome.interrupts.map((interrupt) => { + const raw = interrupt.metadata?.['tanstack:interruptBinding'] + if ( + raw === null || + typeof raw !== 'object' || + Array.isArray(raw) + ) { + throw new Error('Expected a core interrupt binding.') + } + return { + ...interrupt, + metadata: { + ...interrupt.metadata, + 'tanstack:interruptBinding': { + ...Object.fromEntries(Object.entries(raw)), + interruptedRunId: chunk.runId, + generation: opened.generation, + }, + }, + } + }), + }, + } + }, + }) + const emitted: Array = [] + for await (const chunk of chat({ + adapter, + messages: [{ role: 'user', content: 'Look up an account' }], + tools: [tools[1]], + runId: 'core-run', + threadId: 'core-thread', + middleware: [persistence], + }) as AsyncIterable) { + emitted.push(chunk) + } + const terminal = emitted.find( + (chunk) => + chunk.type === EventType.RUN_FINISHED && + chunk.outcome?.type === 'interrupt', + ) + if ( + terminal?.type !== EventType.RUN_FINISHED || + terminal.outcome?.type !== 'interrupt' + ) { + throw new Error('Expected a real core interrupt terminal.') + } + const interrupt = terminal.outcome.interrupts[0] + if (!interrupt) throw new Error('Expected a client-tool interrupt.') + const rawBinding = interrupt.metadata?.['tanstack:interruptBinding'] + if ( + rawBinding === null || + typeof rawBinding !== 'object' || + Array.isArray(rawBinding) + ) { + throw new Error('Expected a public interrupt binding.') + } + const binding: Record = Object.fromEntries( + Object.entries(rawBinding), + ) + const expectedOutputSchemaHash = hashSchemaInput( + lookupDefinition.outputSchema, + ) + const expectedResponseSchema = + convertSchemaToJsonSchema(lookupDefinition.outputSchema) ?? {} + const expectedResponseSchemaHash = digestInterruptJson( + canonicalInterruptJson(expectedResponseSchema), + ) + expect(binding.outputSchemaHash).toBe(expectedOutputSchemaHash) + expect(binding.responseSchemaHash).toBe(expectedResponseSchemaHash) + expect(binding.outputSchemaHash).not.toBe(binding.responseSchemaHash) + + const { manager } = createManager() + manager.hydrate({ + threadId: 'core-thread', + interruptedRunId: 'core-run', + generation: 1, + interrupts: terminal.outcome.interrupts, + }) + expect(manager.getInterrupts()[0]?.kind).toBe('client-tool-execution') + + manager.hydrate({ + threadId: 'core-thread', + interruptedRunId: 'core-run', + generation: 1, + interrupts: [ + { + ...interrupt, + metadata: { + ...interrupt.metadata, + 'tanstack:interruptBinding': { + ...binding, + outputSchemaHash: 'sha256:configured-schema-drift', + }, + }, + }, + ], + }) + expect(manager.getInterrupts()[0]?.kind).toBe('generic') + }) + + it('keeps deprecated approval and client-tool reason aliases compatible', () => { + const approval = normalizeApprovalSchema( + transferDefinition.approvalSchema, + transferDefinition.inputSchema, + ) + const approvalBinding: InterruptBinding = { + kind: 'tool-approval', + interruptId: 'approval-legacy', + interruptedRunId: 'run-legacy', + generation: 1, + toolName: 'transfer', + toolCallId: 'call-approval-legacy', + originalArgs: { cents: 100 }, + inputSchemaHash: hashSchemaInput(transferDefinition.inputSchema), + approvalSchemaHash: approval.approvalSchemaHash, + responseSchemaHash: approval.responseSchemaHash, + } + const outputSchemaHash = hashSchemaInput(lookupDefinition.outputSchema) + const clientBinding: InterruptBinding = { + kind: 'client-tool-execution', + interruptId: 'client-legacy', + interruptedRunId: 'run-legacy', + generation: 1, + toolName: 'lookup', + toolCallId: 'call-client-legacy', + outputSchemaHash, + responseSchemaHash: outputSchemaHash, + } + const { manager } = createManager() + + manager.hydrate({ + threadId: 'thread-legacy', + interruptedRunId: 'run-legacy', + generation: 1, + interrupts: [ + descriptor(approvalBinding, { + reason: 'approval_required', + responseSchema: approval.responseSchema, + }), + descriptor(clientBinding, { reason: 'client_tool_input' }), + ], + }) + + expect(manager.getInterrupts().map((interrupt) => interrupt.kind)).toEqual([ + 'tool-approval', + 'client-tool-execution', + ]) + }) + + it('degrades mismatched tool correlation to generic without trusting wire correlation', () => { + const outputSchemaHash = hashSchemaInput(lookupDefinition.outputSchema) + const binding: InterruptBinding = { + kind: 'client-tool-execution', + interruptId: 'client-1', + interruptedRunId: 'untrusted-run', + generation: 99, + toolName: 'lookup', + toolCallId: 'expected-call', + outputSchemaHash, + responseSchemaHash: outputSchemaHash, + } + const { manager } = createManager() + + manager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 3, + interrupts: [descriptor(binding, { toolCallId: 'other-call' })], + }) + + expect(manager.getInterrupts()[0]).toMatchObject({ + kind: 'generic', + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 3, + }) + }) + + it('keeps cancellation available when a generic response schema is invalid', () => { + const binding: InterruptBinding = { + kind: 'generic', + interruptId: 'generic-1', + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'invalid-schema', + } + const { manager } = createManager() + manager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + interrupts: [ + descriptor(binding, { + responseSchema: { + $schema: 'https://json-schema.org/draft/2019-09/schema', + type: 'object', + }, + }), + ], + }) + + const item = manager.getInterrupts()[0] + expect(item?.kind).toBe('generic') + expect(item?.canResolve).toBe(false) + item?.cancel() + expect(manager.getInterrupts()[0]?.status).toBe('submitting') + }) +}) + +describe('InterruptManager transactions', () => { + it('waits for a complete multi-item batch and submits a singleton immediately', async () => { + const multi = createManager() + multi.manager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + interrupts: [genericDescriptor('one'), genericDescriptor('two')], + }) + + const first = multi.manager.getInterrupts()[0] + if (first?.kind !== 'generic') throw new Error('Expected generic interrupt') + first.resolveInterrupt('first') + expect(multi.submit).not.toHaveBeenCalled() + expect(multi.manager.getInterrupts()[0]?.status).toBe('staged') + multi.manager.getInterrupts()[1]?.cancel() + expect(multi.submit).toHaveBeenCalledTimes(1) + + const single = createManager() + single.manager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + interrupts: [genericDescriptor('only')], + }) + const only = single.manager.getInterrupts()[0] + if (only?.kind !== 'generic') throw new Error('Expected generic interrupt') + only.resolveInterrupt('done') + expect(single.submit).toHaveBeenCalledTimes(1) + await settle() + }) + + it('validates async Standard Schema candidates and preserves a prior valid draft', async () => { + const asyncOutputSchema: StandardSchemaV1 = + { + '~standard': { + version: 1, + vendor: 'test', + validate: async (value) => + isAccountOutput(value) + ? { value } + : { issues: [{ message: 'accountId is required' }] }, + }, + } + const asyncTool = toolDefinition({ + name: 'asyncLookup', + description: 'Async validation', + outputSchema: asyncOutputSchema, + }).client() + const outputSchemaHash = hashSchemaInput(asyncOutputSchema) + const submit = vi.fn( + async (_submission: InterruptManagerSubmission) => undefined, + ) + const manager = new InterruptManager({ + tools: [asyncTool] as const, + submit, + }) + const binding: InterruptBinding = { + kind: 'client-tool-execution', + interruptId: 'async-1', + interruptedRunId: 'run-1', + generation: 1, + toolName: 'asyncLookup', + toolCallId: 'call-1', + outputSchemaHash, + responseSchemaHash: outputSchemaHash, + } + manager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + interrupts: [descriptor(binding), genericDescriptor('other')], + }) + + const first = manager.getInterrupts()[0] + if (first?.kind !== 'client-tool-execution') { + throw new Error('Expected client tool execution interrupt') + } + first.resolveInterrupt({ accountId: 'valid' }) + expect(manager.getInterrupts()[0]?.status).toBe('validating') + await settle() + expect(manager.getInterrupts()[0]?.status).toBe('staged') + + const replacement = manager.getInterrupts()[0] + if (replacement?.kind !== 'client-tool-execution') { + throw new Error('Expected client tool execution interrupt') + } + replacement.resolveInterrupt({ accountId: '' }) + await settle() + expect(manager.getInterrupts()[0]).toMatchObject({ + status: 'error', + errors: [{ code: 'invalid-tool-output' }], + error: { code: 'invalid-tool-output' }, + }) + expect(Object.isFrozen(manager.getInterrupts()[0]?.errors)).toBe(true) + manager.getInterrupts()[1]?.cancel() + expect(submit).not.toHaveBeenCalled() + manager.getInterrupts()[0]?.clearResolution() + expect(manager.getInterrupts()[0]?.status).toBe('pending') + }) + + it('rolls callback transactions back on thrown, returned, thenable, or incomplete work', () => { + const { manager, submit } = createManager() + manager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + interrupts: [genericDescriptor('one'), genericDescriptor('two')], + }) + + manager.resolve(() => { + throw new Error('transaction failed') + }) + Reflect.apply(manager.resolve, manager, [ + (item: ChatInterrupt) => { + item.cancel() + return 'not undefined' + }, + ]) + Reflect.apply(manager.resolve, manager, [() => Promise.resolve()]) + manager.resolve((item) => { + if (item.id === 'one') item.cancel() + return undefined + }) + + expect(submit).not.toHaveBeenCalled() + expect( + manager.getInterrupts().every((item) => item.status === 'pending'), + ).toBe(true) + expect(manager.getInterruptErrors()).toHaveLength(4) + }) + + it('seals transaction items, invokes the callback once per item, and submits atomically', () => { + const { manager, submit } = createManager() + manager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + interrupts: [genericDescriptor('one'), genericDescriptor('two')], + }) + const calls: Array = [] + let lateCancel: (() => void) | undefined + + manager.resolve((item) => { + calls.push(item.id) + if (item.id === 'one') lateCancel = item.cancel + item.cancel() + return undefined + }) + + expect(calls).toEqual(['one', 'two']) + expect(submit).toHaveBeenCalledTimes(1) + expect(() => lateCancel?.()).toThrow('inactive') + expect(() => manager.getInterrupts()[0]?.clearResolution()).toThrow( + 'submitting', + ) + }) + + it('permits boolean bulk resolution only for payloadless tool approvals and cancels all payloadlessly', () => { + const approvalTool = toolDefinition({ + name: 'confirmOnly', + description: 'Confirm only', + needsApproval: true, + }).client() + const approval = normalizeApprovalSchema(undefined, undefined) + const binding: InterruptBinding = { + kind: 'tool-approval', + interruptId: 'approval-1', + interruptedRunId: 'run-1', + generation: 1, + toolName: 'confirmOnly', + toolCallId: 'call-1', + originalArgs: {}, + inputSchemaHash: hashSchemaInput(undefined), + approvalSchemaHash: approval.approvalSchemaHash, + responseSchemaHash: approval.responseSchemaHash, + } + const submit = vi.fn( + async (_submission: InterruptManagerSubmission) => undefined, + ) + const approvals = new InterruptManager({ + tools: [approvalTool] as const, + submit, + }) + approvals.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + interrupts: [ + descriptor(binding, { responseSchema: approval.responseSchema }), + ], + }) + approvals.resolve(true) + expect(submit.mock.calls[0]?.[0].resolutions).toEqual([ + { + interruptId: 'approval-1', + status: 'resolved', + payload: { approved: true }, + }, + ]) + + const generic = createManager() + generic.manager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + interrupts: [genericDescriptor('generic')], + }) + generic.manager.resolve(false) + expect(generic.submit).not.toHaveBeenCalled() + expect(generic.manager.getInterruptErrors()[0]?.code).toBe( + 'unsupported-bulk-operation', + ) + generic.manager.cancel() + expect(generic.submit.mock.calls[0]?.[0].resolutions).toEqual([ + { interruptId: 'generic', status: 'cancelled' }, + ]) + }) + + it('retries the exact frozen batch only for retryable failures and recovers conflicts', async () => { + const retryable = { + scope: 'batch' as const, + code: 'transport' as const, + message: 'offline', + source: 'transport' as const, + retryable: true, + interruptIds: ['generic'], + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + } + const submit = vi + .fn(async (_submission: InterruptManagerSubmission) => undefined) + .mockRejectedValueOnce({ errors: [retryable] }) + .mockResolvedValue(undefined) + const recover = vi.fn(async () => undefined) + const manager = new InterruptManager({ submit, recover }) + manager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + interrupts: [genericDescriptor('generic')], + }) + const retryItem = manager.getInterrupts()[0] + if (retryItem?.kind !== 'generic') { + throw new Error('Expected generic interrupt') + } + retryItem.resolveInterrupt('answer') + await settle() + + const firstSubmission = submit.mock.calls[0]?.[0] + manager.retry() + expect(submit.mock.calls[1]?.[0]).toBe(firstSubmission) + await settle() + manager.getInterrupts()[0]?.clearResolution() + manager.retry() + expect(submit).toHaveBeenCalledTimes(2) + + const conflictSubmit = vi.fn( + async (_submission: InterruptManagerSubmission) => ({ + status: 'conflict' as const, + authoritativeState: { + schemaVersion: 1 as const, + state: 'pending' as const, + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 2, + pendingInterrupts: [], + }, + }), + ) + const conflictManager = new InterruptManager({ + submit: conflictSubmit, + recover, + }) + conflictManager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + interrupts: [genericDescriptor('generic')], + }) + conflictManager.getInterrupts()[0]?.cancel() + await settle() + expect(recover).toHaveBeenCalled() + }) +}) + +function isAccountOutput(value: unknown): value is { accountId: string } { + return ( + value !== null && + typeof value === 'object' && + 'accountId' in value && + typeof value.accountId === 'string' && + value.accountId.length > 0 + ) +} + +describe('ChatClient native interrupts', () => { + it('publishes the shared immutable interrupt state callback', () => { + const onInterruptStateChange = vi.fn() + const interrupt = genericDescriptor('generic-1') + + const client = new ChatClient({ + connection: { async *connect() {} }, + onInterruptStateChange, + initialResumeSnapshot: { + schemaVersion: 2, + resumeState: { threadId: 'thread-1', runId: 'run-1' }, + pendingInterrupts: [interrupt], + interruptState: { + recoveryState: { + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + pendingInterrupts: [interrupt], + }, + drafts: [], + }, + }, + }) + + expect(onInterruptStateChange).toHaveBeenLastCalledWith( + client.getInterruptState(), + ) + const state = onInterruptStateChange.mock.lastCall?.[0] + expect(state?.interrupts).toBe(state?.pendingInterrupts) + }) + + it('owns one immutable interrupt state and resumes with a fresh child run', async () => { + const contexts: Array = [] + const sentMessages: Array | Array> = [] + const binding: InterruptBinding = { + kind: 'generic', + interruptId: 'generic-1', + interruptedRunId: 'placeholder', + generation: 4, + responseSchemaHash: 'none', + } + let call = 0 + const connection: ConnectConnectionAdapter = { + async *connect(messages, _data, _signal, context) { + contexts.push(context) + sentMessages.push(messages) + call++ + const runId = context?.runId ?? `run-${call}` + const threadId = context?.threadId ?? 'thread-1' + yield { + type: EventType.RUN_STARTED, + runId, + threadId, + timestamp: Date.now(), + } + if (call === 1) { + binding.interruptedRunId = runId + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + outcome: { + type: 'interrupt', + interrupts: [descriptor(binding)], + }, + } + return + } + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + outcome: { type: 'success' }, + } + }, + } + const client = new ChatClient({ connection, threadId: 'thread-1' }) + + await client.sendMessage('start') + const state = client.getInterruptState() + expect(Object.isFrozen(state)).toBe(true) + expect(state.interrupts).toBe(state.pendingInterrupts) + expect(client.getInterrupts()).toBe(state.interrupts) + expect(client.getPendingInterrupts()).toBe(state.interrupts) + const item = state.interrupts[0] + if (item?.kind !== 'generic') throw new Error('Expected generic interrupt') + item.resolveInterrupt({ answer: 42 }) + + await vi.waitFor(() => expect(contexts).toHaveLength(2)) + expect(contexts[1]).toMatchObject({ + threadId: 'thread-1', + parentRunId: contexts[0]?.runId, + resume: [ + { + interruptId: 'generic-1', + status: 'resolved', + payload: { answer: 42 }, + }, + ], + }) + expect(contexts[1]?.runId).not.toBe(contexts[0]?.runId) + expect(sentMessages[1]).toEqual([]) + }) + + it('preserves the exact interrupt batch and joins a replay after a child transport failure', async () => { + const contexts: Array = [] + const joined: Array = [] + const binding: InterruptBinding = { + kind: 'generic', + interruptId: 'generic-1', + interruptedRunId: 'placeholder', + generation: 1, + responseSchemaHash: 'none', + } + let calls = 0 + const connection: ConnectConnectionAdapter = { + async *connect(_messages, _data, _signal, context) { + contexts.push(context) + calls++ + const runId = context?.runId ?? `run-${calls}` + if (calls === 1) { + binding.interruptedRunId = runId + yield { + type: EventType.RUN_FINISHED, + threadId: 'thread-1', + runId, + timestamp: Date.now(), + outcome: { + type: 'interrupt', + interrupts: [descriptor(binding)], + }, + } + return + } + if (calls === 2) throw new Error('continuation transport failed') + yield { + type: EventType.RUN_FINISHED, + threadId: 'thread-1', + runId, + timestamp: Date.now(), + outcome: { type: 'success' }, + result: { + replayed: true, + continuationRunId: 'winner-run', + }, + } + }, + async *joinRun(runId) { + joined.push(runId) + yield { + type: EventType.RUN_FINISHED, + threadId: 'thread-1', + runId, + timestamp: Date.now(), + outcome: { type: 'success' }, + } + }, + } + const client = new ChatClient({ connection, threadId: 'thread-1' }) + + await client.sendMessage('start') + client.getInterrupts()[0]?.resolveInterrupt({ answer: 42 }) + + await vi.waitFor(() => expect(contexts).toHaveLength(2)) + await vi.waitFor(() => + expect(client.getInterrupts()[0]).toMatchObject({ + id: 'generic-1', + status: 'error', + }), + ) + expect(Object.isFrozen(client.getInterrupts())).toBe(true) + client.retryInterrupts() + + await vi.waitFor(() => expect(contexts).toHaveLength(3)) + expect(contexts[2]?.resume).toEqual(contexts[1]?.resume) + await vi.waitFor(() => expect(joined).toEqual(['winner-run'])) + }) + + it('exposes an atomic all-item resolveInterrupts callback transaction', async () => { + const contexts: Array = [] + let calls = 0 + const connection: ConnectConnectionAdapter = { + async *connect(_messages, _data, _signal, context) { + contexts.push(context) + calls++ + const runId = context?.runId ?? `run-${calls}` + if (calls === 1) { + yield { + type: EventType.RUN_FINISHED, + threadId: 'thread-1', + runId, + timestamp: Date.now(), + outcome: { + type: 'interrupt', + interrupts: [ + descriptor({ + kind: 'generic', + interruptId: 'first', + interruptedRunId: runId, + generation: 1, + responseSchemaHash: 'none', + }), + descriptor({ + kind: 'generic', + interruptId: 'second', + interruptedRunId: runId, + generation: 1, + responseSchemaHash: 'none', + }), + ], + }, + } + return + } + yield { + type: EventType.RUN_FINISHED, + threadId: 'thread-1', + runId, + timestamp: Date.now(), + outcome: { type: 'success' }, + } + }, + } + const client = new ChatClient({ connection, threadId: 'thread-1' }) + + await client.sendMessage('start') + const visited: Array = [] + client.resolveInterrupts((interrupt) => { + visited.push(interrupt.id) + interrupt.cancel() + return undefined + }) + + await vi.waitFor(() => expect(contexts).toHaveLength(2)) + expect(visited).toEqual(['first', 'second']) + expect(contexts[1]?.resume).toEqual([ + { interruptId: 'first', status: 'cancelled' }, + { interruptId: 'second', status: 'cancelled' }, + ]) + }) + + it('writes a raw V2 snapshot with authoritative recovery state', async () => { + const writes: Array = [] + const persistence: ChatServerPersistence = { + getItem: () => undefined, + setItem: (_id, value) => { + if (value.schemaVersion === 2) writes.push(value) + }, + removeItem: () => undefined, + } + const binding: InterruptBinding = { + kind: 'generic', + interruptId: 'generic-1', + interruptedRunId: 'run-1', + generation: 4, + responseSchemaHash: 'none', + } + const connection: ConnectConnectionAdapter = { + async *connect(_messages, _data, _signal, context) { + binding.interruptedRunId = context?.runId ?? 'run-1' + yield { + type: EventType.RUN_FINISHED, + threadId: context?.threadId ?? 'thread-1', + runId: binding.interruptedRunId, + timestamp: Date.now(), + outcome: { + type: 'interrupt', + interrupts: [descriptor(binding)], + }, + } + }, + } + const client = new ChatClient({ + connection, + threadId: 'thread-1', + persistence: { server: persistence }, + }) + + await client.sendMessage('start') + + expect(writes.at(-1)).toEqual({ + schemaVersion: 2, + resumeState: { + threadId: 'thread-1', + runId: binding.interruptedRunId, + }, + pendingInterrupts: [descriptor(binding)], + interruptState: { + recoveryState: { + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: binding.interruptedRunId, + generation: 4, + pendingInterrupts: [descriptor(binding)], + }, + drafts: [], + }, + }) + expect(JSON.parse(JSON.stringify(writes.at(-1)))).toEqual(writes.at(-1)) + }) + + it('hydrates V2 drafts immediately before reconciling authoritative state', async () => { + const first = genericDescriptor('first') + const second = genericDescriptor('second') + const connect = vi.fn(async function* () {}) + let releaseRecovery: ((state: InterruptRecoveryStateV1) => void) | undefined + const recovery = new Promise((resolve) => { + releaseRecovery = resolve + }) + const loadInterruptState = vi.fn(() => recovery) + const client = new ChatClient({ + connection: { connect, loadInterruptState }, + initialResumeSnapshot: { + schemaVersion: 2, + resumeState: { threadId: 'thread-1', runId: 'run-1' }, + pendingInterrupts: [first, second], + interruptState: { + recoveryState: { + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + pendingInterrupts: [first, second], + }, + drafts: [ + { + interruptId: 'first', + response: { + interruptId: 'first', + status: 'resolved', + payload: { answer: 42 }, + }, + status: 'staged', + }, + ], + }, + }, + }) + + expect(client.getInterrupts().map((item) => item.status)).toEqual([ + 'staged', + 'pending', + ]) + releaseRecovery?.({ + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + pendingInterrupts: [first, second], + }) + await settle() + await settle() + expect(client.getInterrupts()).toHaveLength(2) + expect(client.getInterrupts().map((item) => item.status)).toEqual([ + 'staged', + 'pending', + ]) + expect(loadInterruptState).toHaveBeenCalledWith({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + knownGeneration: 1, + }) + expect(connect).not.toHaveBeenCalled() + }) + + it('hydrates V2 fallback descriptors when recovery is unavailable', () => { + const fallback = genericDescriptor('fallback') + const malformed = JSON.parse( + JSON.stringify({ + schemaVersion: 2, + resumeState: { threadId: 'thread-1', runId: 'run-1' }, + pendingInterrupts: [fallback], + interruptState: { + recoveryState: { + schemaVersion: 1, + state: 'pending', + threadId: 'other-thread', + interruptedRunId: 'run-1', + generation: 99, + pendingInterrupts: [], + }, + drafts: [], + }, + }), + ) + const client = new ChatClient({ + connection: { async *connect() {} }, + initialResumeSnapshot: malformed, + }) + + expect(client.getResumeState()).toEqual({ + threadId: 'thread-1', + runId: 'run-1', + }) + expect(client.getInterrupts().map((item) => item.id)).toEqual(['fallback']) + expect(client.getInterruptState().interruptErrors).toEqual([]) + }) + + it.each([ + ['null recovery state', { recoveryState: null, drafts: [] }], + ['array recovery state', { recoveryState: [], drafts: [] }], + [ + 'invalid drafts', + { + recoveryState: { + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + pendingInterrupts: [], + }, + drafts: { interruptId: 'not-an-array' }, + }, + ], + ])( + 'hydrates fallback descriptors for malformed V2 %s without losing resume state', + (_label, interruptState) => { + const malformed = JSON.parse( + JSON.stringify({ + schemaVersion: 2, + resumeState: { threadId: 'thread-1', runId: 'run-1' }, + pendingInterrupts: [genericDescriptor('fallback')], + interruptState, + }), + ) + + let client: ChatClient | undefined + expect(() => { + client = new ChatClient({ + connection: { async *connect() {} }, + initialResumeSnapshot: malformed, + }) + }).not.toThrow() + expect(client?.getResumeState()).toEqual({ + threadId: 'thread-1', + runId: 'run-1', + }) + expect(client?.getInterrupts().map((item) => item.id)).toEqual([ + 'fallback', + ]) + expect(client?.getInterruptState().interruptErrors).toEqual([]) + }, + ) + + it('hydrates V1 descriptors before explicit recovery replaces them', async () => { + const persisted = genericDescriptor('persisted') + const authoritative = genericDescriptor('authoritative') + let releaseRecovery: ((state: InterruptRecoveryStateV1) => void) | undefined + const recovery = new Promise((resolve) => { + releaseRecovery = resolve + }) + const loadInterruptState = vi.fn(() => recovery) + const client = new ChatClient({ + connection: { async *connect() {}, loadInterruptState }, + initialResumeSnapshot: { + schemaVersion: 1, + resumeState: { threadId: 'thread-1', runId: 'run-1' }, + pendingInterrupts: [persisted], + }, + }) + + expect(client.getInterrupts().map((item) => item.id)).toEqual(['persisted']) + releaseRecovery?.({ + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + pendingInterrupts: [authoritative], + }) + await vi.waitFor(() => + expect(client.getInterrupts()[0]?.id).toBe('authoritative'), + ) + expect(loadInterruptState).toHaveBeenCalledWith({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + knownGeneration: 1, + }) + }) + + it('recovers a committed winner by joining it read-only', async () => { + const contexts: Array = [] + const joined: Array = [] + const binding = genericDescriptor('generic-1') + let interruptedRunId = 'run-1' + let calls = 0 + const connection: ConnectConnectionAdapter = { + async *connect(_messages, _data, _signal, context) { + contexts.push(context) + calls++ + if (calls === 1) { + interruptedRunId = context?.runId ?? interruptedRunId + const metadata = binding.metadata?.['tanstack:interruptBinding'] + if (metadata && typeof metadata === 'object') { + metadata.interruptedRunId = interruptedRunId + } + yield { + type: EventType.RUN_FINISHED, + threadId: 'thread-1', + runId: interruptedRunId, + timestamp: Date.now(), + outcome: { type: 'interrupt', interrupts: [binding] }, + } + return + } + yield { + type: EventType.RUN_ERROR, + threadId: 'thread-1', + runId: context?.runId ?? 'loser-run', + timestamp: Date.now(), + message: 'conflict', + code: 'conflict', + 'tanstack:interruptErrors': [ + { + scope: 'batch', + code: 'conflict', + message: 'another continuation won', + source: 'server', + retryable: false, + interruptIds: ['generic-1'], + threadId: 'thread-1', + interruptedRunId, + generation: 1, + }, + ], + 'tanstack:interruptRecovery': { + schemaVersion: 1, + state: 'committed', + threadId: 'thread-1', + interruptedRunId, + generation: 1, + pendingInterrupts: [], + committed: { + fingerprint: 'winner', + continuationRunId: 'winner-run', + committedAt: '2026-07-13T00:00:00.000Z', + }, + }, + } + }, + async *joinRun(runId) { + joined.push(runId) + yield { + type: EventType.RUN_FINISHED, + threadId: 'thread-1', + runId, + timestamp: Date.now(), + outcome: { type: 'success' }, + } + }, + } + const client = new ChatClient({ connection, threadId: 'thread-1' }) + + await client.sendMessage('start') + client.getInterrupts()[0]?.cancel() + + await vi.waitFor(() => expect(joined).toEqual(['winner-run'])) + expect(contexts).toHaveLength(2) + expect(client.getInterrupts()).toHaveLength(0) + expect(client.getInterruptState().interruptErrors).toMatchObject([ + { code: 'conflict', source: 'server' }, + ]) + }) + + it('replays a persisted committed winner through an explicit continuation loader', async () => { + const connect = vi.fn(async function* () {}) + const continuationLoader = vi.fn(async function* (runId: string) { + yield { + type: EventType.RUN_FINISHED as const, + threadId: 'thread-1', + runId, + timestamp: Date.now(), + outcome: { type: 'success' as const }, + } + }) + const client = new ChatClient({ + connection: { connect }, + continuationLoader, + initialResumeSnapshot: { + schemaVersion: 2, + resumeState: { threadId: 'thread-1', runId: 'interrupted-run' }, + pendingInterrupts: [], + interruptState: { + recoveryState: { + schemaVersion: 1, + state: 'committed', + threadId: 'thread-1', + interruptedRunId: 'interrupted-run', + generation: 1, + pendingInterrupts: [], + committed: { + fingerprint: 'winner', + continuationRunId: 'winner-run', + committedAt: '2026-07-13T00:00:00.000Z', + }, + }, + drafts: [], + }, + }, + }) + + await vi.waitFor(() => + expect(continuationLoader).toHaveBeenCalledWith('winner-run'), + ) + await vi.waitFor(() => expect(client.getResumeState()).toBeNull()) + expect(connect).not.toHaveBeenCalled() + }) + + it('reports recovery-unavailable when a committed winner cannot be joined', async () => { + const client = new ChatClient({ + connection: { async *connect() {} }, + initialResumeSnapshot: { + schemaVersion: 2, + resumeState: { threadId: 'thread-1', runId: 'interrupted-run' }, + pendingInterrupts: [], + interruptState: { + recoveryState: { + schemaVersion: 1, + state: 'committed', + threadId: 'thread-1', + interruptedRunId: 'interrupted-run', + generation: 1, + pendingInterrupts: [], + committed: { + fingerprint: 'winner', + continuationRunId: 'winner-run', + committedAt: '2026-07-13T00:00:00.000Z', + }, + }, + drafts: [], + }, + }, + }) + + await vi.waitFor(() => + expect(client.getInterruptState().interruptErrors[0]?.code).toBe( + 'recovery-unavailable', + ), + ) + }) + + it('loads authoritative pending state only through the explicit recovery operation', async () => { + const original = genericDescriptor('generic-1') + const replacement = descriptor({ + kind: 'generic', + interruptId: 'generic-2', + interruptedRunId: 'run-1', + generation: 2, + responseSchemaHash: 'none', + }) + let interruptedRunId = 'run-1' + let calls = 0 + const loadInterruptState = vi.fn(async () => ({ + schemaVersion: 1 as const, + state: 'pending' as const, + threadId: 'thread-1', + interruptedRunId, + generation: 2, + pendingInterrupts: [replacement], + })) + const connection: ConnectConnectionAdapter = { + async *connect(_messages, _data, _signal, context) { + calls++ + if (calls === 1) { + interruptedRunId = context?.runId ?? interruptedRunId + const metadata = original.metadata?.['tanstack:interruptBinding'] + if (metadata && typeof metadata === 'object') { + metadata.interruptedRunId = interruptedRunId + } + const nextMetadata = + replacement.metadata?.['tanstack:interruptBinding'] + if (nextMetadata && typeof nextMetadata === 'object') { + nextMetadata.interruptedRunId = interruptedRunId + } + yield { + type: EventType.RUN_FINISHED, + threadId: 'thread-1', + runId: interruptedRunId, + timestamp: Date.now(), + outcome: { type: 'interrupt', interrupts: [original] }, + } + return + } + yield { + type: EventType.RUN_ERROR, + threadId: 'thread-1', + runId: context?.runId ?? 'loser-run', + timestamp: Date.now(), + message: 'stale', + code: 'stale', + 'tanstack:interruptErrors': [ + { + scope: 'batch', + code: 'stale', + message: 'stale generation', + source: 'server', + retryable: false, + interruptIds: ['generic-1'], + threadId: 'thread-1', + interruptedRunId, + generation: 1, + }, + ], + } + }, + loadInterruptState, + } + const client = new ChatClient({ connection, threadId: 'thread-1' }) + + await client.sendMessage('start') + client.getInterrupts()[0]?.cancel() + + await vi.waitFor(() => + expect(client.getInterrupts()[0]).toMatchObject({ + id: 'generic-2', + generation: 2, + }), + ) + expect(loadInterruptState).toHaveBeenCalledWith({ + threadId: 'thread-1', + interruptedRunId, + knownGeneration: 1, + }) + }) +}) diff --git a/packages/ai-client/tests/chat-client-resume.test.ts b/packages/ai-client/tests/chat-client-resume.test.ts index f2e7e4839..7ecebca85 100644 --- a/packages/ai-client/tests/chat-client-resume.test.ts +++ b/packages/ai-client/tests/chat-client-resume.test.ts @@ -210,7 +210,8 @@ describe('ChatClient resume', () => { await client.resumeInterrupts(resumeItems) expect(contexts[1]?.threadId).toBe(resumeState?.threadId) - expect(contexts[1]?.runId).toBe(resumeState?.runId) + expect(contexts[1]?.runId).not.toBe(resumeState?.runId) + expect(contexts[1]?.parentRunId).toBe(resumeState?.runId) expect(contexts[1]?.resume).toEqual(resumeItems) expect(client.getPendingInterrupts()).toEqual([]) expect(client.getResumeState()).toBeNull() @@ -598,7 +599,11 @@ describe('ChatClient resume', () => { id: 'approval-1', reason: 'approval_required', toolCallId: 'tool-1', - metadata: { kind: 'approval' }, + metadata: { + kind: 'approval', + toolName: 'lookup', + input: { query: 'first' }, + }, }, ], }, @@ -650,13 +655,21 @@ describe('ChatClient resume', () => { id: 'approval-1', reason: 'approval_required', toolCallId: 'tool-1', - metadata: { kind: 'approval' }, + metadata: { + kind: 'approval', + toolName: 'lookup', + input: { query: 'first' }, + }, }, { id: 'approval-2', reason: 'approval_required', toolCallId: 'tool-2', - metadata: { kind: 'approval' }, + metadata: { + kind: 'approval', + toolName: 'lookup', + input: { query: 'second' }, + }, }, ], }, @@ -682,7 +695,7 @@ describe('ChatClient resume', () => { }, { interruptId: 'approval-2', - status: 'cancelled', + status: 'resolved', payload: { approved: false }, }, ]) @@ -712,7 +725,11 @@ describe('ChatClient resume', () => { id: 'approval-1', reason: 'approval_required', toolCallId: 'tool-1', - metadata: { kind: 'approval' }, + metadata: { + kind: 'approval', + toolName: 'lookup', + input: { query: 'restored' }, + }, }, ], }, @@ -729,7 +746,8 @@ describe('ChatClient resume', () => { await client.addToolApprovalResponse({ id: 'approval-1', approved: true }) expect(contexts[0]?.threadId).toBe('thread-1') - expect(contexts[0]?.runId).toBe('run-1') + expect(contexts[0]?.runId).not.toBe('run-1') + expect(contexts[0]?.parentRunId).toBe('run-1') expect(contexts[0]?.resume).toEqual([ { interruptId: 'approval-1', @@ -805,8 +823,30 @@ describe('ChatClient resume', () => { await client.sendMessage('hi') expect(persistence.setItem).toHaveBeenLastCalledWith('thread-1', { + schemaVersion: 2, resumeState: client.getResumeState(), - pendingInterrupts: client.getPendingInterrupts(), + pendingInterrupts: [ + expect.objectContaining({ + id: 'interrupt-1', + reason: 'client_tool_input', + }), + ], + interruptState: { + recoveryState: { + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: client.getResumeState()?.runId, + generation: 0, + pendingInterrupts: [ + expect.objectContaining({ + id: 'interrupt-1', + reason: 'client_tool_input', + }), + ], + }, + drafts: [], + }, }) expect(onResumeStateChange).toHaveBeenCalled() }) @@ -1300,6 +1340,11 @@ describe('ChatClient resume', () => { id: 'interrupt-tool-1', reason: 'client_tool_input', toolCallId: 'tool-call-1', + metadata: { + kind: 'client_tool', + toolName: 'lookup', + input: { query: 'first' }, + }, }, ], }, @@ -1318,7 +1363,8 @@ describe('ChatClient resume', () => { }) expect(contexts[1]?.threadId).toBe(resumeState?.threadId) - expect(contexts[1]?.runId).toBe(resumeState?.runId) + expect(contexts[1]?.runId).not.toBe(resumeState?.runId) + expect(contexts[1]?.parentRunId).toBe(resumeState?.runId) expect(contexts[1]?.resume).toEqual([ { interruptId: 'interrupt-tool-1', @@ -1350,11 +1396,21 @@ describe('ChatClient resume', () => { id: 'interrupt-tool-1', reason: 'client_tool_input', toolCallId: 'tool-call-1', + metadata: { + kind: 'client_tool', + toolName: 'lookup', + input: { query: 'first' }, + }, }, { id: 'interrupt-tool-2', reason: 'client_tool_input', toolCallId: 'tool-call-2', + metadata: { + kind: 'client_tool', + toolName: 'lookup', + input: { query: 'second' }, + }, }, ], }, diff --git a/packages/ai-client/tests/connection-adapters.test.ts b/packages/ai-client/tests/connection-adapters.test.ts index 4c98d25d1..78b13a041 100644 --- a/packages/ai-client/tests/connection-adapters.test.ts +++ b/packages/ai-client/tests/connection-adapters.test.ts @@ -440,6 +440,7 @@ describe('connection-adapters', () => { { threadId: 'thread-1', runId: 'run-1', + parentRunId: 'interrupted-run', resume: [ { interruptId: 'interrupt-1', @@ -1230,6 +1231,60 @@ describe('connection-adapters', () => { expect(received[0].error?.message).toBe('already failed') } }) + + it('preserves read-only join and explicit interrupt recovery operations', async () => { + const recovery = { + schemaVersion: 1 as const, + state: 'pending' as const, + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 2, + pendingInterrupts: [], + } + const loadInterruptState = vi.fn(async () => recovery) + const base = { + async *connect() { + yield { + type: EventType.RUN_FINISHED, + threadId: 'thread-1', + runId: 'run-new', + timestamp: Date.now(), + } as const + }, + async *joinRun(runId: string) { + yield { + type: EventType.RUN_FINISHED, + threadId: 'thread-1', + runId, + timestamp: Date.now(), + } as const + }, + loadInterruptState, + } + const adapter = normalizeConnectionAdapter(base) + const abortController = new AbortController() + const received = (async () => { + for await (const chunk of adapter.subscribe(abortController.signal)) { + abortController.abort() + return chunk + } + return undefined + })() + + await adapter.joinRun?.('winner-run') + expect(await received).toMatchObject({ + type: EventType.RUN_FINISHED, + runId: 'winner-run', + }) + await expect( + adapter.loadInterruptState?.({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + knownGeneration: 1, + }), + ).resolves.toBe(recovery) + expect(loadInterruptState).toHaveBeenCalledTimes(1) + }) }) describe('rpcStream', () => { @@ -1354,6 +1409,7 @@ describe('connection-adapters', () => { { threadId: 'thread-1', runId: 'run-1', + parentRunId: 'interrupted-run', resume: [ { interruptId: 'interrupt-1', @@ -1374,6 +1430,7 @@ describe('connection-adapters', () => { payload: { approved: true }, }, ]) + expect(input.parentRunId).toBe('interrupted-run') }) }) }) diff --git a/packages/ai-client/tests/interrupt-recovery-adapters.test.ts b/packages/ai-client/tests/interrupt-recovery-adapters.test.ts new file mode 100644 index 000000000..3e5d846fc --- /dev/null +++ b/packages/ai-client/tests/interrupt-recovery-adapters.test.ts @@ -0,0 +1,111 @@ +import { describe, expect, it, vi } from 'vitest' +import { EventType } from '@tanstack/ai/client' +import { + createInterruptContinuationLoader, + createInterruptStateFetcher, + fetchServerSentEvents, +} from '../src/connection-adapters' +import type { InterruptRecoveryStateV1 } from '@tanstack/ai/client' + +const recoveryState: InterruptRecoveryStateV1 = { + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 2, + pendingInterrupts: [], +} + +function continuationResponse(): Response { + return new Response( + `data: ${JSON.stringify({ + type: EventType.RUN_FINISHED, + threadId: 'thread-1', + runId: 'winner-run', + timestamp: 1, + outcome: { type: 'success' }, + })}\n\n`, + { headers: { 'content-type': 'text/event-stream' } }, + ) +} + +describe('interrupt recovery fetch adapters', () => { + it('fetches and strictly parses authoritative state from an explicit URL', async () => { + const fetchClient = vi.fn(async () => + Response.json(recoveryState), + ) + const fetchState = createInterruptStateFetcher('/api/recovery?tenant=one', { + fetchClient, + headers: { authorization: 'Bearer test' }, + }) + + await expect( + fetchState({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + knownGeneration: 2, + }), + ).resolves.toEqual(recoveryState) + const [url, init] = fetchClient.mock.calls[0] ?? [] + expect(String(url)).toBe( + '/api/recovery?tenant=one&threadId=thread-1&interruptedRunId=run-1&knownGeneration=2', + ) + expect(init?.method).toBe('POST') + expect(new Headers(init?.headers).get('authorization')).toBe('Bearer test') + }) + + it('rejects malformed recovery JSON instead of trusting the response shape', async () => { + const fetchState = createInterruptStateFetcher('/api/recovery', { + fetchClient: vi.fn(async () => + Response.json({ ...recoveryState, pendingInterrupts: 'not-an-array' }), + ), + }) + + await expect( + fetchState({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + knownGeneration: 2, + }), + ).rejects.toThrow('Invalid interrupt recovery response') + }) + + it('wires explicit recovery and continuation operations without posting a new run', async () => { + const recoveryFetch = vi.fn(async () => + Response.json(recoveryState), + ) + const continuationFetch = vi.fn(async () => + continuationResponse(), + ) + const chatFetch = vi.fn() + const interruptStateFetcher = createInterruptStateFetcher('/api/recovery', { + fetchClient: recoveryFetch, + }) + const continuationLoader = createInterruptContinuationLoader( + '/api/continuation', + { fetchClient: continuationFetch }, + ) + const adapter = fetchServerSentEvents('/api/chat', { + fetchClient: chatFetch, + interruptStateFetcher, + continuationLoader, + }) + + await expect( + adapter.loadInterruptState?.({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + knownGeneration: 2, + }), + ).resolves.toEqual(recoveryState) + const chunks = [] + for await (const chunk of adapter.joinRun('winner-run')) chunks.push(chunk) + + expect(chunks).toHaveLength(1) + expect(continuationFetch).toHaveBeenCalledWith( + '/api/continuation?offset=-1&runId=winner-run', + expect.objectContaining({ method: 'GET' }), + ) + expect(chatFetch).not.toHaveBeenCalled() + }) +}) diff --git a/packages/ai-client/tests/interrupts-types.test-d.ts b/packages/ai-client/tests/interrupts-types.test-d.ts new file mode 100644 index 000000000..85f05bd54 --- /dev/null +++ b/packages/ai-client/tests/interrupts-types.test-d.ts @@ -0,0 +1,217 @@ +import { expectTypeOf } from 'vitest' +import { toolDefinition } from '@tanstack/ai/client' +import { z } from 'zod' +import { + createInterruptContinuationLoader, + createInterruptStateFetcher, + fetchServerSentEvents, +} from '../src/index' +import type { JSONSchema, ServerTool } from '@tanstack/ai' +import type { + ClientTool, + InterruptBinding, + Interrupt as WireInterrupt, + InputSchemaOf, + ItemInterruptError, + NoSchema, + OutputSchemaOf, +} from '@tanstack/ai/client' +import type { + BoundInterrupts, + ChatClient, + ChatInterrupt, + ClientToolExecutionInterrupt, + ToolApprovalInterrupt, +} from '../src/index' + +const transfer = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema: z.object({ cents: z.number() }), + outputSchema: z.object({ receipt: z.string() }), + approvalSchema: { + approve: z.object({ note: z.string() }), + reject: z.object({ reason: z.string() }), + }, +}).client() + +const confirm = toolDefinition({ + name: 'confirm', + description: 'Approval without custom schemas', + needsApproval: true, +}).client() + +const lookup = toolDefinition({ + name: 'lookup', + description: 'Client lookup', + outputSchema: z.object({ accountId: z.string() }), +}).client() + +const raw = toolDefinition({ + name: 'raw', + description: 'Raw JSON Schema remains unknown', + outputSchema: { type: 'object' }, +}).client() + +type Tools = readonly [ + typeof transfer, + typeof confirm, + typeof lookup, + typeof raw, +] +type Interrupt = ChatInterrupt + +expectTypeOf>().toEqualTypeOf() +expectTypeOf>().toEqualTypeOf() +expectTypeOf>().toEqualTypeOf() +expectTypeOf>().toEqualTypeOf() + +const manuallyAnnotatedClientTool: ClientTool = { + __toolSide: 'client', + name: 'manual', + description: 'Manually annotated schema-less tool', +} +expectTypeOf(manuallyAnnotatedClientTool.inputSchema).toEqualTypeOf() + +const manualJsonSchema = { + type: 'object', + properties: {}, +} satisfies JSONSchema +const manuallyAnnotatedServerTool: ServerTool = { + __toolSide: 'server', + name: 'manual-server', + description: 'Manually annotated tool with an explicit schema generic', + inputSchema: manualJsonSchema, + execute: async () => undefined, +} +expectTypeOf(manuallyAnnotatedServerTool.inputSchema).toEqualTypeOf< + typeof manualJsonSchema | undefined +>() + +const publicRecoveryConnection = fetchServerSentEvents('/api/chat', { + interruptStateFetcher: createInterruptStateFetcher('/api/recovery'), + continuationLoader: createInterruptContinuationLoader('/api/continuation'), +}) +expectTypeOf(publicRecoveryConnection.loadInterruptState).not.toBeUndefined() + +declare const approvalBinding: Extract< + InterruptBinding, + { kind: 'tool-approval' } +> +const canonicalApprovalDescriptor = { + id: approvalBinding.interruptId, + reason: 'tool_call' as const, + toolCallId: approvalBinding.toolCallId, + metadata: { 'tanstack:interruptBinding': approvalBinding }, +} satisfies WireInterrupt +expectTypeOf(canonicalApprovalDescriptor.reason).toEqualTypeOf<'tool_call'>() + +declare const clientToolBinding: Extract< + InterruptBinding, + { kind: 'client-tool-execution' } +> +const canonicalClientToolDescriptor = { + id: clientToolBinding.interruptId, + reason: 'tanstack:client_tool_execution' as const, + toolCallId: clientToolBinding.toolCallId, + metadata: { 'tanstack:interruptBinding': clientToolBinding }, +} satisfies WireInterrupt +expectTypeOf( + canonicalClientToolDescriptor.reason, +).toEqualTypeOf<'tanstack:client_tool_execution'>() + +declare const transferApproval: Extract< + Interrupt, + { kind: 'tool-approval'; toolName: 'transfer' } +> +transferApproval.resolveInterrupt(true, { + editedArgs: { cents: 100 }, + payload: { note: 'approved' }, +}) +transferApproval.resolveInterrupt(false, { payload: { reason: 'declined' } }) +// @ts-expect-error rejection never permits edited arguments +transferApproval.resolveInterrupt(false, { editedArgs: { cents: 1 } }) +// @ts-expect-error approve payload is inferred from the approve branch +transferApproval.resolveInterrupt(true, { payload: { reason: 'wrong branch' } }) + +declare const confirmation: Extract< + Interrupt, + { kind: 'tool-approval'; toolName: 'confirm' } +> +confirmation.resolveInterrupt(true) +confirmation.resolveInterrupt(false) +// @ts-expect-error omitted input schema forbids edits +confirmation.resolveInterrupt(true, { editedArgs: { value: 1 } }) +// @ts-expect-error omitted approval branch forbids payload +confirmation.resolveInterrupt(false, { payload: { reason: 'no schema' } }) + +declare const lookupExecution: Extract< + Interrupt, + { kind: 'client-tool-execution'; toolName: 'lookup' } +> +lookupExecution.resolveInterrupt({ accountId: 'account-1' }) +// @ts-expect-error client tool output is inferred +lookupExecution.resolveInterrupt({ account: 'account-1' }) + +const structuralLookup = toolDefinition({ + name: 'structuralLookup', + description: 'Client lookup with a structural Standard Schema', + outputSchema: { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { accountId: '' }, + output: { accountId: '' }, + }, + validate: () => ({ value: { accountId: '' } }), + }, + }, +}).client() + +declare const structuralLookupExecution: Extract< + ChatInterrupt, + { kind: 'client-tool-execution'; toolName: 'structuralLookup' } +> +structuralLookupExecution.resolveInterrupt({ accountId: 'account-1' }) +// @ts-expect-error structural Standard Schema output remains inferred +structuralLookupExecution.resolveInterrupt({ account: 'account-1' }) + +declare const rawExecution: Extract< + Interrupt, + { kind: 'client-tool-execution'; toolName: 'raw' } +> +expectTypeOf(rawExecution.errors).toEqualTypeOf< + ReadonlyArray +>() +expectTypeOf(rawExecution.error).toEqualTypeOf() +expectTypeOf(rawExecution.resolveInterrupt) + .parameter(0) + .toEqualTypeOf() + +expectTypeOf>().toEqualTypeOf< + readonly ChatInterrupt[] +>() + +declare const client: ChatClient +client.updateOptions({ + onInterruptStateChange: (state) => { + expectTypeOf(state.interrupts).toEqualTypeOf(state.pendingInterrupts) + }, +}) +client.resolveInterrupts((interrupt) => { + interrupt.cancel() + return undefined +}) +client.resumeInterruptsUnsafe([ + { interruptId: 'generic-1', status: 'cancelled' }, +]) +// @ts-expect-error the low-level recovery escape hatch uses the locked name +client.unsafeResumeInterrupts([]) +expectTypeOf>().toMatchTypeOf< + ToolApprovalInterrupt +>() +expectTypeOf< + Extract +>().toMatchTypeOf>() diff --git a/packages/ai-durable-stream/src/durable-stream.ts b/packages/ai-durable-stream/src/durable-stream.ts index 0c8621a78..087882bcb 100644 --- a/packages/ai-durable-stream/src/durable-stream.ts +++ b/packages/ai-durable-stream/src/durable-stream.ts @@ -152,25 +152,23 @@ async function* readLines( const reader = body.getReader() const decoder = new TextDecoder() let buffer = '' - let completed = false - let cancelled = false - let readFailed = false + let shouldCancelReader = true try { for (;;) { let result: ReadableStreamReadResult | typeof READ_ABORTED try { result = await readWithAbort(reader, signal) } catch (error) { - readFailed = true + shouldCancelReader = false throw new ResponseBodyReadFailure(error) } if (result === READ_ABORTED) { - cancelled = true + shouldCancelReader = false await reader.cancel(signal?.reason) return } if (result.done) { - completed = true + shouldCancelReader = false break } buffer += decoder.decode(result.value, { stream: true }) @@ -186,7 +184,7 @@ async function* readLines( } } finally { try { - if (!completed && !cancelled && !readFailed) await reader.cancel() + if (shouldCancelReader) await reader.cancel() } finally { reader.releaseLock() } diff --git a/packages/ai-gemini/tests/text-interactions-adapter.test.ts b/packages/ai-gemini/tests/text-interactions-adapter.test.ts index c043efb63..62d57f420 100644 --- a/packages/ai-gemini/tests/text-interactions-adapter.test.ts +++ b/packages/ai-gemini/tests/text-interactions-adapter.test.ts @@ -1,7 +1,16 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' import { z } from 'zod' -import { chat } from '@tanstack/ai' -import type { StreamChunk, Tool } from '@tanstack/ai' +import { + InterruptPersistenceCapability, + chat, + defineChatMiddleware, + provideInterruptPersistence, +} from '@tanstack/ai' +import type { + InterruptPersistenceGateway, + StreamChunk, + Tool, +} from '@tanstack/ai' import { GeminiTextInteractionsAdapter } from '../src/experimental/text-interactions/adapter' import type { GeminiTextInteractionsProviderOptions } from '../src/experimental/text-interactions/adapter' import type { @@ -53,6 +62,43 @@ const collectChunks = async (stream: AsyncIterable) => { return chunks } +const testInterruptGateway: InterruptPersistenceGateway = { + openInterruptBatch: async (input) => ({ + generation: 1, + descriptors: input.descriptors, + }), + commitInterruptResolutions: async (input) => ({ + status: 'committed', + continuationRunId: input.continuationRunId, + }), + getInterruptRecoveryState: async (input) => ({ + schemaVersion: 1, + state: 'missing', + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + generation: input.knownGeneration, + pendingInterrupts: [], + }), +} + +const testInterruptPersistence = defineChatMiddleware({ + name: 'gemini-test-interrupt-persistence', + provides: [InterruptPersistenceCapability], + setup(ctx) { + provideInterruptPersistence(ctx, testInterruptGateway) + }, + async onChunk(_ctx, chunk) { + if (chunk.type === 'RUN_FINISHED' && chunk.outcome?.type === 'interrupt') { + await testInterruptGateway.openInterruptBatch({ + threadId: chunk.threadId, + interruptedRunId: chunk.runId, + descriptors: chunk.outcome.interrupts, + bindings: [], + }) + } + }, +}) + describe('GeminiTextInteractionsAdapter', () => { beforeEach(() => { // `mockReset()` (vs `clearAllMocks`) also clears `mockResolvedValue` @@ -416,6 +462,7 @@ describe('GeminiTextInteractionsAdapter', () => { adapter, messages: [{ role: 'user', content: 'Weather in Madrid?' }], tools: [weatherTool], + middleware: [testInterruptPersistence], }), ) @@ -428,19 +475,20 @@ describe('GeminiTextInteractionsAdapter', () => { }), ]) - const startEvent = chunks.find((c) => c.type === 'TOOL_CALL_START') as any - expect(startEvent).toBeDefined() - expect(startEvent.toolCallId).toBe('call_1') - expect(startEvent.toolName).toBe('lookup_weather') + const startEvent = chunks.find((c) => c.type === 'TOOL_CALL_START') + expect(startEvent).toMatchObject({ + toolCallId: 'call_1', + toolName: 'lookup_weather', + }) - const argsEvent = chunks.find((c) => c.type === 'TOOL_CALL_ARGS') as any - expect(argsEvent.args).toBe('{"location":"Madrid"}') + const argsEvent = chunks.find((c) => c.type === 'TOOL_CALL_ARGS') + expect(argsEvent).toMatchObject({ args: '{"location":"Madrid"}' }) - const endEvent = chunks.find((c) => c.type === 'TOOL_CALL_END') as any - expect(endEvent.input).toEqual({ location: 'Madrid' }) + const endEvent = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(endEvent).toMatchObject({ input: { location: 'Madrid' } }) - const finished = chunks.find((c) => c.type === 'RUN_FINISHED') as any - expect(finished.finishReason).toBe('tool_calls') + const finished = chunks.find((c) => c.type === 'RUN_FINISHED') + expect(finished).toMatchObject({ finishReason: 'tool_calls' }) }) it('accumulates multi-fragment arguments_delta into the final tool input', async () => { @@ -492,24 +540,27 @@ describe('GeminiTextInteractionsAdapter', () => { adapter, messages: [{ role: 'user', content: 'Weather in Berlin?' }], tools: [{ name: 'lookup_weather', description: 'Return the weather' }], + middleware: [testInterruptPersistence], }), ) - const startEvent = chunks.find((c) => c.type === 'TOOL_CALL_START') as any - expect(startEvent.toolCallId).toBe('call_frag') - expect(startEvent.toolName).toBe('lookup_weather') + const startEvent = chunks.find((c) => c.type === 'TOOL_CALL_START') + expect(startEvent).toMatchObject({ + toolCallId: 'call_frag', + toolName: 'lookup_weather', + }) // The last TOOL_CALL_ARGS event carries the fully-accumulated buffer. - const argsEvents = chunks.filter( - (c) => c.type === 'TOOL_CALL_ARGS', - ) as Array + const argsEvents = chunks.filter((c) => c.type === 'TOOL_CALL_ARGS') expect(argsEvents.length).toBeGreaterThan(1) - expect(argsEvents.at(-1)!.args).toBe( - '{"location":"Berlin","unit":"celsius"}', - ) + expect(argsEvents.at(-1)).toMatchObject({ + args: '{"location":"Berlin","unit":"celsius"}', + }) - const endEvent = chunks.find((c) => c.type === 'TOOL_CALL_END') as any - expect(endEvent.input).toEqual({ location: 'Berlin', unit: 'celsius' }) + const endEvent = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(endEvent).toMatchObject({ + input: { location: 'Berlin', unit: 'celsius' }, + }) }) it('preserves the last good args when the arguments stream truncates mid-fragment', async () => { @@ -557,13 +608,14 @@ describe('GeminiTextInteractionsAdapter', () => { adapter, messages: [{ role: 'user', content: 'Weather in London?' }], tools: [{ name: 'lookup_weather', description: 'Return the weather' }], + middleware: [testInterruptPersistence], }), ) // No parse-failure RUN_ERROR, and the completed key survives truncation. expect(chunks.find((c) => c.type === 'RUN_ERROR')).toBeUndefined() - const endEvent = chunks.find((c) => c.type === 'TOOL_CALL_END') as any - expect(endEvent.input).toEqual({ city: 'London' }) + const endEvent = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(endEvent).toMatchObject({ input: { city: 'London' } }) }) it('translates thought_summary deltas into REASONING events', async () => { diff --git a/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts b/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts index 84c74522e..f6f458f15 100644 --- a/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts +++ b/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts @@ -1,11 +1,21 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' -import { EventType, chat } from '@tanstack/ai' +import { + EventType, + InterruptPersistenceCapability, + chat, + defineChatMiddleware, + provideInterruptPersistence, +} from '@tanstack/ai' import { resolveDebugOption } from '@tanstack/ai/adapter-internals' import { ResponsesRequest$outboundSchema } from '@openrouter/sdk/models' import { createOpenRouterResponsesText } from '../src/adapters/responses-text' import { webSearchTool } from '../src/tools/web-search-tool' import { webFetchTool } from '../src/tools/web-fetch-tool' -import type { StreamChunk, Tool } from '@tanstack/ai' +import type { + InterruptPersistenceGateway, + StreamChunk, + Tool, +} from '@tanstack/ai' const testLogger = resolveDebugOption(false) let mockSend: any @@ -34,6 +44,43 @@ const weatherTool: Tool = { description: 'Return the forecast for a location', } +const testInterruptGateway: InterruptPersistenceGateway = { + openInterruptBatch: async (input) => ({ + generation: 1, + descriptors: input.descriptors, + }), + commitInterruptResolutions: async (input) => ({ + status: 'committed', + continuationRunId: input.continuationRunId, + }), + getInterruptRecoveryState: async (input) => ({ + schemaVersion: 1, + state: 'missing', + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + generation: input.knownGeneration, + pendingInterrupts: [], + }), +} + +const testInterruptPersistence = defineChatMiddleware({ + name: 'openrouter-test-interrupt-persistence', + provides: [InterruptPersistenceCapability], + setup(ctx) { + provideInterruptPersistence(ctx, testInterruptGateway) + }, + async onChunk(_ctx, chunk) { + if (chunk.type === 'RUN_FINISHED' && chunk.outcome?.type === 'interrupt') { + await testInterruptGateway.openInterruptBatch({ + threadId: chunk.threadId, + interruptedRunId: chunk.runId, + descriptors: chunk.outcome.interrupts, + bindings: [], + }) + } + }, +}) + function createAsyncIterable(chunks: Array): AsyncIterable { return { [Symbol.asyncIterator]() { @@ -691,26 +738,27 @@ describe('OpenRouter responses adapter — stream event bridge', () => { adapter, messages: [{ role: 'user', content: 'hi' }], tools: [weatherTool], + middleware: [testInterruptPersistence], })) { chunks.push(c) } - const start = chunks.find((c) => c.type === 'TOOL_CALL_START') as any + const start = chunks.find((c) => c.type === 'TOOL_CALL_START') expect(start).toMatchObject({ type: 'TOOL_CALL_START', toolCallId: 'item_1', toolCallName: 'lookup_weather', }) - const args = chunks.filter((c) => c.type === 'TOOL_CALL_ARGS') as any[] + const args = chunks.filter((c) => c.type === 'TOOL_CALL_ARGS') expect(args.length).toBe(1) - expect(args[0]!.delta).toBe('{"location":"Berlin"}') + expect(args[0]).toMatchObject({ delta: '{"location":"Berlin"}' }) - const end = chunks.find((c) => c.type === 'TOOL_CALL_END') as any - expect(end.input).toEqual({ location: 'Berlin' }) + const end = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(end).toMatchObject({ input: { location: 'Berlin' } }) - const finished = chunks.find((c) => c.type === 'RUN_FINISHED') as any - expect(finished.finishReason).toBe('tool_calls') + const finished = chunks.find((c) => c.type === 'RUN_FINISHED') + expect(finished).toMatchObject({ finishReason: 'tool_calls' }) }) it('emits parentMessageId on tool-first tool calls matching the assistant message id', async () => { diff --git a/packages/ai-persistence-cloudflare/migrations/0001_tanstack_ai_interrupt_batches.sql b/packages/ai-persistence-cloudflare/migrations/0001_tanstack_ai_interrupt_batches.sql new file mode 100644 index 000000000..d81b74cee --- /dev/null +++ b/packages/ai-persistence-cloudflare/migrations/0001_tanstack_ai_interrupt_batches.sql @@ -0,0 +1,68 @@ +CREATE TABLE `interrupt_batches` ( + `interrupted_run_id` text PRIMARY KEY NOT NULL, + `thread_id` text NOT NULL, + `generation` integer NOT NULL, + `expected_interrupt_ids_json` text NOT NULL, + `fingerprint` text NOT NULL, + `canonical_resolutions` text NOT NULL, + `resolutions_json` text NOT NULL, + `continuation_run_id` text NOT NULL, + `committed_at` integer NOT NULL +); +--> statement-breakpoint +CREATE UNIQUE INDEX `interrupt_batches_continuation_run_id_unique` ON `interrupt_batches` (`continuation_run_id`); +--> statement-breakpoint +CREATE TABLE `__new_interrupts` ( + `interrupt_id` text PRIMARY KEY NOT NULL, + `run_id` text NOT NULL, + `thread_id` text NOT NULL, + `generation` integer NOT NULL, + `status` text NOT NULL, + `requested_at` integer NOT NULL, + `resolved_at` integer, + `payload_json` text NOT NULL, + `binding_json` text NOT NULL, + `response_schema_hash` text NOT NULL, + `response_json` text +); +--> statement-breakpoint +INSERT INTO `__new_interrupts` ( + `interrupt_id`, + `run_id`, + `thread_id`, + `generation`, + `status`, + `requested_at`, + `resolved_at`, + `payload_json`, + `binding_json`, + `response_schema_hash`, + `response_json` +) +SELECT + `interrupt_id`, + `run_id`, + `thread_id`, + CASE WHEN `status` = 'pending' THEN 1 ELSE 0 END, + `status`, + `requested_at`, + `resolved_at`, + `payload_json`, + json_object( + 'kind', 'generic', + 'interruptId', `interrupt_id`, + 'interruptedRunId', `run_id`, + 'generation', CASE WHEN `status` = 'pending' THEN 1 ELSE 0 END, + 'responseSchemaHash', 'legacy:unknown' + ), + 'legacy:unknown', + `response_json` +FROM `interrupts`; +--> statement-breakpoint +DROP TABLE `interrupts`; +--> statement-breakpoint +ALTER TABLE `__new_interrupts` RENAME TO `interrupts`; +--> statement-breakpoint +CREATE INDEX `interrupts_thread_status_requested_at_idx` ON `interrupts` (`thread_id`,`status`,`requested_at`); +--> statement-breakpoint +CREATE INDEX `interrupts_run_status_requested_at_idx` ON `interrupts` (`run_id`,`status`,`requested_at`); diff --git a/packages/ai-persistence-cloudflare/src/assets/0001_tanstack_ai_interrupt_batches.sql b/packages/ai-persistence-cloudflare/src/assets/0001_tanstack_ai_interrupt_batches.sql new file mode 100644 index 000000000..d81b74cee --- /dev/null +++ b/packages/ai-persistence-cloudflare/src/assets/0001_tanstack_ai_interrupt_batches.sql @@ -0,0 +1,68 @@ +CREATE TABLE `interrupt_batches` ( + `interrupted_run_id` text PRIMARY KEY NOT NULL, + `thread_id` text NOT NULL, + `generation` integer NOT NULL, + `expected_interrupt_ids_json` text NOT NULL, + `fingerprint` text NOT NULL, + `canonical_resolutions` text NOT NULL, + `resolutions_json` text NOT NULL, + `continuation_run_id` text NOT NULL, + `committed_at` integer NOT NULL +); +--> statement-breakpoint +CREATE UNIQUE INDEX `interrupt_batches_continuation_run_id_unique` ON `interrupt_batches` (`continuation_run_id`); +--> statement-breakpoint +CREATE TABLE `__new_interrupts` ( + `interrupt_id` text PRIMARY KEY NOT NULL, + `run_id` text NOT NULL, + `thread_id` text NOT NULL, + `generation` integer NOT NULL, + `status` text NOT NULL, + `requested_at` integer NOT NULL, + `resolved_at` integer, + `payload_json` text NOT NULL, + `binding_json` text NOT NULL, + `response_schema_hash` text NOT NULL, + `response_json` text +); +--> statement-breakpoint +INSERT INTO `__new_interrupts` ( + `interrupt_id`, + `run_id`, + `thread_id`, + `generation`, + `status`, + `requested_at`, + `resolved_at`, + `payload_json`, + `binding_json`, + `response_schema_hash`, + `response_json` +) +SELECT + `interrupt_id`, + `run_id`, + `thread_id`, + CASE WHEN `status` = 'pending' THEN 1 ELSE 0 END, + `status`, + `requested_at`, + `resolved_at`, + `payload_json`, + json_object( + 'kind', 'generic', + 'interruptId', `interrupt_id`, + 'interruptedRunId', `run_id`, + 'generation', CASE WHEN `status` = 'pending' THEN 1 ELSE 0 END, + 'responseSchemaHash', 'legacy:unknown' + ), + 'legacy:unknown', + `response_json` +FROM `interrupts`; +--> statement-breakpoint +DROP TABLE `interrupts`; +--> statement-breakpoint +ALTER TABLE `__new_interrupts` RENAME TO `interrupts`; +--> statement-breakpoint +CREATE INDEX `interrupts_thread_status_requested_at_idx` ON `interrupts` (`thread_id`,`status`,`requested_at`); +--> statement-breakpoint +CREATE INDEX `interrupts_run_status_requested_at_idx` ON `interrupts` (`run_id`,`status`,`requested_at`); diff --git a/packages/ai-persistence-cloudflare/src/d1.ts b/packages/ai-persistence-cloudflare/src/d1.ts index 011e70ce4..13a63d517 100644 --- a/packages/ai-persistence-cloudflare/src/d1.ts +++ b/packages/ai-persistence-cloudflare/src/d1.ts @@ -1,13 +1,1008 @@ +import { + canonicalInterruptJson, + canonicalizeInterruptResolutions, + cloneAndDeepFreezeJson, +} from '@tanstack/ai' import { drizzle } from 'drizzle-orm/d1' import { drizzlePersistence, schema } from '@tanstack/ai-persistence-drizzle' +import { + InterruptStoreCorruptionError, + hasExactInterruptIds, + projectInterruptRecovery, +} from '@tanstack/ai-persistence' +import type { + InterruptBinding, + InterruptRecoveryQuery, + RunAgentResumeItem, +} from '@tanstack/ai' +import type { + InterruptBatchRecord, + InterruptRecord, + InterruptStore, +} from '@tanstack/ai-persistence' + +interface D1InterruptRow { + interrupt_id: string + run_id: string + thread_id: string + generation: number + status: string + requested_at: number + resolved_at: number | null + payload_json: string + binding_json: string + response_schema_hash: string + response_json: string | null +} + +interface D1InterruptBatchRow { + interrupted_run_id: string + thread_id: string + generation: number + expected_interrupt_ids_json: string + fingerprint: string + canonical_resolutions: string + resolutions_json: string + continuation_run_id: string + committed_at: number +} + +type D1RecoveryRow = D1InterruptRow | D1InterruptBatchRow + +export interface D1InterruptStoreOptions { + /** Override wall-clock time, primarily for deterministic runtimes/tests. */ + clock?: () => number +} + +function corrupt(message: string, options?: ErrorOptions): never { + throw new InterruptStoreCorruptionError(message, options) +} + +function parseJson(value: string, field: string): unknown { + try { + const parsed: unknown = JSON.parse(value) + return parsed + } catch (error) { + return corrupt(`Stored ${field} is not valid JSON.`, { cause: error }) + } +} + +function decodeBinding( + value: unknown, + correlation: { + interruptId: string + interruptedRunId: string + generation: number + responseSchemaHash: string + }, +): InterruptBinding { + if ( + value === null || + typeof value !== 'object' || + !('kind' in value) || + !('interruptId' in value) || + !('interruptedRunId' in value) || + !('generation' in value) || + !('responseSchemaHash' in value) || + typeof value.interruptId !== 'string' || + typeof value.interruptedRunId !== 'string' || + typeof value.generation !== 'number' || + typeof value.responseSchemaHash !== 'string' + ) { + return corrupt('Stored interrupt binding is malformed.') + } + const expiresAt = 'expiresAt' in value ? value.expiresAt : undefined + if ( + expiresAt !== undefined && + (typeof expiresAt !== 'string' || !Number.isFinite(Date.parse(expiresAt))) + ) { + return corrupt('Stored interrupt expiry is malformed.') + } + if ( + value.interruptId !== correlation.interruptId || + value.interruptedRunId !== correlation.interruptedRunId || + value.generation !== correlation.generation || + value.responseSchemaHash !== correlation.responseSchemaHash + ) { + return corrupt('Stored interrupt binding correlation is inconsistent.') + } + if (value.kind === 'generic') { + return { + kind: 'generic', + interruptId: value.interruptId, + interruptedRunId: value.interruptedRunId, + generation: value.generation, + responseSchemaHash: value.responseSchemaHash, + ...(expiresAt !== undefined ? { expiresAt } : {}), + } + } + if ( + value.kind === 'client-tool-execution' && + 'toolName' in value && + 'toolCallId' in value && + 'outputSchemaHash' in value && + typeof value.toolName === 'string' && + typeof value.toolCallId === 'string' && + typeof value.outputSchemaHash === 'string' + ) { + return { + kind: 'client-tool-execution', + interruptId: value.interruptId, + interruptedRunId: value.interruptedRunId, + generation: value.generation, + toolName: value.toolName, + toolCallId: value.toolCallId, + outputSchemaHash: value.outputSchemaHash, + responseSchemaHash: value.responseSchemaHash, + ...(expiresAt !== undefined ? { expiresAt } : {}), + } + } + if ( + value.kind === 'tool-approval' && + 'toolName' in value && + 'toolCallId' in value && + 'originalArgs' in value && + 'inputSchemaHash' in value && + 'approvalSchemaHash' in value && + typeof value.toolName === 'string' && + typeof value.toolCallId === 'string' && + typeof value.inputSchemaHash === 'string' && + typeof value.approvalSchemaHash === 'string' + ) { + return { + kind: 'tool-approval', + interruptId: value.interruptId, + interruptedRunId: value.interruptedRunId, + generation: value.generation, + toolName: value.toolName, + toolCallId: value.toolCallId, + originalArgs: value.originalArgs, + inputSchemaHash: value.inputSchemaHash, + approvalSchemaHash: value.approvalSchemaHash, + responseSchemaHash: value.responseSchemaHash, + ...(expiresAt !== undefined ? { expiresAt } : {}), + } + } + return corrupt('Stored interrupt binding kind is malformed.') +} + +function validateDescriptor(value: unknown, interruptId: string): void { + if ( + value === null || + typeof value !== 'object' || + !('id' in value) || + !('reason' in value) || + value.id !== interruptId || + typeof value.reason !== 'string' + ) { + corrupt('Stored interrupt descriptor correlation is inconsistent.') + } +} + +function decodeStatus(value: string): InterruptRecord['status'] { + if (value === 'pending' || value === 'resolved' || value === 'cancelled') { + return value + } + return corrupt('Stored interrupt status is malformed.') +} + +function mapInterrupt(row: D1InterruptRow): InterruptRecord { + const payload = parseJson(row.payload_json, 'interrupt payload') + if (row.generation > 0) validateDescriptor(payload, row.interrupt_id) + const binding = decodeBinding( + parseJson(row.binding_json, 'interrupt binding'), + { + interruptId: row.interrupt_id, + interruptedRunId: row.run_id, + generation: row.generation, + responseSchemaHash: row.response_schema_hash, + }, + ) + return { + interruptId: row.interrupt_id, + runId: row.run_id, + threadId: row.thread_id, + generation: row.generation, + status: decodeStatus(row.status), + requestedAt: row.requested_at, + ...(row.resolved_at !== null ? { resolvedAt: row.resolved_at } : {}), + payload, + binding, + ...(row.response_json !== null + ? { response: parseJson(row.response_json, 'interrupt response') } + : {}), + } +} + +function decodeStringArray(value: unknown, field: string): Array { + if (!Array.isArray(value)) return corrupt(`Stored ${field} is malformed.`) + const strings: Array = [] + for (const item of value) { + if (typeof item !== 'string') { + return corrupt(`Stored ${field} is malformed.`) + } + strings.push(item) + } + if (!hasExactInterruptIds(strings, strings)) { + return corrupt(`Stored ${field} contains duplicate IDs.`) + } + return strings +} + +function decodeResolution(value: unknown): RunAgentResumeItem { + if ( + value === null || + typeof value !== 'object' || + !('interruptId' in value) || + !('status' in value) || + typeof value.interruptId !== 'string' || + (value.status !== 'resolved' && value.status !== 'cancelled') + ) { + return corrupt('Stored interrupt resolution is malformed.') + } + return { + interruptId: value.interruptId, + status: value.status, + ...('payload' in value ? { payload: value.payload } : {}), + } +} + +function decodeResolutions(value: unknown): Array { + if (!Array.isArray(value)) { + return corrupt('Stored interrupt resolutions are malformed.') + } + return value.map(decodeResolution) +} + +function mapBatch(row: D1InterruptBatchRow): InterruptBatchRecord { + const expectedInterruptIds = decodeStringArray( + parseJson(row.expected_interrupt_ids_json, 'expected interrupt IDs'), + 'expected interrupt IDs', + ) + const resolutions = decodeResolutions( + parseJson(row.resolutions_json, 'interrupt resolutions'), + ) + const candidate = canonicalizeInterruptResolutions(resolutions) + if ( + !hasExactInterruptIds( + expectedInterruptIds, + resolutions.map((resolution) => resolution.interruptId), + ) || + candidate.fingerprint !== row.fingerprint || + candidate.canonicalResolutions !== row.canonical_resolutions + ) { + return corrupt('Stored interrupt batch identity is inconsistent.') + } + return { + threadId: row.thread_id, + interruptedRunId: row.interrupted_run_id, + generation: row.generation, + expectedInterruptIds, + fingerprint: row.fingerprint, + canonicalResolutions: row.canonical_resolutions, + resolutions, + continuationRunId: row.continuation_run_id, + committedAt: row.committed_at, + } +} + +function isInterruptRow(row: D1RecoveryRow): row is D1InterruptRow { + return 'interrupt_id' in row +} + +function isBatchRow(row: D1RecoveryRow): row is D1InterruptBatchRow { + return 'interrupted_run_id' in row +} + +function inspectResults( + results: Array>, + expectedCount: number, +): void { + if (results.length !== expectedCount) { + corrupt( + `D1 batch returned ${results.length} results for ${expectedCount} statements.`, + ) + } + for (const result of results) { + if ( + !Array.isArray(result.results) || + typeof result.meta.changes !== 'number' + ) { + corrupt('D1 batch returned a malformed statement result.') + } + } +} + +function mutationChanges(result: D1Result, label: string): number { + if (typeof result.meta.changes !== 'number') { + return corrupt(`D1 ${label} returned malformed change metadata.`) + } + return result.meta.changes +} + +function requireResult( + results: Array>, + index: number, + label: string, +): D1Result { + const result = results[index] + return result ?? corrupt(`D1 ${label} result is missing.`) +} + +function bind( + d1: D1Database, + sql: string, + values: ReadonlyArray, +): D1PreparedStatement { + return d1.prepare(sql).bind(...values) +} + +/** Direct, atomic D1 interrupt store. */ +export function createD1InterruptStore( + d1: D1Database, + options: D1InterruptStoreOptions = {}, +): InterruptStore { + const clock = options.clock ?? Date.now + + const readBatchAndRows = async (interruptedRunId: string) => { + const statements = [ + bind(d1, 'SELECT * FROM interrupt_batches WHERE interrupted_run_id = ?', [ + interruptedRunId, + ]), + bind( + d1, + 'SELECT * FROM interrupts WHERE run_id = ? ORDER BY requested_at, interrupt_id', + [interruptedRunId], + ), + ] + const results = await d1.batch(statements) + inspectResults(results, statements.length) + const batchRows = results[0]?.results ?? [] + const interruptRows = results[1]?.results ?? [] + if (batchRows.length > 1) { + return corrupt('D1 returned malformed interrupt batch rows.') + } + let batch: InterruptBatchRecord | null = null + if (batchRows[0]) { + if (!isBatchRow(batchRows[0])) { + return corrupt('D1 returned malformed interrupt batch rows.') + } + batch = mapBatch(batchRows[0]) + } + const rows: Array = [] + for (const row of interruptRows) { + if (!isInterruptRow(row)) { + return corrupt('D1 returned malformed interrupt rows.') + } + rows.push(mapInterrupt(row)) + } + return { + batch, + rows, + } + } + + const readRecovery = async (input: InterruptRecoveryQuery) => { + const snapshot = await readBatchAndRows(input.interruptedRunId) + const rows = snapshot.rows.filter((row) => row.threadId === input.threadId) + const batch = + snapshot.batch?.threadId === input.threadId ? snapshot.batch : null + const generations = new Set(rows.map((row) => row.generation)) + if (generations.size > 1) { + return corrupt('Stored interrupt rows mix generations.') + } + if ( + batch && + (rows.length === 0 || + rows.some((row) => row.generation !== batch.generation) || + !hasExactInterruptIds( + batch.expectedInterruptIds, + rows.map((row) => row.interruptId), + )) + ) { + return corrupt('Stored interrupt batch correlation is inconsistent.') + } + return projectInterruptRecovery({ + query: input, + rows, + batch, + now: clock(), + includeResolutions: true, + }) + } + + const store: InterruptStore = { + async create(record) { + const generation = record.generation ?? 0 + const binding = decodeBinding( + cloneAndDeepFreezeJson( + record.binding ?? { + kind: 'generic', + interruptId: record.interruptId, + interruptedRunId: record.runId, + generation, + responseSchemaHash: 'legacy:unknown', + }, + ), + { + interruptId: record.interruptId, + interruptedRunId: record.runId, + generation, + responseSchemaHash: + record.binding?.responseSchemaHash ?? 'legacy:unknown', + }, + ) + const result = await bind( + d1, + `INSERT INTO interrupts ( + interrupt_id, run_id, thread_id, generation, status, requested_at, + resolved_at, payload_json, binding_json, response_schema_hash, + response_json + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(interrupt_id) DO NOTHING`, + [ + record.interruptId, + record.runId, + record.threadId, + generation, + record.status, + record.requestedAt, + null, + canonicalInterruptJson(record.payload), + canonicalInterruptJson(binding), + binding.responseSchemaHash, + record.response === undefined + ? null + : canonicalInterruptJson(record.response), + ], + ).run() + const changes = mutationChanges(result, 'legacy interrupt insert') + if (changes !== 0 && changes !== 1) { + corrupt('D1 legacy interrupt insert changed an unexpected row count.') + } + }, + async resolve(interruptId, response) { + const hasResponse = response !== undefined + const result = await bind( + d1, + `UPDATE interrupts SET + status = 'resolved', resolved_at = ? + ${hasResponse ? ', response_json = ?' : ''} + WHERE interrupt_id = ?`, + [ + clock(), + ...(hasResponse ? [canonicalInterruptJson(response)] : []), + interruptId, + ], + ).run() + const changes = mutationChanges(result, 'legacy interrupt resolve') + if (changes !== 0 && changes !== 1) { + corrupt('D1 legacy interrupt resolve changed an unexpected row count.') + } + }, + async cancel(interruptId) { + const result = await bind( + d1, + `UPDATE interrupts SET status = 'cancelled', resolved_at = ? + WHERE interrupt_id = ?`, + [clock(), interruptId], + ).run() + const changes = mutationChanges(result, 'legacy interrupt cancellation') + if (changes !== 0 && changes !== 1) { + corrupt( + 'D1 legacy interrupt cancellation changed an unexpected row count.', + ) + } + }, + async get(interruptId) { + const result = await bind( + d1, + 'SELECT * FROM interrupts WHERE interrupt_id = ?', + [interruptId], + ).all() + inspectResults([result], 1) + if (result.results.length > 1) { + return corrupt('D1 returned duplicate interrupt IDs.') + } + return result.results[0] ? mapInterrupt(result.results[0]) : null + }, + async list(threadId) { + const result = await bind( + d1, + 'SELECT * FROM interrupts WHERE thread_id = ? ORDER BY requested_at, interrupt_id', + [threadId], + ).all() + inspectResults([result], 1) + return result.results.map(mapInterrupt) + }, + async listPending(threadId) { + const result = await bind( + d1, + `SELECT * FROM interrupts + WHERE thread_id = ? AND status = 'pending' + ORDER BY requested_at, interrupt_id`, + [threadId], + ).all() + inspectResults([result], 1) + return result.results.map(mapInterrupt) + }, + async listByRun(runId) { + const snapshot = await readBatchAndRows(runId) + return snapshot.rows + }, + async listPendingByRun(runId) { + const result = await bind( + d1, + `SELECT * FROM interrupts + WHERE run_id = ? AND status = 'pending' + ORDER BY requested_at, interrupt_id`, + [runId], + ).all() + inspectResults([result], 1) + return result.results.map(mapInterrupt) + }, + async openInterruptBatch(input) { + const descriptors = cloneAndDeepFreezeJson(input.descriptors) + const unopenedBindings = cloneAndDeepFreezeJson(input.bindings) + const descriptorIds = descriptors.map((descriptor) => descriptor.id) + if ( + descriptorIds.length === 0 || + !hasExactInterruptIds( + descriptorIds, + unopenedBindings.map((binding) => binding.interruptId), + ) + ) { + throw new TypeError( + 'Interrupt descriptors and bindings must contain the same exact nonempty IDs.', + ) + } + for (const descriptor of descriptors) { + validateDescriptor(descriptor, descriptor.id) + } + + const snapshot = await readBatchAndRows(input.interruptedRunId) + if (snapshot.batch) { + throw new Error('Cannot reopen a committed interrupt batch.') + } + if (snapshot.rows.length > 0) { + const generation = snapshot.rows[0]?.generation + const requestedBindings = unopenedBindings + .map((binding) => ({ + ...binding, + interruptedRunId: input.interruptedRunId, + generation, + })) + .sort((left, right) => + left.interruptId.localeCompare(right.interruptId), + ) + const existingBindings = snapshot.rows + .map((row) => row.binding) + .sort((left, right) => + left.interruptId.localeCompare(right.interruptId), + ) + const existingDescriptors = snapshot.rows + .map((row) => row.payload) + .sort((left, right) => { + if ( + left !== null && + typeof left === 'object' && + 'id' in left && + typeof left.id === 'string' && + right !== null && + typeof right === 'object' && + 'id' in right && + typeof right.id === 'string' + ) { + return left.id.localeCompare(right.id) + } + return 0 + }) + const requestedDescriptors = [...descriptors].sort((left, right) => + left.id.localeCompare(right.id), + ) + if ( + generation !== undefined && + snapshot.rows.every( + (row) => + row.threadId === input.threadId && + row.generation === generation && + row.status === 'pending', + ) && + hasExactInterruptIds( + descriptorIds, + snapshot.rows.map((row) => row.interruptId), + ) && + canonicalInterruptJson(existingDescriptors) === + canonicalInterruptJson(requestedDescriptors) && + canonicalInterruptJson(existingBindings) === + canonicalInterruptJson(requestedBindings) + ) { + return { generation, descriptors } + } + throw new Error('An incompatible interrupt batch is already open.') + } + + const generation = 1 + const bindingById = new Map( + unopenedBindings.map((binding) => [binding.interruptId, binding]), + ) + const requestedAt = clock() + const encoded = descriptors.map((descriptor) => { + const unopened = bindingById.get(descriptor.id) + if (!unopened) { + throw new Error(`Missing binding for interrupt: ${descriptor.id}`) + } + const proposedBinding = cloneAndDeepFreezeJson({ + ...unopened, + interruptedRunId: input.interruptedRunId, + generation, + }) + const binding = decodeBinding(proposedBinding, { + interruptId: descriptor.id, + interruptedRunId: input.interruptedRunId, + generation, + responseSchemaHash: proposedBinding.responseSchemaHash, + }) + return { + descriptor, + binding, + payloadJson: canonicalInterruptJson(descriptor), + bindingJson: canonicalInterruptJson(binding), + } + }) + const insertStatements = encoded.map((item) => + bind( + d1, + `INSERT INTO interrupts ( + interrupt_id, run_id, thread_id, generation, status, requested_at, + payload_json, binding_json, response_schema_hash + ) VALUES (?, ?, ?, ?, 'pending', ?, ?, ?, ?) + ON CONFLICT(interrupt_id) DO NOTHING`, + [ + item.descriptor.id, + input.interruptedRunId, + input.threadId, + generation, + requestedAt, + item.payloadJson, + item.bindingJson, + item.binding.responseSchemaHash, + ], + ), + ) + const exactClauses = encoded + .map( + () => + '(interrupt_id = ? AND payload_json = ? AND binding_json = ? AND response_schema_hash = ?)', + ) + .join(' OR ') + const assertionValues: Array = [ + input.interruptedRunId, + input.interruptedRunId, + encoded.length, + input.interruptedRunId, + input.threadId, + generation, + ] + for (const item of encoded) { + assertionValues.push( + item.descriptor.id, + item.payloadJson, + item.bindingJson, + item.binding.responseSchemaHash, + ) + } + assertionValues.push(encoded.length) + const assertion = bind( + d1, + `INSERT INTO interrupts ( + interrupt_id, run_id, thread_id, generation, status, requested_at, + payload_json, binding_json, response_schema_hash + ) + SELECT NULL, '', '', 0, 'pending', 0, '{}', '{}', '' + WHERE NOT ( + NOT EXISTS ( + SELECT 1 FROM interrupt_batches WHERE interrupted_run_id = ? + ) + AND (SELECT COUNT(*) FROM interrupts WHERE run_id = ?) = ? + AND ( + SELECT COUNT(*) FROM interrupts + WHERE run_id = ? AND thread_id = ? AND generation = ? + AND status = 'pending' AND (${exactClauses}) + ) = ? + )`, + assertionValues, + ) + const statements = [...insertStatements, assertion] + const results = await d1.batch(statements) + inspectResults(results, statements.length) + for (const result of results.slice(0, -1)) { + const changes = mutationChanges(result, 'interrupt open insert') + if (changes !== 0 && changes !== 1) { + corrupt('D1 interrupt open changed an unexpected row count.') + } + } + const openAssertionResult = requireResult( + results, + results.length - 1, + 'interrupt open assertion', + ) + if ( + mutationChanges(openAssertionResult, 'interrupt open assertion') !== 0 + ) { + corrupt('D1 interrupt open assertion unexpectedly changed a row.') + } + const opened = await readBatchAndRows(input.interruptedRunId) + if ( + opened.batch || + !hasExactInterruptIds( + descriptorIds, + opened.rows.map((row) => row.interruptId), + ) + ) { + return corrupt('D1 interrupt open did not persist the exact batch.') + } + return { generation, descriptors } + }, + async commitInterruptResolutions(input) { + const candidate = canonicalizeInterruptResolutions(input.resolutions) + if ( + candidate.fingerprint !== input.fingerprint || + candidate.canonicalResolutions !== input.canonicalResolutions + ) { + throw new TypeError( + 'Interrupt batch identity does not match its resolutions.', + ) + } + const resolutionIds = candidate.resolutions.map( + (resolution) => resolution.interruptId, + ) + if (!hasExactInterruptIds(input.expectedInterruptIds, resolutionIds)) { + return { + status: 'conflict', + authoritativeState: await readRecovery({ + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + knownGeneration: input.expectedGeneration, + }), + } + } + + const snapshot = await readBatchAndRows(input.interruptedRunId) + if (snapshot.batch) { + return snapshot.batch.threadId === input.threadId && + snapshot.batch.fingerprint === candidate.fingerprint && + snapshot.batch.canonicalResolutions === candidate.canonicalResolutions + ? { + status: 'replayed', + continuationRunId: snapshot.batch.continuationRunId, + } + : { + status: 'conflict', + authoritativeState: await readRecovery({ + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + knownGeneration: input.expectedGeneration, + }), + } + } + const pending = snapshot.rows.filter((row) => row.status === 'pending') + if ( + !hasExactInterruptIds( + input.expectedInterruptIds, + pending.map((row) => row.interruptId), + ) || + pending.some( + (row) => + row.threadId !== input.threadId || + row.generation !== input.expectedGeneration || + (row.binding.expiresAt !== undefined && + Date.parse(row.binding.expiresAt) <= clock()), + ) + ) { + return { + status: 'conflict', + authoritativeState: await readRecovery({ + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + knownGeneration: input.expectedGeneration, + }), + } + } + + const committedAt = clock() + const idPlaceholders = input.expectedInterruptIds.map(() => '?').join(',') + const winnerInsert = bind( + d1, + `INSERT INTO interrupt_batches ( + interrupted_run_id, thread_id, generation, + expected_interrupt_ids_json, fingerprint, canonical_resolutions, + resolutions_json, continuation_run_id, committed_at + ) + SELECT ?, ?, ?, ?, ?, ?, ?, ?, ? + WHERE NOT EXISTS ( + SELECT 1 FROM interrupt_batches WHERE interrupted_run_id = ? + ) + AND ( + SELECT COUNT(*) FROM interrupts + WHERE run_id = ? AND status = 'pending' + ) = ? + AND ( + SELECT COUNT(*) FROM interrupts + WHERE run_id = ? AND thread_id = ? AND generation = ? + AND status = 'pending' + AND interrupt_id IN (${idPlaceholders}) + AND ( + json_extract(binding_json, '$.expiresAt') IS NULL + OR json_extract(binding_json, '$.expiresAt') > ? + ) + ) = ? + ON CONFLICT(interrupted_run_id) DO NOTHING`, + [ + input.interruptedRunId, + input.threadId, + input.expectedGeneration, + canonicalInterruptJson(input.expectedInterruptIds), + candidate.fingerprint, + candidate.canonicalResolutions, + candidate.canonicalResolutions, + input.continuationRunId, + committedAt, + input.interruptedRunId, + input.interruptedRunId, + input.expectedInterruptIds.length, + input.interruptedRunId, + input.threadId, + input.expectedGeneration, + ...input.expectedInterruptIds, + new Date(committedAt).toISOString(), + input.expectedInterruptIds.length, + ], + ) + const transitionStatements = candidate.resolutions.map((resolution) => { + const status = + resolution.status === 'cancelled' ? 'cancelled' : 'resolved' + return bind( + d1, + `UPDATE interrupts + SET status = ?, resolved_at = ?, response_json = ? + WHERE interrupt_id = ? AND run_id = ? AND thread_id = ? + AND generation = ? AND status = 'pending' + AND EXISTS ( + SELECT 1 FROM interrupt_batches + WHERE interrupted_run_id = ? AND thread_id = ? + AND generation = ? AND fingerprint = ? + AND canonical_resolutions = ? AND continuation_run_id = ? + )`, + [ + status, + committedAt, + canonicalInterruptJson(resolution), + resolution.interruptId, + input.interruptedRunId, + input.threadId, + input.expectedGeneration, + input.interruptedRunId, + input.threadId, + input.expectedGeneration, + candidate.fingerprint, + candidate.canonicalResolutions, + input.continuationRunId, + ], + ) + }) + const finalClauses = candidate.resolutions + .map( + () => + '(interrupt_id = ? AND status = ? AND resolved_at = ? AND response_json = ?)', + ) + .join(' OR ') + const assertionValues: Array = [ + input.interruptedRunId, + input.threadId, + input.expectedGeneration, + candidate.fingerprint, + candidate.canonicalResolutions, + input.continuationRunId, + input.interruptedRunId, + ] + for (const resolution of candidate.resolutions) { + assertionValues.push( + resolution.interruptId, + resolution.status === 'cancelled' ? 'cancelled' : 'resolved', + committedAt, + canonicalInterruptJson(resolution), + ) + } + assertionValues.push(candidate.resolutions.length) + const assertion = bind( + d1, + `INSERT INTO interrupts ( + interrupt_id, run_id, thread_id, generation, status, requested_at, + payload_json, binding_json, response_schema_hash + ) + SELECT NULL, '', '', 0, 'pending', 0, '{}', '{}', '' + WHERE EXISTS ( + SELECT 1 FROM interrupt_batches + WHERE interrupted_run_id = ? AND thread_id = ? AND generation = ? + AND fingerprint = ? AND canonical_resolutions = ? + AND continuation_run_id = ? + ) + AND ( + SELECT COUNT(*) FROM interrupts + WHERE run_id = ? AND (${finalClauses}) + ) <> ?`, + assertionValues, + ) + const statements = [winnerInsert, ...transitionStatements, assertion] + const results = await d1.batch(statements) + inspectResults(results, statements.length) + const insertChanges = mutationChanges( + requireResult(results, 0, 'winner insert'), + 'winner insert', + ) + if (insertChanges === 0) { + const authoritative = await readBatchAndRows(input.interruptedRunId) + if ( + authoritative.batch?.threadId === input.threadId && + authoritative.batch.fingerprint === candidate.fingerprint && + authoritative.batch.canonicalResolutions === + candidate.canonicalResolutions + ) { + return { + status: 'replayed', + continuationRunId: authoritative.batch.continuationRunId, + } + } + return { + status: 'conflict', + authoritativeState: await readRecovery({ + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + knownGeneration: input.expectedGeneration, + }), + } + } + if (insertChanges !== 1) { + return corrupt('D1 winner insert changed an unexpected row count.') + } + for (const result of results.slice(1, -1)) { + if (mutationChanges(result, 'interrupt transition') !== 1) { + throw new Error('A D1 interrupt transition changed no rows.') + } + } + const commitAssertionResult = requireResult( + results, + results.length - 1, + 'commit assertion', + ) + if (mutationChanges(commitAssertionResult, 'commit assertion') !== 0) { + return corrupt('D1 commit assertion unexpectedly changed a row.') + } + return { + status: 'committed', + continuationRunId: input.continuationRunId, + } + }, + getInterruptRecoveryState(input) { + return readRecovery(input) + }, + } + return store +} /** Create the structured stores owned by a migrated Cloudflare D1 binding. */ -export function createD1Stores(d1: D1Database) { - const persistence = drizzlePersistence(drizzle(d1, { schema })) +export function createD1Stores( + d1: D1Database, + interruptOptions?: D1InterruptStoreOptions, +) { + const persistence = drizzlePersistence(drizzle(d1, { schema }), { + interrupts: false, + }) return { messages: persistence.stores.messages, runs: persistence.stores.runs, - interrupts: persistence.stores.interrupts, + interrupts: createD1InterruptStore(d1, interruptOptions), metadata: persistence.stores.metadata, } } diff --git a/packages/ai-persistence-cloudflare/src/index.ts b/packages/ai-persistence-cloudflare/src/index.ts index 76e31cc7f..d95c501c1 100644 --- a/packages/ai-persistence-cloudflare/src/index.ts +++ b/packages/ai-persistence-cloudflare/src/index.ts @@ -15,7 +15,8 @@ import type { LockStore } from '@tanstack/ai' import type { R2BucketBinding } from './bindings' import type { DurableObjectLockStoreOptions } from './locks' -export { createD1Stores } from './d1' +export { createD1InterruptStore, createD1Stores } from './d1' +export type { D1InterruptStoreOptions } from './d1' export { CloudflareLockDurableObject, createDurableObjectLockStore, diff --git a/packages/ai-persistence-cloudflare/src/migrations.ts b/packages/ai-persistence-cloudflare/src/migrations.ts index 2399a8c10..5f3a2f948 100644 --- a/packages/ai-persistence-cloudflare/src/migrations.ts +++ b/packages/ai-persistence-cloudflare/src/migrations.ts @@ -1,4 +1,5 @@ import initialMigrationSql from './assets/0000_tanstack_ai_initial.sql?raw' +import interruptBatchesMigrationSql from './assets/0001_tanstack_ai_interrupt_batches.sql?raw' export interface D1Migration { id: string @@ -13,4 +14,9 @@ export const d1Migrations: ReadonlyArray = [ filename: '0000_tanstack_ai_initial.sql', sql: initialMigrationSql, }, + { + id: '0001_tanstack_ai_interrupt_batches', + filename: '0001_tanstack_ai_interrupt_batches.sql', + sql: interruptBatchesMigrationSql, + }, ] diff --git a/packages/ai-persistence-cloudflare/tests/api-types.test-d.ts b/packages/ai-persistence-cloudflare/tests/api-types.test-d.ts index 4bc9c92f6..62d99c09b 100644 --- a/packages/ai-persistence-cloudflare/tests/api-types.test-d.ts +++ b/packages/ai-persistence-cloudflare/tests/api-types.test-d.ts @@ -4,6 +4,8 @@ import { composePersistence } from '@tanstack/ai-persistence' import { CloudflareLockDurableObject, cloudflarePersistence, + createD1InterruptStore, + createD1Stores, } from '../src/index' import type { ArtifactStore, @@ -14,6 +16,7 @@ import type { RunStore, } from '@tanstack/ai-persistence' import type { LockStore } from '@tanstack/ai' +import type { D1InterruptStoreOptions } from '../src/index' declare const d1: D1Database declare const r2: R2Bucket @@ -22,6 +25,12 @@ declare const durableObjectState: DurableObjectState new CloudflareLockDurableObject(durableObjectState) +declare const d1InterruptOptions: D1InterruptStoreOptions +expectTypeOf( + createD1InterruptStore(d1, d1InterruptOptions), +).toEqualTypeOf() +expectTypeOf(createD1Stores(d1).interrupts).toEqualTypeOf() + expectTypeOf(cloudflarePersistence({}).stores).toEqualTypeOf<{}>() expectTypeOf(cloudflarePersistence({ d1 }).stores).toEqualTypeOf<{ messages: MessageStore diff --git a/packages/ai-persistence-cloudflare/tests/migration-cli.test.ts b/packages/ai-persistence-cloudflare/tests/migration-cli.test.ts index 5e5fd529b..b88f285db 100644 --- a/packages/ai-persistence-cloudflare/tests/migration-cli.test.ts +++ b/packages/ai-persistence-cloudflare/tests/migration-cli.test.ts @@ -30,19 +30,25 @@ describe('Cloudflare D1 migrations CLI', () => { }, }) - expect(output).toBe(`${d1Migrations[0]?.sql.trimEnd()}\n`) + expect(output).toBe( + `${d1Migrations.map((migration) => migration.sql.trimEnd()).join('\n\n')}\n`, + ) }) it('copies migrations without overwriting divergent files by default', async () => { const directory = await createTemporaryDirectory() await runCloudflareMigrationsCli(['--out', directory]) - const migration = d1Migrations[0] + for (const migration of d1Migrations) { + expect(await readFile(join(directory, migration.filename), 'utf8')).toBe( + migration.sql, + ) + } + + const migration = d1Migrations[1] expect(migration).toBeDefined() if (!migration) return const destination = join(directory, migration.filename) - expect(await readFile(destination, 'utf8')).toBe(migration.sql) - await writeFile(destination, 'user-owned contents', 'utf8') await expect( runCloudflareMigrationsCli(['--out', directory]), diff --git a/packages/ai-persistence-cloudflare/tests/migrations.test.ts b/packages/ai-persistence-cloudflare/tests/migrations.test.ts index 666e303fd..66109c1ff 100644 --- a/packages/ai-persistence-cloudflare/tests/migrations.test.ts +++ b/packages/ai-persistence-cloudflare/tests/migrations.test.ts @@ -1,11 +1,32 @@ import { readFile } from 'node:fs/promises' import { fileURLToPath } from 'node:url' import { describe, expect, it } from 'vitest' +import { Miniflare } from 'miniflare' import { d1Migrations } from '../src/index' +interface MigrationBindings { + AI_DB: D1Database +} + +function statementsFor(sql: string): Array { + return sql + .split('--> statement-breakpoint') + .map((statement) => statement.trim()) + .filter((statement) => statement.length > 0) +} + +async function applyMigration( + d1: D1Database, + migration: (typeof d1Migrations)[number], +): Promise { + await d1.batch( + statementsFor(migration.sql).map((statement) => d1.prepare(statement)), + ) +} + describe('D1 migrations', () => { it('exports the canonical structured-state migration', () => { - expect(d1Migrations).toHaveLength(1) + expect(d1Migrations).toHaveLength(2) expect(d1Migrations[0]).toMatchObject({ id: '0000_tanstack_ai_initial', filename: '0000_tanstack_ai_initial.sql', @@ -16,7 +37,12 @@ describe('D1 migrations', () => { expect(sql).toMatch(/CREATE TABLE [`"]interrupts[`"]/) expect(sql).toMatch(/CREATE TABLE [`"]metadata[`"]/) expect(sql).not.toMatch(/CREATE TABLE [`"]artifacts[`"]/) - expect(sql).not.toMatch(/CREATE TABLE [`"]blobs[`"]/) + expect(sql).not.toContain('CREATE TABLE `blobs`') + expect(d1Migrations[1]).toMatchObject({ + id: '0001_tanstack_ai_interrupt_batches', + filename: '0001_tanstack_ai_interrupt_batches.sql', + }) + expect(d1Migrations[1]?.sql).toContain('CREATE TABLE `interrupt_batches`') }) it('keeps the published migration equal to the embedded asset', async () => { @@ -34,5 +60,111 @@ describe('D1 migrations', () => { ) expect(published).toBe(embedded) expect(d1Migrations[0]?.sql).toBe(embedded) + + const interruptEmbedded = await readFile( + fileURLToPath( + new URL( + '../src/assets/0001_tanstack_ai_interrupt_batches.sql', + import.meta.url, + ), + ), + 'utf8', + ) + const interruptPublished = await readFile( + fileURLToPath( + new URL( + '../migrations/0001_tanstack_ai_interrupt_batches.sql', + import.meta.url, + ), + ), + 'utf8', + ) + const drizzlePublished = await readFile( + fileURLToPath( + new URL( + '../../ai-persistence-drizzle/drizzle/0001_tanstack_ai_interrupt_batches.sql', + import.meta.url, + ), + ), + 'utf8', + ) + expect(interruptPublished).not.toBe('') + expect(interruptPublished).toBe(interruptEmbedded) + expect(interruptPublished).toBe(drizzlePublished) + expect(d1Migrations[1]?.sql).toBe(interruptEmbedded) + }) + + it('backfills safe legacy identity for pending and completed rows', async () => { + const miniflare = new Miniflare({ + compatibilityDate: '2026-06-24', + d1Databases: ['AI_DB'], + modules: true, + script: 'export default { fetch() { return new Response("ok") } }', + }) + try { + const bindings = await miniflare.getBindings() + const initial = d1Migrations[0] + const interruptBatches = d1Migrations[1] + expect(initial).toBeDefined() + expect(interruptBatches).toBeDefined() + if (!initial || !interruptBatches) return + await applyMigration(bindings.AI_DB, initial) + await bindings.AI_DB.batch([ + bindings.AI_DB.prepare( + `INSERT INTO interrupts + (interrupt_id, run_id, thread_id, status, requested_at, payload_json) + VALUES (?, ?, ?, ?, ?, ?)`, + ).bind( + 'pending-int', + 'pending-run', + 'thread-legacy', + 'pending', + 1, + '{"id":"pending-int","reason":"confirmation"}', + ), + bindings.AI_DB.prepare( + `INSERT INTO interrupts + (interrupt_id, run_id, thread_id, status, requested_at, payload_json) + VALUES (?, ?, ?, ?, ?, ?)`, + ).bind( + 'resolved-int', + 'resolved-run', + 'thread-legacy', + 'resolved', + 2, + '{"id":"resolved-int","reason":"confirmation"}', + ), + ]) + + await applyMigration(bindings.AI_DB, interruptBatches) + + const rows = await bindings.AI_DB.prepare( + `SELECT interrupt_id, generation, binding_json, response_schema_hash + FROM interrupts ORDER BY interrupt_id`, + ).all<{ + interrupt_id: string + generation: number + binding_json: string + response_schema_hash: string + }>() + expect(rows.results).toEqual([ + { + interrupt_id: 'pending-int', + generation: 1, + binding_json: + '{"kind":"generic","interruptId":"pending-int","interruptedRunId":"pending-run","generation":1,"responseSchemaHash":"legacy:unknown"}', + response_schema_hash: 'legacy:unknown', + }, + { + interrupt_id: 'resolved-int', + generation: 0, + binding_json: + '{"kind":"generic","interruptId":"resolved-int","interruptedRunId":"resolved-run","generation":0,"responseSchemaHash":"legacy:unknown"}', + response_schema_hash: 'legacy:unknown', + }, + ]) + } finally { + await miniflare.dispose() + } }) }) diff --git a/packages/ai-persistence-cloudflare/tests/runtime.conformance.test.ts b/packages/ai-persistence-cloudflare/tests/runtime.conformance.test.ts index 9f9101250..70263a3f7 100644 --- a/packages/ai-persistence-cloudflare/tests/runtime.conformance.test.ts +++ b/packages/ai-persistence-cloudflare/tests/runtime.conformance.test.ts @@ -1,19 +1,42 @@ /// import { afterAll, beforeAll, describe, expect, it } from 'vitest' import { Miniflare } from 'miniflare' -import { runPersistenceConformance } from '@tanstack/ai-persistence/testkit' -import { cloudflarePersistence, d1Migrations } from '../src/index' +import { + runInterruptStoreConformance, + runPersistenceConformance, +} from '@tanstack/ai-persistence/testkit' import { composePersistence } from '@tanstack/ai-persistence' +import { cloudflarePersistence, d1Migrations } from '../src/index' +import { createD1InterruptStore } from '../src/d1' import type { AIPersistence, InterruptStore } from '@tanstack/ai-persistence' +import type { InterruptConformanceHarness } from '@tanstack/ai-persistence/testkit' interface RuntimeBindings { AI_DB: D1Database AI_BUCKET: R2Bucket } +function migrationStatements(sql: string): Array { + return sql + .split('--> statement-breakpoint') + .map((statement) => statement.trim()) + .filter((statement) => statement.length > 0) +} + +async function applyMigrations(d1: D1Database): Promise { + for (const migration of d1Migrations) { + await d1.batch( + migrationStatements(migration.sql).map((statement) => + d1.prepare(statement), + ), + ) + } +} + describe('Cloudflare persistence on Miniflare bindings', () => { let miniflare: Miniflare let persistence: AIPersistence + let d1: D1Database beforeAll(async () => { miniflare = new Miniflare({ @@ -24,15 +47,8 @@ describe('Cloudflare persistence on Miniflare bindings', () => { script: 'export default { fetch() { return new Response("ok") } }', }) const bindings = await miniflare.getBindings() - for (const migration of d1Migrations) { - const statements = migration.sql - .split('--> statement-breakpoint') - .map((statement) => statement.trim()) - .filter((statement) => statement.length > 0) - await bindings.AI_DB.batch( - statements.map((statement) => bindings.AI_DB.prepare(statement)), - ) - } + d1 = bindings.AI_DB + await applyMigrations(d1) persistence = cloudflarePersistence({ d1: bindings.AI_DB, r2: bindings.AI_BUCKET, @@ -45,16 +61,88 @@ describe('Cloudflare persistence on Miniflare bindings', () => { runPersistenceConformance('cloudflare-d1-r2', () => persistence) + runInterruptStoreConformance( + async (): Promise => { + await d1.exec('DROP TRIGGER IF EXISTS fail_interrupt_transition') + await d1.batch([ + d1.prepare('DELETE FROM interrupt_batches'), + d1.prepare('DELETE FROM interrupts'), + ]) + let now = Date.parse('2026-07-13T10:00:00.000Z') + let pendingFailureSetup: Promise | undefined + const createBase = () => + createD1InterruptStore(d1, { + clock: () => now, + }) + const base = createBase() + const store: InterruptStore = { + create: (record) => base.create(record), + resolve: (interruptId, response) => base.resolve(interruptId, response), + cancel: (interruptId) => base.cancel(interruptId), + get: (interruptId) => base.get(interruptId), + list: (threadId) => base.list(threadId), + listPending: (threadId) => base.listPending(threadId), + listByRun: (runId) => base.listByRun(runId), + listPendingByRun: (runId) => base.listPendingByRun(runId), + openInterruptBatch: (input) => base.openInterruptBatch(input), + async commitInterruptResolutions(input) { + if (pendingFailureSetup) { + await pendingFailureSetup + pendingFailureSetup = undefined + } + return base.commitInterruptResolutions(input) + }, + getInterruptRecoveryState: (input) => + base.getInterruptRecoveryState(input), + } + + return { + getStore: () => store, + advanceBy(milliseconds) { + now += milliseconds + }, + async inspect(interruptedRunId) { + const [rows, batch] = await Promise.all([ + d1 + .prepare( + 'SELECT status FROM interrupts WHERE run_id = ? ORDER BY interrupt_id', + ) + .bind(interruptedRunId) + .all<{ status: string }>(), + d1 + .prepare( + 'SELECT COUNT(*) AS count FROM interrupt_batches WHERE interrupted_run_id = ?', + ) + .bind(interruptedRunId) + .first('count'), + ]) + return { + statuses: rows.results.map((row) => row.status), + batchCount: batch ?? -1, + } + }, + failTransitionOnce(interruptId) { + const escapedInterruptId = interruptId.replaceAll("'", "''") + pendingFailureSetup = d1.exec(` + CREATE TRIGGER fail_interrupt_transition + BEFORE UPDATE OF status ON interrupts + WHEN OLD.interrupt_id = '${escapedInterruptId}' + AND NEW.status <> OLD.status + BEGIN + SELECT RAISE(ABORT, 'injected interrupt transition failure'); + END; + `) + }, + reopen: () => Promise.resolve(createBase()), + } + }, + ) + it('composes a custom interrupt store while retaining D1 runs', () => { + const baseInterrupts = persistence.stores.interrupts + if (!baseInterrupts) throw new Error('D1 interrupt store missing') const customInterrupts: InterruptStore = { - create: () => Promise.resolve(), - resolve: () => Promise.resolve(), - cancel: () => Promise.resolve(), - get: () => Promise.resolve(null), - list: () => Promise.resolve([]), - listPending: () => Promise.resolve([]), - listByRun: () => Promise.resolve([]), - listPendingByRun: () => Promise.resolve([]), + ...baseInterrupts, } const composed = composePersistence(persistence, { overrides: { interrupts: customInterrupts }, diff --git a/packages/ai-persistence-drizzle/drizzle/0001_tanstack_ai_interrupt_batches.sql b/packages/ai-persistence-drizzle/drizzle/0001_tanstack_ai_interrupt_batches.sql new file mode 100644 index 000000000..d81b74cee --- /dev/null +++ b/packages/ai-persistence-drizzle/drizzle/0001_tanstack_ai_interrupt_batches.sql @@ -0,0 +1,68 @@ +CREATE TABLE `interrupt_batches` ( + `interrupted_run_id` text PRIMARY KEY NOT NULL, + `thread_id` text NOT NULL, + `generation` integer NOT NULL, + `expected_interrupt_ids_json` text NOT NULL, + `fingerprint` text NOT NULL, + `canonical_resolutions` text NOT NULL, + `resolutions_json` text NOT NULL, + `continuation_run_id` text NOT NULL, + `committed_at` integer NOT NULL +); +--> statement-breakpoint +CREATE UNIQUE INDEX `interrupt_batches_continuation_run_id_unique` ON `interrupt_batches` (`continuation_run_id`); +--> statement-breakpoint +CREATE TABLE `__new_interrupts` ( + `interrupt_id` text PRIMARY KEY NOT NULL, + `run_id` text NOT NULL, + `thread_id` text NOT NULL, + `generation` integer NOT NULL, + `status` text NOT NULL, + `requested_at` integer NOT NULL, + `resolved_at` integer, + `payload_json` text NOT NULL, + `binding_json` text NOT NULL, + `response_schema_hash` text NOT NULL, + `response_json` text +); +--> statement-breakpoint +INSERT INTO `__new_interrupts` ( + `interrupt_id`, + `run_id`, + `thread_id`, + `generation`, + `status`, + `requested_at`, + `resolved_at`, + `payload_json`, + `binding_json`, + `response_schema_hash`, + `response_json` +) +SELECT + `interrupt_id`, + `run_id`, + `thread_id`, + CASE WHEN `status` = 'pending' THEN 1 ELSE 0 END, + `status`, + `requested_at`, + `resolved_at`, + `payload_json`, + json_object( + 'kind', 'generic', + 'interruptId', `interrupt_id`, + 'interruptedRunId', `run_id`, + 'generation', CASE WHEN `status` = 'pending' THEN 1 ELSE 0 END, + 'responseSchemaHash', 'legacy:unknown' + ), + 'legacy:unknown', + `response_json` +FROM `interrupts`; +--> statement-breakpoint +DROP TABLE `interrupts`; +--> statement-breakpoint +ALTER TABLE `__new_interrupts` RENAME TO `interrupts`; +--> statement-breakpoint +CREATE INDEX `interrupts_thread_status_requested_at_idx` ON `interrupts` (`thread_id`,`status`,`requested_at`); +--> statement-breakpoint +CREATE INDEX `interrupts_run_status_requested_at_idx` ON `interrupts` (`run_id`,`status`,`requested_at`); diff --git a/packages/ai-persistence-drizzle/drizzle/meta/0001_snapshot.json b/packages/ai-persistence-drizzle/drizzle/meta/0001_snapshot.json new file mode 100644 index 000000000..24a99472d --- /dev/null +++ b/packages/ai-persistence-drizzle/drizzle/meta/0001_snapshot.json @@ -0,0 +1,446 @@ +{ + "version": "6", + "dialect": "sqlite", + "id": "a87024ed-c5cd-41f2-aefd-fa31b9664b8d", + "prevId": "5e224922-d0ca-4cac-bf7b-5813ad6f6d5d", + "tables": { + "artifacts": { + "name": "artifacts", + "columns": { + "artifact_id": { + "name": "artifact_id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "thread_id": { + "name": "thread_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "name": { + "name": "name", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "mime_type": { + "name": "mime_type", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "size": { + "name": "size", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "external_url": { + "name": "external_url", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "created_at": { + "name": "created_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "blobs": { + "name": "blobs", + "columns": { + "key": { + "name": "key", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "content_type": { + "name": "content_type", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "size": { + "name": "size", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "etag": { + "name": "etag", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "custom_metadata_json": { + "name": "custom_metadata_json", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "created_at": { + "name": "created_at", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "updated_at": { + "name": "updated_at", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "body": { + "name": "body", + "type": "blob", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "interrupt_batches": { + "name": "interrupt_batches", + "columns": { + "interrupted_run_id": { + "name": "interrupted_run_id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "thread_id": { + "name": "thread_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "generation": { + "name": "generation", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "expected_interrupt_ids_json": { + "name": "expected_interrupt_ids_json", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "fingerprint": { + "name": "fingerprint", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "canonical_resolutions": { + "name": "canonical_resolutions", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "resolutions_json": { + "name": "resolutions_json", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "continuation_run_id": { + "name": "continuation_run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "committed_at": { + "name": "committed_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": { + "interrupt_batches_continuation_run_id_unique": { + "name": "interrupt_batches_continuation_run_id_unique", + "columns": ["continuation_run_id"], + "isUnique": true + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "interrupts": { + "name": "interrupts", + "columns": { + "interrupt_id": { + "name": "interrupt_id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "thread_id": { + "name": "thread_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "generation": { + "name": "generation", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "requested_at": { + "name": "requested_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "resolved_at": { + "name": "resolved_at", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "payload_json": { + "name": "payload_json", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "binding_json": { + "name": "binding_json", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "response_schema_hash": { + "name": "response_schema_hash", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "response_json": { + "name": "response_json", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": { + "interrupts_thread_status_requested_at_idx": { + "name": "interrupts_thread_status_requested_at_idx", + "columns": ["thread_id", "status", "requested_at"], + "isUnique": false + }, + "interrupts_run_status_requested_at_idx": { + "name": "interrupts_run_status_requested_at_idx", + "columns": ["run_id", "status", "requested_at"], + "isUnique": false + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "messages": { + "name": "messages", + "columns": { + "thread_id": { + "name": "thread_id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "messages_json": { + "name": "messages_json", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "metadata": { + "name": "metadata", + "columns": { + "scope": { + "name": "scope", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "key": { + "name": "key", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "value_json": { + "name": "value_json", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "metadata_scope_key_pk": { + "columns": ["scope", "key"], + "name": "metadata_scope_key_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "runs": { + "name": "runs", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "thread_id": { + "name": "thread_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "started_at": { + "name": "started_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "finished_at": { + "name": "finished_at", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "error": { + "name": "error", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "usage_json": { + "name": "usage_json", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + } + }, + "views": {}, + "enums": {}, + "_meta": { + "schemas": {}, + "tables": {}, + "columns": {} + }, + "internal": { + "indexes": {} + } +} diff --git a/packages/ai-persistence-drizzle/drizzle/meta/_journal.json b/packages/ai-persistence-drizzle/drizzle/meta/_journal.json index 7083785c2..da27e9e0f 100644 --- a/packages/ai-persistence-drizzle/drizzle/meta/_journal.json +++ b/packages/ai-persistence-drizzle/drizzle/meta/_journal.json @@ -8,6 +8,13 @@ "when": 1783594836281, "tag": "0000_tanstack_ai_initial", "breakpoints": true + }, + { + "idx": 1, + "version": "6", + "when": 1783960294769, + "tag": "0001_tanstack_ai_interrupt_batches", + "breakpoints": true } ] } diff --git a/packages/ai-persistence-drizzle/src/assets/0001_tanstack_ai_interrupt_batches.sql b/packages/ai-persistence-drizzle/src/assets/0001_tanstack_ai_interrupt_batches.sql new file mode 100644 index 000000000..d81b74cee --- /dev/null +++ b/packages/ai-persistence-drizzle/src/assets/0001_tanstack_ai_interrupt_batches.sql @@ -0,0 +1,68 @@ +CREATE TABLE `interrupt_batches` ( + `interrupted_run_id` text PRIMARY KEY NOT NULL, + `thread_id` text NOT NULL, + `generation` integer NOT NULL, + `expected_interrupt_ids_json` text NOT NULL, + `fingerprint` text NOT NULL, + `canonical_resolutions` text NOT NULL, + `resolutions_json` text NOT NULL, + `continuation_run_id` text NOT NULL, + `committed_at` integer NOT NULL +); +--> statement-breakpoint +CREATE UNIQUE INDEX `interrupt_batches_continuation_run_id_unique` ON `interrupt_batches` (`continuation_run_id`); +--> statement-breakpoint +CREATE TABLE `__new_interrupts` ( + `interrupt_id` text PRIMARY KEY NOT NULL, + `run_id` text NOT NULL, + `thread_id` text NOT NULL, + `generation` integer NOT NULL, + `status` text NOT NULL, + `requested_at` integer NOT NULL, + `resolved_at` integer, + `payload_json` text NOT NULL, + `binding_json` text NOT NULL, + `response_schema_hash` text NOT NULL, + `response_json` text +); +--> statement-breakpoint +INSERT INTO `__new_interrupts` ( + `interrupt_id`, + `run_id`, + `thread_id`, + `generation`, + `status`, + `requested_at`, + `resolved_at`, + `payload_json`, + `binding_json`, + `response_schema_hash`, + `response_json` +) +SELECT + `interrupt_id`, + `run_id`, + `thread_id`, + CASE WHEN `status` = 'pending' THEN 1 ELSE 0 END, + `status`, + `requested_at`, + `resolved_at`, + `payload_json`, + json_object( + 'kind', 'generic', + 'interruptId', `interrupt_id`, + 'interruptedRunId', `run_id`, + 'generation', CASE WHEN `status` = 'pending' THEN 1 ELSE 0 END, + 'responseSchemaHash', 'legacy:unknown' + ), + 'legacy:unknown', + `response_json` +FROM `interrupts`; +--> statement-breakpoint +DROP TABLE `interrupts`; +--> statement-breakpoint +ALTER TABLE `__new_interrupts` RENAME TO `interrupts`; +--> statement-breakpoint +CREATE INDEX `interrupts_thread_status_requested_at_idx` ON `interrupts` (`thread_id`,`status`,`requested_at`); +--> statement-breakpoint +CREATE INDEX `interrupts_run_status_requested_at_idx` ON `interrupts` (`run_id`,`status`,`requested_at`); diff --git a/packages/ai-persistence-drizzle/src/index.ts b/packages/ai-persistence-drizzle/src/index.ts index 0b2e29076..98578488b 100644 --- a/packages/ai-persistence-drizzle/src/index.ts +++ b/packages/ai-persistence-drizzle/src/index.ts @@ -16,25 +16,69 @@ import { createRunStore, } from './stores' import type { LockStore } from '@tanstack/ai' -import type { DrizzleDb } from './stores' +import type { DrizzleDb, DrizzleTransactionExecutor } from './stores' export { schema } from './schema' export { sqliteMigrations } from './migrations' export type { SqliteMigration } from './migrations' -export type { DrizzleDb } from './stores' +export type { DrizzleDb, DrizzleTransactionExecutor } from './stores' -/** Wire TanStack AI persistence stores over a migrated Drizzle SQLite database. */ -export function drizzlePersistence(db: DrizzleDb) { +export interface DrizzlePersistenceInterruptOptions { + /** Atomic transaction boundary required by the interrupt store. */ + interrupts: DrizzleTransactionExecutor + /** Override wall-clock time, primarily for deterministic runtimes/tests. */ + clock?: () => number +} + +export interface DrizzlePersistenceWithoutInterruptsOptions { + /** Explicitly omit interrupts while retaining all other Drizzle stores. */ + interrupts: false +} + +function createBaseStores(db: DrizzleDb) { const locks: LockStore = new InMemoryLockStore() + return { + messages: createMessageStore(db), + runs: createRunStore(db), + metadata: createMetadataStore(db), + artifacts: createArtifactStore(db), + blobs: createBlobStore(db), + locks, + } +} + +type DrizzleBaseStores = ReturnType +type DrizzleStoresWithInterrupts = DrizzleBaseStores & { + interrupts: ReturnType +} + +/** Wire TanStack AI persistence stores over a migrated Drizzle SQLite database. */ +export function drizzlePersistence( + db: DrizzleDb, + options: DrizzlePersistenceInterruptOptions, +): { stores: DrizzleStoresWithInterrupts } +export function drizzlePersistence( + db: DrizzleDb, + options: DrizzlePersistenceWithoutInterruptsOptions, +): { stores: DrizzleBaseStores } +export function drizzlePersistence(db: DrizzleDb): never +export function drizzlePersistence( + db: DrizzleDb, + options?: + | DrizzlePersistenceInterruptOptions + | DrizzlePersistenceWithoutInterruptsOptions, +): { stores: DrizzleBaseStores | DrizzleStoresWithInterrupts } { + if (!options) { + throw new Error( + 'Drizzle interrupts require an explicit transaction executor. Pass { interrupts: executor } or explicitly disable them with { interrupts: false }.', + ) + } + const stores = createBaseStores(db) + if (options.interrupts === false) return { stores } return { stores: { - messages: createMessageStore(db), - runs: createRunStore(db), - interrupts: createInterruptStore(db), - metadata: createMetadataStore(db), - artifacts: createArtifactStore(db), - blobs: createBlobStore(db), - locks, + ...stores, + interrupts: createInterruptStore(db, options.interrupts, options.clock), }, } } diff --git a/packages/ai-persistence-drizzle/src/migrations.ts b/packages/ai-persistence-drizzle/src/migrations.ts index 889f6426f..1df3a525c 100644 --- a/packages/ai-persistence-drizzle/src/migrations.ts +++ b/packages/ai-persistence-drizzle/src/migrations.ts @@ -1,4 +1,5 @@ import initialMigrationSql from './assets/0000_tanstack_ai_initial.sql?raw' +import interruptBatchesMigrationSql from './assets/0001_tanstack_ai_interrupt_batches.sql?raw' /** A canonical SQLite migration bundled with the adapter. */ export interface SqliteMigration { @@ -17,4 +18,9 @@ export const sqliteMigrations: ReadonlyArray = [ filename: '0000_tanstack_ai_initial.sql', sql: initialMigrationSql, }, + { + id: '0001_tanstack_ai_interrupt_batches', + filename: '0001_tanstack_ai_interrupt_batches.sql', + sql: interruptBatchesMigrationSql, + }, ] diff --git a/packages/ai-persistence-drizzle/src/schema.ts b/packages/ai-persistence-drizzle/src/schema.ts index b9cae92c7..ca9b02d99 100644 --- a/packages/ai-persistence-drizzle/src/schema.ts +++ b/packages/ai-persistence-drizzle/src/schema.ts @@ -16,13 +16,19 @@ */ import { customType, + index, integer, primaryKey, sqliteTable, text, + uniqueIndex, } from 'drizzle-orm/sqlite-core' -import type { InterruptRecord, RunStatus } from '@tanstack/ai-persistence' -import type { ModelMessage, TokenUsage } from '@tanstack/ai' +import type { + InterruptBatchRecord, + InterruptRecord, + RunStatus, +} from '@tanstack/ai-persistence' +import type { InterruptBinding, ModelMessage, TokenUsage } from '@tanstack/ai' const bytes = customType<{ data: Uint8Array @@ -62,18 +68,65 @@ export const runs = sqliteTable('runs', { }) /** Interrupt / approval records (`InterruptStore`). */ -export const interrupts = sqliteTable('interrupts', { - interruptId: text('interrupt_id').primaryKey(), - runId: text('run_id').notNull(), - threadId: text('thread_id').notNull(), - status: text('status').$type().notNull(), - requestedAt: integer('requested_at').notNull(), - resolvedAt: integer('resolved_at'), - payloadJson: text('payload_json', { mode: 'json' }) - .$type>() - .notNull(), - responseJson: text('response_json', { mode: 'json' }).$type(), -}) +export const interrupts = sqliteTable( + 'interrupts', + { + interruptId: text('interrupt_id').primaryKey(), + runId: text('run_id').notNull(), + threadId: text('thread_id').notNull(), + generation: integer('generation').notNull(), + status: text('status').$type().notNull(), + requestedAt: integer('requested_at').notNull(), + resolvedAt: integer('resolved_at'), + payloadJson: text('payload_json', { mode: 'json' }) + .$type() + .notNull(), + bindingJson: text('binding_json', { mode: 'json' }) + .$type() + .notNull(), + responseSchemaHash: text('response_schema_hash').notNull(), + responseJson: text('response_json', { mode: 'json' }).$type(), + }, + (table) => [ + index('interrupts_thread_status_requested_at_idx').on( + table.threadId, + table.status, + table.requestedAt, + ), + index('interrupts_run_status_requested_at_idx').on( + table.runId, + table.status, + table.requestedAt, + ), + ], +) + +/** Durable winner records for atomic interrupt resolution batches. */ +export const interruptBatches = sqliteTable( + 'interrupt_batches', + { + interruptedRunId: text('interrupted_run_id').primaryKey(), + threadId: text('thread_id').notNull(), + generation: integer('generation').notNull(), + expectedInterruptIdsJson: text('expected_interrupt_ids_json', { + mode: 'json', + }) + .$type() + .notNull(), + fingerprint: text('fingerprint').notNull(), + canonicalResolutions: text('canonical_resolutions').notNull(), + resolutionsJson: text('resolutions_json', { mode: 'json' }) + .$type() + .notNull(), + continuationRunId: text('continuation_run_id').notNull(), + committedAt: integer('committed_at').notNull(), + }, + (table) => [ + uniqueIndex('interrupt_batches_continuation_run_id_unique').on( + table.continuationRunId, + ), + ], +) /** Scoped key/value metadata (`MetadataStore`). */ export const metadata = sqliteTable( @@ -117,6 +170,7 @@ export const schema = { messages, runs, interrupts, + interruptBatches, metadata, artifacts, blobs, diff --git a/packages/ai-persistence-drizzle/src/sqlite.ts b/packages/ai-persistence-drizzle/src/sqlite.ts index 0774c3ff0..3a7c87bc0 100644 --- a/packages/ai-persistence-drizzle/src/sqlite.ts +++ b/packages/ai-persistence-drizzle/src/sqlite.ts @@ -7,6 +7,7 @@ import { drizzle } from 'drizzle-orm/sqlite-proxy' import { sqliteMigrations } from './migrations' import { applySqliteMigrations } from './sqlite-migrations' import { drizzlePersistence } from './index' +import type { DrizzleTransactionExecutor } from './index' export { SqliteMigrationError } from './sqlite-migrations' @@ -15,6 +16,41 @@ export interface SqlitePersistenceOptions { url: string /** Apply the bundled TanStack AI migrations before creating stores. */ migrate?: boolean + /** Override wall-clock time, primarily for deterministic runtimes/tests. */ + clock?: () => number +} + +function createSqliteTransactionExecutor( + database: DatabaseSync, +): DrizzleTransactionExecutor { + let tail = Promise.resolve() + return { + transaction(work) { + const result = tail.then(async () => { + database.exec('BEGIN IMMEDIATE') + try { + const value = await work() + database.exec('COMMIT') + return value + } catch (error) { + try { + database.exec('ROLLBACK') + } catch (rollbackError) { + throw new AggregateError( + [error, rollbackError], + 'Interrupt transaction failed and could not be rolled back.', + ) + } + throw error + } + }) + tail = result.then( + () => undefined, + () => undefined, + ) + return result + }, + } } /** Build persistence over Node's built-in SQLite driver. */ @@ -43,7 +79,11 @@ export function sqlitePersistence(options: SqlitePersistenceOptions) { return Promise.resolve({ rows: rows.map((row) => Object.values(row)) }) }) - const persistence = drizzlePersistence(db) + const transactionExecutor = createSqliteTransactionExecutor(sqlite) + const persistence = drizzlePersistence(db, { + interrupts: transactionExecutor, + ...(options.clock !== undefined ? { clock: options.clock } : {}), + }) let closed = false return { ...persistence, diff --git a/packages/ai-persistence-drizzle/src/stores.ts b/packages/ai-persistence-drizzle/src/stores.ts index fcaa30dcd..3d00e99c4 100644 --- a/packages/ai-persistence-drizzle/src/stores.ts +++ b/packages/ai-persistence-drizzle/src/stores.ts @@ -5,10 +5,21 @@ * (`@tanstack/ai-persistence`'s `memory.ts`). JSON columns are handled by * Drizzle's `text({ mode: 'json' })`; blob bytes by `blob({ mode: 'buffer' })`. */ +import { + canonicalInterruptJson, + canonicalizeInterruptResolutions, + cloneAndDeepFreezeJson, +} from '@tanstack/ai' import { and, asc, eq, gt, gte, lt } from 'drizzle-orm' +import { + InterruptStoreCorruptionError, + hasExactInterruptIds, + projectInterruptRecovery, +} from '@tanstack/ai-persistence' import { artifacts, blobs, + interruptBatches, interrupts, messages, metadata, @@ -16,7 +27,12 @@ import { } from './schema' import type { SQL } from 'drizzle-orm' import type { BaseSQLiteDatabase } from 'drizzle-orm/sqlite-core' -import type { ModelMessage } from '@tanstack/ai' +import type { + InterruptBinding, + InterruptRecoveryQuery, + ModelMessage, + RunAgentResumeItem, +} from '@tanstack/ai' import type { ArtifactRecord, ArtifactStore, @@ -26,6 +42,7 @@ import type { BlobObject, BlobRecord, BlobStore, + InterruptBatchRecord, InterruptRecord, InterruptStore, MessageStore, @@ -46,6 +63,11 @@ export type DrizzleDb = Pick< 'select' | 'insert' | 'update' | 'delete' > +/** Required atomic transaction boundary for Drizzle interrupt operations. */ +export interface DrizzleTransactionExecutor { + transaction: (work: () => Promise) => Promise +} + const textEncoder = new TextEncoder() const textDecoder = new TextDecoder() @@ -193,49 +215,329 @@ export function createRunStore(db: DrizzleDb): RunStore { return store } +function corrupt(message: string): never { + throw new InterruptStoreCorruptionError(message) +} + +function decodeBinding( + value: unknown, + correlation: { + interruptId: string + interruptedRunId: string + generation: number + responseSchemaHash: string + }, +): InterruptBinding { + if ( + value === null || + typeof value !== 'object' || + !('kind' in value) || + !('interruptId' in value) || + !('interruptedRunId' in value) || + !('generation' in value) || + !('responseSchemaHash' in value) || + typeof value.interruptId !== 'string' || + typeof value.interruptedRunId !== 'string' || + typeof value.generation !== 'number' || + typeof value.responseSchemaHash !== 'string' + ) { + return corrupt('Stored interrupt binding is malformed.') + } + const expiresAt = 'expiresAt' in value ? value.expiresAt : undefined + if ( + expiresAt !== undefined && + (typeof expiresAt !== 'string' || !Number.isFinite(Date.parse(expiresAt))) + ) { + return corrupt('Stored interrupt expiry is malformed.') + } + if ( + value.interruptId !== correlation.interruptId || + value.interruptedRunId !== correlation.interruptedRunId || + value.generation !== correlation.generation || + value.responseSchemaHash !== correlation.responseSchemaHash + ) { + return corrupt('Stored interrupt binding correlation is inconsistent.') + } + + if (value.kind === 'generic') { + return { + kind: 'generic', + interruptId: value.interruptId, + interruptedRunId: value.interruptedRunId, + generation: value.generation, + responseSchemaHash: value.responseSchemaHash, + ...(expiresAt !== undefined ? { expiresAt } : {}), + } + } + if ( + value.kind === 'client-tool-execution' && + 'toolName' in value && + 'toolCallId' in value && + 'outputSchemaHash' in value && + typeof value.toolName === 'string' && + typeof value.toolCallId === 'string' && + typeof value.outputSchemaHash === 'string' + ) { + return { + kind: 'client-tool-execution', + interruptId: value.interruptId, + interruptedRunId: value.interruptedRunId, + generation: value.generation, + toolName: value.toolName, + toolCallId: value.toolCallId, + outputSchemaHash: value.outputSchemaHash, + responseSchemaHash: value.responseSchemaHash, + ...(expiresAt !== undefined ? { expiresAt } : {}), + } + } + if ( + value.kind === 'tool-approval' && + 'toolName' in value && + 'toolCallId' in value && + 'originalArgs' in value && + 'inputSchemaHash' in value && + 'approvalSchemaHash' in value && + typeof value.toolName === 'string' && + typeof value.toolCallId === 'string' && + typeof value.inputSchemaHash === 'string' && + typeof value.approvalSchemaHash === 'string' + ) { + return { + kind: 'tool-approval', + interruptId: value.interruptId, + interruptedRunId: value.interruptedRunId, + generation: value.generation, + toolName: value.toolName, + toolCallId: value.toolCallId, + originalArgs: value.originalArgs, + inputSchemaHash: value.inputSchemaHash, + approvalSchemaHash: value.approvalSchemaHash, + responseSchemaHash: value.responseSchemaHash, + ...(expiresAt !== undefined ? { expiresAt } : {}), + } + } + return corrupt('Stored interrupt binding kind is malformed.') +} + +function validateNativeDescriptor(value: unknown, interruptId: string): void { + if ( + value === null || + typeof value !== 'object' || + !('id' in value) || + !('reason' in value) || + value.id !== interruptId || + typeof value.reason !== 'string' + ) { + corrupt('Stored interrupt descriptor correlation is inconsistent.') + } +} + function mapInterrupt(row: typeof interrupts.$inferSelect): InterruptRecord { + const binding = decodeBinding(row.bindingJson, { + interruptId: row.interruptId, + interruptedRunId: row.runId, + generation: row.generation, + responseSchemaHash: row.responseSchemaHash, + }) + if (row.generation > 0) { + validateNativeDescriptor(row.payloadJson, row.interruptId) + } return { interruptId: row.interruptId, runId: row.runId, threadId: row.threadId, + generation: row.generation, status: row.status, requestedAt: row.requestedAt, ...(row.resolvedAt != null ? { resolvedAt: row.resolvedAt } : {}), payload: row.payloadJson, + binding, ...(row.responseJson != null ? { response: row.responseJson } : {}), } } -export function createInterruptStore(db: DrizzleDb): InterruptStore { +function decodeStringArray(value: unknown, field: string): Array { + if (!Array.isArray(value)) return corrupt(`Stored ${field} is malformed.`) + const strings: Array = [] + for (const item of value) { + if (typeof item !== 'string') { + return corrupt(`Stored ${field} is malformed.`) + } + strings.push(item) + } + if (!hasExactInterruptIds(strings, strings)) { + return corrupt(`Stored ${field} contains duplicate IDs.`) + } + return strings +} + +function decodeResolution(value: unknown): RunAgentResumeItem { + if ( + value === null || + typeof value !== 'object' || + !('interruptId' in value) || + !('status' in value) || + typeof value.interruptId !== 'string' || + (value.status !== 'resolved' && value.status !== 'cancelled') + ) { + return corrupt('Stored interrupt resolution is malformed.') + } return { + interruptId: value.interruptId, + status: value.status, + ...('payload' in value ? { payload: value.payload } : {}), + } +} + +function decodeResolutions(value: unknown): Array { + if (!Array.isArray(value)) { + return corrupt('Stored interrupt resolutions are malformed.') + } + return value.map(decodeResolution) +} + +function mapInterruptBatch( + row: typeof interruptBatches.$inferSelect, +): InterruptBatchRecord { + const expectedInterruptIds = decodeStringArray( + row.expectedInterruptIdsJson, + 'expected interrupt IDs', + ) + const resolutions = decodeResolutions(row.resolutionsJson) + const candidate = canonicalizeInterruptResolutions(resolutions) + if ( + !hasExactInterruptIds( + expectedInterruptIds, + resolutions.map((resolution) => resolution.interruptId), + ) || + candidate.fingerprint !== row.fingerprint || + candidate.canonicalResolutions !== row.canonicalResolutions + ) { + return corrupt('Stored interrupt batch identity is inconsistent.') + } + return { + threadId: row.threadId, + interruptedRunId: row.interruptedRunId, + generation: row.generation, + expectedInterruptIds, + fingerprint: row.fingerprint, + canonicalResolutions: row.canonicalResolutions, + resolutions, + continuationRunId: row.continuationRunId, + committedAt: row.committedAt, + } +} + +export function createInterruptStore( + db: DrizzleDb, + transactionExecutor: DrizzleTransactionExecutor, + clock: () => number = Date.now, +): InterruptStore { + const rowsForRun = async (runId: string) => + ( + await db + .select() + .from(interrupts) + .where(eq(interrupts.runId, runId)) + .orderBy(asc(interrupts.requestedAt), asc(interrupts.interruptId)) + ).map(mapInterrupt) + + const readBatch = async ( + interruptedRunId: string, + ): Promise => { + const rows = await db + .select() + .from(interruptBatches) + .where(eq(interruptBatches.interruptedRunId, interruptedRunId)) + return rows[0] ? mapInterruptBatch(rows[0]) : null + } + + const recovery = async (input: InterruptRecoveryQuery) => { + const rows = (await rowsForRun(input.interruptedRunId)).filter( + (row) => row.threadId === input.threadId, + ) + const batch = await readBatch(input.interruptedRunId) + const correlatedBatch = batch?.threadId === input.threadId ? batch : null + const generations = new Set(rows.map((row) => row.generation)) + if (generations.size > 1) { + return corrupt('Stored interrupt rows mix generations.') + } + if ( + correlatedBatch && + (rows.length === 0 || + rows.some( + (row) => + row.generation !== correlatedBatch.generation || + row.runId !== correlatedBatch.interruptedRunId, + ) || + !hasExactInterruptIds( + correlatedBatch.expectedInterruptIds, + rows.map((row) => row.interruptId), + )) + ) { + return corrupt('Stored interrupt batch correlation is inconsistent.') + } + return projectInterruptRecovery({ + query: input, + rows, + batch: correlatedBatch, + now: clock(), + includeResolutions: true, + }) + } + + const store: InterruptStore = { async create(record) { + const generation = record.generation ?? 0 + const binding = decodeBinding( + cloneAndDeepFreezeJson( + record.binding ?? { + kind: 'generic', + interruptId: record.interruptId, + interruptedRunId: record.runId, + generation, + responseSchemaHash: 'legacy:unknown', + }, + ), + { + interruptId: record.interruptId, + interruptedRunId: record.runId, + generation, + responseSchemaHash: + record.binding?.responseSchemaHash ?? 'legacy:unknown', + }, + ) await db .insert(interrupts) .values({ interruptId: record.interruptId, runId: record.runId, threadId: record.threadId, + generation, status: record.status, requestedAt: record.requestedAt, - payloadJson: record.payload, + payloadJson: cloneAndDeepFreezeJson(record.payload), + bindingJson: binding, + responseSchemaHash: binding.responseSchemaHash, responseJson: record.response ?? null, }) .onConflictDoNothing({ target: interrupts.interruptId }) }, async resolve(interruptId, response) { + const set: Partial = { + status: 'resolved', + resolvedAt: clock(), + } + if (response !== undefined) set.responseJson = response await db .update(interrupts) - .set({ - status: 'resolved', - resolvedAt: Date.now(), - responseJson: response ?? null, - }) + .set(set) .where(eq(interrupts.interruptId, interruptId)) }, async cancel(interruptId) { await db .update(interrupts) - .set({ status: 'cancelled', resolvedAt: Date.now() }) + .set({ status: 'cancelled', resolvedAt: clock() }) .where(eq(interrupts.interruptId, interruptId)) }, async get(interruptId) { @@ -268,12 +570,7 @@ export function createInterruptStore(db: DrizzleDb): InterruptStore { return rows.map(mapInterrupt) }, async listByRun(runId) { - const rows = await db - .select() - .from(interrupts) - .where(eq(interrupts.runId, runId)) - .orderBy(asc(interrupts.requestedAt)) - return rows.map(mapInterrupt) + return rowsForRun(runId) }, async listPendingByRun(runId) { const rows = await db @@ -285,7 +582,286 @@ export function createInterruptStore(db: DrizzleDb): InterruptStore { .orderBy(asc(interrupts.requestedAt)) return rows.map(mapInterrupt) }, + async openInterruptBatch(input) { + const descriptors = cloneAndDeepFreezeJson(input.descriptors) + const unopenedBindings = cloneAndDeepFreezeJson(input.bindings) + const descriptorIds = descriptors.map((descriptor) => descriptor.id) + const bindingIds = unopenedBindings.map((binding) => binding.interruptId) + if ( + descriptorIds.length === 0 || + !hasExactInterruptIds(descriptorIds, bindingIds) + ) { + throw new TypeError( + 'Interrupt descriptors and bindings must contain the same exact nonempty IDs.', + ) + } + for (const descriptor of descriptors) { + validateNativeDescriptor(descriptor, descriptor.id) + } + + return transactionExecutor.transaction(async () => { + if (await readBatch(input.interruptedRunId)) { + throw new Error('Cannot reopen a committed interrupt batch.') + } + const existing = await rowsForRun(input.interruptedRunId) + if (existing.length > 0) { + const generation = existing[0]?.generation + const requestedBindings = unopenedBindings + .map((binding) => ({ + ...binding, + interruptedRunId: input.interruptedRunId, + generation, + })) + .sort((left, right) => + left.interruptId.localeCompare(right.interruptId), + ) + const existingBindings = existing + .map((row) => row.binding) + .sort((left, right) => + left.interruptId.localeCompare(right.interruptId), + ) + const existingDescriptors = existing + .map((row) => row.payload) + .sort((left, right) => { + if ( + left !== null && + typeof left === 'object' && + 'id' in left && + typeof left.id === 'string' && + right !== null && + typeof right === 'object' && + 'id' in right && + typeof right.id === 'string' + ) { + return left.id.localeCompare(right.id) + } + return 0 + }) + const requestedDescriptors = [...descriptors].sort((left, right) => + left.id.localeCompare(right.id), + ) + if ( + generation !== undefined && + existing.every( + (row) => + row.threadId === input.threadId && + row.generation === generation && + row.status === 'pending', + ) && + hasExactInterruptIds( + descriptorIds, + existing.map((row) => row.interruptId), + ) && + canonicalInterruptJson(existingDescriptors) === + canonicalInterruptJson(requestedDescriptors) && + canonicalInterruptJson(existingBindings) === + canonicalInterruptJson(requestedBindings) + ) { + return { generation, descriptors } + } + throw new Error('An incompatible interrupt batch is already open.') + } + + const generation = 1 + const byId = new Map( + unopenedBindings.map((binding) => [binding.interruptId, binding]), + ) + const requestedAt = clock() + const values: Array = [] + for (const descriptor of descriptors) { + const unopened = byId.get(descriptor.id) + if (!unopened) { + throw new Error(`Missing binding for interrupt: ${descriptor.id}`) + } + const proposedBinding = cloneAndDeepFreezeJson({ + ...unopened, + interruptedRunId: input.interruptedRunId, + generation, + }) + const binding = decodeBinding(proposedBinding, { + interruptId: descriptor.id, + interruptedRunId: input.interruptedRunId, + generation, + responseSchemaHash: proposedBinding.responseSchemaHash, + }) + values.push({ + interruptId: descriptor.id, + runId: input.interruptedRunId, + threadId: input.threadId, + generation, + status: 'pending', + requestedAt, + payloadJson: descriptor, + bindingJson: binding, + responseSchemaHash: binding.responseSchemaHash, + }) + } + await db.insert(interrupts).values(values).onConflictDoNothing() + const inserted = await rowsForRun(input.interruptedRunId) + if ( + !hasExactInterruptIds( + descriptorIds, + inserted.map((row) => row.interruptId), + ) || + inserted.some( + (row) => + row.threadId !== input.threadId || + row.generation !== generation || + row.status !== 'pending', + ) + ) { + throw new Error('Interrupt IDs conflict with an existing batch.') + } + return { generation, descriptors } + }) + }, + async commitInterruptResolutions(input) { + return transactionExecutor.transaction(async () => { + const candidate = canonicalizeInterruptResolutions(input.resolutions) + if ( + candidate.fingerprint !== input.fingerprint || + candidate.canonicalResolutions !== input.canonicalResolutions + ) { + throw new TypeError( + 'Interrupt batch identity does not match its resolutions.', + ) + } + + const existingWinner = await readBatch(input.interruptedRunId) + if (existingWinner) { + return existingWinner.threadId === input.threadId && + existingWinner.fingerprint === input.fingerprint && + existingWinner.canonicalResolutions === input.canonicalResolutions + ? { + status: 'replayed', + continuationRunId: existingWinner.continuationRunId, + } + : { + status: 'conflict', + authoritativeState: await recovery({ + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + knownGeneration: input.expectedGeneration, + }), + } + } + + const pending = (await rowsForRun(input.interruptedRunId)).filter( + (row) => row.status === 'pending', + ) + const resolutionIds = candidate.resolutions.map( + (resolution) => resolution.interruptId, + ) + if ( + !hasExactInterruptIds( + input.expectedInterruptIds, + pending.map((row) => row.interruptId), + ) || + !hasExactInterruptIds(input.expectedInterruptIds, resolutionIds) || + pending.some( + (row) => + row.threadId !== input.threadId || + row.generation !== input.expectedGeneration || + (row.binding.expiresAt !== undefined && + Date.parse(row.binding.expiresAt) <= clock()), + ) + ) { + return { + status: 'conflict', + authoritativeState: await recovery({ + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + knownGeneration: input.expectedGeneration, + }), + } + } + + const committedAt = clock() + await db + .insert(interruptBatches) + .values({ + interruptedRunId: input.interruptedRunId, + threadId: input.threadId, + generation: input.expectedGeneration, + expectedInterruptIdsJson: cloneAndDeepFreezeJson( + input.expectedInterruptIds, + ), + fingerprint: candidate.fingerprint, + canonicalResolutions: candidate.canonicalResolutions, + resolutionsJson: candidate.resolutions, + continuationRunId: input.continuationRunId, + committedAt, + }) + .onConflictDoNothing({ target: interruptBatches.interruptedRunId }) + + const winner = await readBatch(input.interruptedRunId) + if (!winner) { + return corrupt('Interrupt winner insert did not persist a row.') + } + if ( + winner.threadId !== input.threadId || + winner.fingerprint !== input.fingerprint || + winner.canonicalResolutions !== input.canonicalResolutions + ) { + return { + status: 'conflict', + authoritativeState: await recovery({ + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + knownGeneration: input.expectedGeneration, + }), + } + } + if (winner.continuationRunId !== input.continuationRunId) { + return { + status: 'replayed', + continuationRunId: winner.continuationRunId, + } + } + + for (const resolution of candidate.resolutions) { + const status = + resolution.status === 'cancelled' ? 'cancelled' : 'resolved' + await db + .update(interrupts) + .set({ + status, + resolvedAt: committedAt, + responseJson: resolution, + }) + .where( + and( + eq(interrupts.interruptId, resolution.interruptId), + eq(interrupts.runId, input.interruptedRunId), + eq(interrupts.threadId, input.threadId), + eq(interrupts.generation, input.expectedGeneration), + eq(interrupts.status, 'pending'), + ), + ) + const transitioned = await store.get(resolution.interruptId) + if ( + !transitioned || + transitioned.status !== status || + transitioned.resolvedAt !== committedAt || + canonicalInterruptJson(transitioned.response) !== + canonicalInterruptJson(resolution) + ) { + throw new Error( + `Pending interrupt changed: ${resolution.interruptId}`, + ) + } + } + return { + status: 'committed', + continuationRunId: input.continuationRunId, + } + }) + }, + async getInterruptRecoveryState(input) { + return transactionExecutor.transaction(() => recovery(input)) + }, } + return store } export function createMetadataStore(db: DrizzleDb): MetadataStore { diff --git a/packages/ai-persistence-drizzle/tests/api-types.test-d.ts b/packages/ai-persistence-drizzle/tests/api-types.test-d.ts index 781c3993c..3d121913c 100644 --- a/packages/ai-persistence-drizzle/tests/api-types.test-d.ts +++ b/packages/ai-persistence-drizzle/tests/api-types.test-d.ts @@ -1,4 +1,6 @@ import { expectTypeOf } from 'vitest' +import { drizzlePersistence } from '../src/index' +import { sqlitePersistence } from '../src/sqlite' import type { DrizzleD1Database } from 'drizzle-orm/d1' import type { ArtifactStore, @@ -9,19 +11,33 @@ import type { RunStore, } from '@tanstack/ai-persistence' import type { LockStore } from '@tanstack/ai' -import { drizzlePersistence } from '../src/index' -import { sqlitePersistence } from '../src/sqlite' +import type { DrizzleTransactionExecutor } from '../src/index' declare const d1Database: DrizzleD1Database -const d1Persistence = drizzlePersistence(d1Database) -expectTypeOf(d1Persistence.stores.messages).toEqualTypeOf() -expectTypeOf(d1Persistence.stores.runs).toEqualTypeOf() -expectTypeOf(d1Persistence.stores.interrupts).toEqualTypeOf() -expectTypeOf(d1Persistence.stores.metadata).toEqualTypeOf() -expectTypeOf(d1Persistence.stores.artifacts).toEqualTypeOf() -expectTypeOf(d1Persistence.stores.blobs).toEqualTypeOf() -expectTypeOf(d1Persistence.stores.locks).toEqualTypeOf() +declare const transactionExecutor: DrizzleTransactionExecutor +const genericPersistence = drizzlePersistence(d1Database, { + interrupts: transactionExecutor, +}) +expectTypeOf(genericPersistence.stores.messages).toEqualTypeOf() +expectTypeOf(genericPersistence.stores.runs).toEqualTypeOf() +expectTypeOf( + genericPersistence.stores.interrupts, +).toEqualTypeOf() +expectTypeOf(genericPersistence.stores.metadata).toEqualTypeOf() +expectTypeOf(genericPersistence.stores.artifacts).toEqualTypeOf() +expectTypeOf(genericPersistence.stores.blobs).toEqualTypeOf() +expectTypeOf(genericPersistence.stores.locks).toEqualTypeOf() + +const nonInterruptPersistence = drizzlePersistence(d1Database, { + interrupts: false, +}) +expectTypeOf(nonInterruptPersistence.stores.runs).toEqualTypeOf() +// @ts-expect-error an explicit non-interrupt factory has no interrupt store +nonInterruptPersistence.stores.interrupts + +const missingExecutor = drizzlePersistence(d1Database) +expectTypeOf(missingExecutor).toEqualTypeOf() const sqlite = sqlitePersistence({ url: ':memory:', migrate: true }) -expectTypeOf(sqlite.stores).toEqualTypeOf() +expectTypeOf(sqlite.stores).toEqualTypeOf() expectTypeOf(sqlite.close).toEqualTypeOf<() => void>() diff --git a/packages/ai-persistence-drizzle/tests/drizzle.conformance.test.ts b/packages/ai-persistence-drizzle/tests/drizzle.conformance.test.ts index caaef2ee3..ed9e77a9a 100644 --- a/packages/ai-persistence-drizzle/tests/drizzle.conformance.test.ts +++ b/packages/ai-persistence-drizzle/tests/drizzle.conformance.test.ts @@ -1,6 +1,132 @@ -import { runPersistenceConformance } from '@tanstack/ai-persistence/testkit' +import { mkdtempSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { DatabaseSync } from 'node:sqlite' +import { afterAll, describe, expect, it } from 'vitest' +import { drizzle } from 'drizzle-orm/sqlite-proxy' +import { + runInterruptStoreConformance, + runPersistenceConformance, +} from '@tanstack/ai-persistence/testkit' +import { drizzlePersistence } from '../src/index' import { sqlitePersistence } from '../src/sqlite' +import type { InterruptConformanceHarness } from '@tanstack/ai-persistence/testkit' + +const cleanup: Array<() => void> = [] + +function createNativeDrizzleDatabase(database: DatabaseSync) { + return drizzle((sql, params, method) => { + const statement = database.prepare(sql) + if (method === 'run') { + statement.run(...params) + return Promise.resolve({ rows: [] }) + } + if (method === 'get') { + const row = statement.get(...params) + return Promise.resolve({ rows: row ? Object.values(row) : [] }) + } + const rows = statement.all(...params) + return Promise.resolve({ rows: rows.map((row) => Object.values(row)) }) + }) +} + +function createSqliteHarness(): Promise { + const directory = mkdtempSync(join(tmpdir(), 'tanstack-ai-drizzle-cas-')) + const filename = join(directory, 'state.sqlite') + let now = Date.parse('2026-07-13T10:00:00.000Z') + const openConnections: Array> = [] + + const open = () => { + const persistence = sqlitePersistence({ + url: filename, + migrate: true, + clock: () => now, + }) + openConnections.push(persistence) + return persistence + } + const persistence = open() + + cleanup.push(() => { + for (const connection of openConnections.splice(0)) connection.close() + rmSync(directory, { recursive: true, force: true }) + }) + + return Promise.resolve({ + getStore: () => persistence.stores.interrupts, + advanceBy(milliseconds) { + now += milliseconds + }, + async inspect(interruptedRunId) { + const statuses = ( + await persistence.stores.interrupts.listByRun(interruptedRunId) + ).map((row) => row.status) + const database = new DatabaseSync(filename) + try { + const batch = database + .prepare( + 'SELECT COUNT(*) AS count FROM interrupt_batches WHERE interrupted_run_id = ?', + ) + .get(interruptedRunId) + return { + statuses, + batchCount: typeof batch?.count === 'number' ? batch.count : -1, + } + } finally { + database.close() + } + }, + failTransitionOnce(interruptId) { + const database = new DatabaseSync(filename) + try { + database.exec(` + CREATE TRIGGER fail_interrupt_transition + BEFORE UPDATE OF status ON interrupts + WHEN OLD.interrupt_id = '${interruptId}' AND NEW.status <> OLD.status + BEGIN + SELECT RAISE(ABORT, 'injected interrupt transition failure'); + END; + `) + } finally { + database.close() + } + }, + reopen: () => Promise.resolve(open().stores.interrupts), + }) +} runPersistenceConformance('drizzle-sqlite', () => sqlitePersistence({ url: ':memory:', migrate: true }), ) + +runInterruptStoreConformance(createSqliteHarness) + +describe('drizzlePersistence interrupt transactions', () => { + it('rejects an implicit non-atomic interrupt store', () => { + const database = new DatabaseSync(':memory:') + cleanup.push(() => database.close()) + + expect(() => + drizzlePersistence(createNativeDrizzleDatabase(database)), + ).toThrow(/interrupt.*transaction executor/i) + }) + + it('allows callers to explicitly omit interrupts', () => { + const database = new DatabaseSync(':memory:') + cleanup.push(() => database.close()) + + const persistence = drizzlePersistence( + createNativeDrizzleDatabase(database), + { + interrupts: false, + }, + ) + + expect('interrupts' in persistence.stores).toBe(false) + expect(persistence.stores.runs).toBeDefined() + }) +}) + +afterAll(() => { + for (const dispose of cleanup.splice(0).reverse()) dispose() +}) diff --git a/packages/ai-persistence-drizzle/tests/migration-cli.test.ts b/packages/ai-persistence-drizzle/tests/migration-cli.test.ts index 2743a2cbb..e79dcb5ed 100644 --- a/packages/ai-persistence-drizzle/tests/migration-cli.test.ts +++ b/packages/ai-persistence-drizzle/tests/migration-cli.test.ts @@ -30,18 +30,25 @@ describe('drizzle migrations CLI', () => { }, }) - expect(output).toBe(`${sqliteMigrations[0]?.sql.trimEnd()}\n`) + expect(output).toBe( + `${sqliteMigrations.map((migration) => migration.sql.trimEnd()).join('\n\n')}\n`, + ) }) it('copies migrations without overwriting divergent files by default', async () => { const directory = await createTemporaryDirectory() await runDrizzleMigrationsCli(['--out', directory]) - const migration = sqliteMigrations[0] + for (const migration of sqliteMigrations) { + expect(await readFile(join(directory, migration.filename), 'utf8')).toBe( + migration.sql, + ) + } + + const migration = sqliteMigrations[1] expect(migration).toBeDefined() if (!migration) return const destination = join(directory, migration.filename) - expect(await readFile(destination, 'utf8')).toBe(migration.sql) await writeFile(destination, 'user-owned contents', 'utf8') await expect(runDrizzleMigrationsCli(['--out', directory])).rejects.toThrow( diff --git a/packages/ai-persistence-drizzle/tests/migrations.test.ts b/packages/ai-persistence-drizzle/tests/migrations.test.ts index 0e9d92945..d96e7ac7c 100644 --- a/packages/ai-persistence-drizzle/tests/migrations.test.ts +++ b/packages/ai-persistence-drizzle/tests/migrations.test.ts @@ -1,4 +1,5 @@ import { DatabaseSync } from 'node:sqlite' +import { readFileSync } from 'node:fs' import { afterEach, describe, expect, it } from 'vitest' import { sqliteMigrations } from '../src/migrations' import { applySqliteMigrations } from '../src/sqlite-migrations' @@ -19,6 +20,7 @@ describe('sqlite migrations', () => { it('exposes an ordered canonical manifest and applies it idempotently', () => { expect(sqliteMigrations.map((migration) => migration.id)).toEqual([ '0000_tanstack_ai_initial', + '0001_tanstack_ai_interrupt_batches', ]) const database = createDatabase() @@ -37,6 +39,7 @@ describe('sqlite migrations', () => { 'artifacts', 'blobs', 'interrupts', + 'interrupt_batches', 'messages', 'metadata', 'runs', @@ -46,7 +49,100 @@ describe('sqlite migrations', () => { database .prepare('SELECT migration_id FROM __tanstack_ai_migrations') .all(), - ).toEqual([{ migration_id: '0000_tanstack_ai_initial' }]) + ).toEqual([ + { migration_id: '0000_tanstack_ai_initial' }, + { migration_id: '0001_tanstack_ai_interrupt_batches' }, + ]) + }) + + it('ships byte-identical nonempty generated and runtime migration assets', () => { + const migration = sqliteMigrations[1] + expect(migration?.sql.trim()).not.toBe('') + expect(migration?.sql).toBe( + readFileSync( + new URL( + '../drizzle/0001_tanstack_ai_interrupt_batches.sql', + import.meta.url, + ), + 'utf8', + ), + ) + }) + + it('backfills safe legacy interrupt identity and creates durable indexes', () => { + const database = createDatabase() + const initialMigration = sqliteMigrations[0] + expect(initialMigration).toBeDefined() + if (!initialMigration) return + applySqliteMigrations(database, [initialMigration]) + database + .prepare( + `INSERT INTO interrupts + (interrupt_id, run_id, thread_id, status, requested_at, payload_json) + VALUES (?, ?, ?, ?, ?, ?)`, + ) + .run( + 'pending-int', + 'pending-run', + 'thread-legacy', + 'pending', + 1, + '{"id":"pending-int","reason":"confirmation"}', + ) + database + .prepare( + `INSERT INTO interrupts + (interrupt_id, run_id, thread_id, status, requested_at, payload_json) + VALUES (?, ?, ?, ?, ?, ?)`, + ) + .run( + 'resolved-int', + 'resolved-run', + 'thread-legacy', + 'resolved', + 2, + '{"id":"resolved-int","reason":"confirmation"}', + ) + + applySqliteMigrations(database, sqliteMigrations) + + expect( + database + .prepare( + `SELECT interrupt_id, generation, binding_json, response_schema_hash + FROM interrupts ORDER BY interrupt_id`, + ) + .all(), + ).toEqual([ + { + interrupt_id: 'pending-int', + generation: 1, + binding_json: + '{"kind":"generic","interruptId":"pending-int","interruptedRunId":"pending-run","generation":1,"responseSchemaHash":"legacy:unknown"}', + response_schema_hash: 'legacy:unknown', + }, + { + interrupt_id: 'resolved-int', + generation: 0, + binding_json: + '{"kind":"generic","interruptId":"resolved-int","interruptedRunId":"resolved-run","generation":0,"responseSchemaHash":"legacy:unknown"}', + response_schema_hash: 'legacy:unknown', + }, + ]) + + const indexes = database + .prepare( + "SELECT name FROM sqlite_master WHERE type = 'index' AND tbl_name IN ('interrupts', 'interrupt_batches') ORDER BY name", + ) + .all() + .map((row) => row.name) + expect(indexes).toEqual( + expect.arrayContaining([ + 'interrupt_batches_continuation_run_id_unique', + 'interrupts_run_status_requested_at_idx', + 'interrupts_thread_status_requested_at_idx', + ]), + ) }) it('rolls back migration SQL and bookkeeping together, then permits retry', () => { diff --git a/packages/ai-persistence-prisma/prisma/tanstack-ai.prisma b/packages/ai-persistence-prisma/prisma/tanstack-ai.prisma index 382893ed3..c0d9ebfef 100644 --- a/packages/ai-persistence-prisma/prisma/tanstack-ai.prisma +++ b/packages/ai-persistence-prisma/prisma/tanstack-ai.prisma @@ -30,15 +30,36 @@ model Interrupt { interruptId String @id @map("interrupt_id") runId String @map("run_id") threadId String @map("thread_id") + generation Int @default(0) status String requestedAt BigInt @map("requested_at") resolvedAt BigInt? @map("resolved_at") payloadJson String @map("payload_json") + bindingJson String? @map("binding_json") + schemaHash String? @map("schema_hash") responseJson String? @map("response_json") + @@index([threadId, status]) + @@index([runId, generation, status]) @@map("interrupts") } +/// Immutable winner for an atomically committed interrupt batch. +model InterruptBatch { + interruptedRunId String @id @map("interrupted_run_id") + threadId String @map("thread_id") + generation Int + expectedInterruptIdsJson String @map("expected_interrupt_ids_json") + fingerprint String + canonicalResolutions String @map("canonical_resolutions") + resolutionsJson String @map("resolutions_json") + continuationRunId String @map("continuation_run_id") + committedAt BigInt @map("committed_at") + + @@index([threadId]) + @@map("interrupt_batches") +} + /// Scoped key/value metadata (`MetadataStore`). model Metadata { scope String diff --git a/packages/ai-persistence-prisma/src/assets/tanstack-ai.prisma b/packages/ai-persistence-prisma/src/assets/tanstack-ai.prisma index 382893ed3..c0d9ebfef 100644 --- a/packages/ai-persistence-prisma/src/assets/tanstack-ai.prisma +++ b/packages/ai-persistence-prisma/src/assets/tanstack-ai.prisma @@ -30,15 +30,36 @@ model Interrupt { interruptId String @id @map("interrupt_id") runId String @map("run_id") threadId String @map("thread_id") + generation Int @default(0) status String requestedAt BigInt @map("requested_at") resolvedAt BigInt? @map("resolved_at") payloadJson String @map("payload_json") + bindingJson String? @map("binding_json") + schemaHash String? @map("schema_hash") responseJson String? @map("response_json") + @@index([threadId, status]) + @@index([runId, generation, status]) @@map("interrupts") } +/// Immutable winner for an atomically committed interrupt batch. +model InterruptBatch { + interruptedRunId String @id @map("interrupted_run_id") + threadId String @map("thread_id") + generation Int + expectedInterruptIdsJson String @map("expected_interrupt_ids_json") + fingerprint String + canonicalResolutions String @map("canonical_resolutions") + resolutionsJson String @map("resolutions_json") + continuationRunId String @map("continuation_run_id") + committedAt BigInt @map("committed_at") + + @@index([threadId]) + @@map("interrupt_batches") +} + /// Scoped key/value metadata (`MetadataStore`). model Metadata { scope String diff --git a/packages/ai-persistence-prisma/src/stores.ts b/packages/ai-persistence-prisma/src/stores.ts index c1a94e0c9..b8fd9c238 100644 --- a/packages/ai-persistence-prisma/src/stores.ts +++ b/packages/ai-persistence-prisma/src/stores.ts @@ -7,8 +7,26 @@ * provider-neutral Prisma `String` fields, so they are serialized with * `JSON.stringify`/`JSON.parse` here; blob bytes use Prisma's `Bytes`. */ -import type { PrismaClient } from '@prisma/client' -import type { ModelMessage } from '@tanstack/ai' +import { + canonicalInterruptJson, + canonicalizeInterruptResolutions, + cloneAndDeepFreezeJson, +} from '@tanstack/ai' +import { + InterruptStoreCorruptionError, + hasExactInterruptIds, + projectInterruptRecovery, +} from '@tanstack/ai-persistence' +import type { Prisma, PrismaClient } from '@prisma/client' +import type { + CommitInterruptResolutionsInput, + InterruptBinding, + InterruptCommitResult, + InterruptRecoveryQuery, + InterruptRecoveryStateV1, + ModelMessage, + RunAgentResumeItem, +} from '@tanstack/ai' import type { ArtifactRecord, ArtifactStore, @@ -18,6 +36,7 @@ import type { BlobObject, BlobRecord, BlobStore, + InterruptBatchRecord, InterruptRecord, InterruptStore, MessageStore, @@ -197,42 +216,473 @@ interface InterruptRow { interruptId: string runId: string threadId: string + generation: number status: string requestedAt: bigint resolvedAt: bigint | null payloadJson: string + bindingJson: string | null + schemaHash: string | null responseJson: string | null } +interface InterruptBatchRow { + interruptedRunId: string + threadId: string + generation: number + expectedInterruptIdsJson: string + fingerprint: string + canonicalResolutions: string + resolutionsJson: string + continuationRunId: string + committedAt: bigint +} + +function parseStoredJson(value: string, label: string): unknown { + try { + return JSON.parse(value) + } catch (error) { + throw new InterruptStoreCorruptionError( + `Stored ${label} is not valid JSON.`, + { cause: error }, + ) + } +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value) +} + +function hasOptionalString( + value: Record, + key: string, +): boolean { + return value[key] === undefined || typeof value[key] === 'string' +} + +function isInterruptBinding(value: unknown): value is InterruptBinding { + if ( + !isRecord(value) || + typeof value.kind !== 'string' || + typeof value.interruptId !== 'string' || + typeof value.interruptedRunId !== 'string' || + typeof value.generation !== 'number' || + !Number.isInteger(value.generation) || + typeof value.responseSchemaHash !== 'string' || + !hasOptionalString(value, 'expiresAt') + ) { + return false + } + if (value.kind === 'generic') return true + if (value.kind === 'client-tool-execution') { + return ( + typeof value.toolName === 'string' && + typeof value.toolCallId === 'string' && + typeof value.outputSchemaHash === 'string' + ) + } + if (value.kind === 'tool-approval') { + return ( + typeof value.toolName === 'string' && + typeof value.toolCallId === 'string' && + 'originalArgs' in value && + typeof value.inputSchemaHash === 'string' && + typeof value.approvalSchemaHash === 'string' + ) + } + return false +} + +function decodeBinding(row: InterruptRow): InterruptBinding { + const fallback: InterruptBinding = { + kind: 'generic', + interruptId: row.interruptId, + interruptedRunId: row.runId, + generation: row.generation, + responseSchemaHash: row.schemaHash ?? 'legacy:unknown', + } + const parsed = + row.bindingJson === null + ? fallback + : parseStoredJson(row.bindingJson, `binding for ${row.interruptId}`) + if (!isInterruptBinding(parsed)) { + throw new InterruptStoreCorruptionError( + `Stored binding for interrupt ${row.interruptId} is malformed.`, + ) + } + if ( + parsed.interruptId !== row.interruptId || + parsed.interruptedRunId !== row.runId || + parsed.generation !== row.generation + ) { + throw new InterruptStoreCorruptionError( + `Stored binding correlation does not match interrupt ${row.interruptId}.`, + ) + } + if (row.schemaHash !== null && parsed.responseSchemaHash !== row.schemaHash) { + throw new InterruptStoreCorruptionError( + `Stored schema hash does not match interrupt ${row.interruptId}.`, + ) + } + return parsed +} + +function decodeInterruptPayload(row: InterruptRow): unknown { + const payload = parseStoredJson( + row.payloadJson, + `payload for ${row.interruptId}`, + ) + if (isRecord(payload) && 'id' in payload && payload.id !== row.interruptId) { + throw new InterruptStoreCorruptionError( + `Stored payload ID does not match interrupt ${row.interruptId}.`, + ) + } + if ( + row.generation > 0 && + (!isRecord(payload) || payload.id !== row.interruptId) + ) { + throw new InterruptStoreCorruptionError( + `Stored native interrupt ${row.interruptId} has no matching payload ID.`, + ) + } + return payload +} + function mapInterrupt(row: InterruptRow): InterruptRecord { - return { + if ( + !Number.isInteger(row.generation) || + row.generation < 0 || + (row.status !== 'pending' && + row.status !== 'resolved' && + row.status !== 'cancelled') + ) { + throw new InterruptStoreCorruptionError( + `Stored interrupt ${row.interruptId} has invalid lifecycle fields.`, + ) + } + const record: InterruptRecord = { interruptId: row.interruptId, runId: row.runId, threadId: row.threadId, - status: row.status as InterruptRecord['status'], + generation: row.generation, + status: row.status, requestedAt: Number(row.requestedAt), ...(row.resolvedAt != null ? { resolvedAt: Number(row.resolvedAt) } : {}), - payload: JSON.parse(row.payloadJson) as Record, + payload: decodeInterruptPayload(row), + binding: decodeBinding(row), ...(row.responseJson != null - ? { response: JSON.parse(row.responseJson) as unknown } + ? { + response: parseStoredJson( + row.responseJson, + `response for ${row.interruptId}`, + ), + } : {}), } + return cloneAndDeepFreezeJson(record) +} + +function isResumeItem(value: unknown): value is RunAgentResumeItem { + return ( + isRecord(value) && + typeof value.interruptId === 'string' && + (value.status === 'resolved' || value.status === 'cancelled') + ) +} + +function decodeStringArray( + value: string, + label: string, +): ReadonlyArray { + const parsed = parseStoredJson(value, label) + if ( + !Array.isArray(parsed) || + !parsed.every((item): item is string => typeof item === 'string') || + !hasExactInterruptIds(parsed, parsed) + ) { + throw new InterruptStoreCorruptionError( + `Stored ${label} is not an exact set of string IDs.`, + ) + } + return cloneAndDeepFreezeJson(parsed) } -export function createInterruptStore(prisma: PrismaClient): InterruptStore { +function decodeInterruptBatch( + row: InterruptBatchRow, + expectedInterruptedRunId: string, +): InterruptBatchRecord { + if ( + row.interruptedRunId !== expectedInterruptedRunId || + !Number.isInteger(row.generation) || + row.generation < 0 + ) { + throw new InterruptStoreCorruptionError( + `Stored interrupt batch correlation does not match ${expectedInterruptedRunId}.`, + ) + } + const expectedInterruptIds = decodeStringArray( + row.expectedInterruptIdsJson, + `expected IDs for ${row.interruptedRunId}`, + ) + const resolutions = parseStoredJson( + row.resolutionsJson, + `resolutions for ${row.interruptedRunId}`, + ) + if (!Array.isArray(resolutions) || !resolutions.every(isResumeItem)) { + throw new InterruptStoreCorruptionError( + `Stored resolutions for ${row.interruptedRunId} are malformed.`, + ) + } + const canonical = canonicalizeInterruptResolutions(resolutions) + if ( + canonical.fingerprint !== row.fingerprint || + canonical.canonicalResolutions !== row.canonicalResolutions || + !hasExactInterruptIds( + expectedInterruptIds, + canonical.resolutions.map((resolution) => resolution.interruptId), + ) + ) { + throw new InterruptStoreCorruptionError( + `Stored interrupt batch identity is corrupt for ${row.interruptedRunId}.`, + ) + } + return cloneAndDeepFreezeJson({ + threadId: row.threadId, + interruptedRunId: row.interruptedRunId, + generation: row.generation, + expectedInterruptIds, + fingerprint: row.fingerprint, + canonicalResolutions: row.canonicalResolutions, + resolutions: canonical.resolutions, + continuationRunId: row.continuationRunId, + committedAt: Number(row.committedAt), + }) +} + +function legacyBindingForRow(row: InterruptRow): InterruptBinding { return { + kind: 'generic', + interruptId: row.interruptId, + interruptedRunId: row.runId, + generation: row.generation, + responseSchemaHash: row.schemaHash ?? 'legacy:unknown', + } +} + +async function upgradeLegacyPendingRows( + transaction: Prisma.TransactionClient, + rows: ReadonlyArray, +): Promise> { + const upgraded: Array = [] + for (const row of rows) { + if (row.status !== 'pending' || row.bindingJson !== null) { + upgraded.push(row) + continue + } + const binding = legacyBindingForRow(row) + const bindingJson = canonicalInterruptJson(binding) + const result = await transaction.interrupt.updateMany({ + where: { + interruptId: row.interruptId, + runId: row.runId, + generation: row.generation, + status: 'pending', + bindingJson: null, + }, + data: { + bindingJson, + schemaHash: binding.responseSchemaHash, + }, + }) + if (result.count !== 1) { + throw new Error(`Pending interrupt changed: ${row.interruptId}`) + } + upgraded.push({ + ...row, + bindingJson, + schemaHash: binding.responseSchemaHash, + }) + } + return upgraded +} + +async function loadRowsForRun( + transaction: Prisma.TransactionClient, + interruptedRunId: string, +): Promise> { + const rows = await transaction.interrupt.findMany({ + where: { runId: interruptedRunId }, + orderBy: [{ requestedAt: 'asc' }, { interruptId: 'asc' }], + }) + const upgraded = await upgradeLegacyPendingRows(transaction, rows) + return upgraded.map(mapInterrupt) +} + +function validateBatchRowCorrelation( + batch: InterruptBatchRecord, + rows: ReadonlyArray, +): void { + if ( + !hasExactInterruptIds( + batch.expectedInterruptIds, + rows.map((row) => row.interruptId), + ) || + rows.some( + (row) => + row.runId !== batch.interruptedRunId || + row.threadId !== batch.threadId || + row.generation !== batch.generation || + row.status === 'pending', + ) + ) { + throw new InterruptStoreCorruptionError( + `Stored interrupt batch correlation is corrupt for ${batch.interruptedRunId}.`, + ) + } +} + +async function recoveryInTransaction( + transaction: Prisma.TransactionClient, + input: InterruptRecoveryQuery, + now: number, +): Promise { + const rows = await loadRowsForRun(transaction, input.interruptedRunId) + const batchRow = await transaction.interruptBatch.findUnique({ + where: { interruptedRunId: input.interruptedRunId }, + }) + const batch = batchRow + ? decodeInterruptBatch(batchRow, input.interruptedRunId) + : null + if (batch?.threadId === input.threadId) { + validateBatchRowCorrelation(batch, rows) + } + return projectInterruptRecovery({ + query: input, + rows: rows.filter((row) => row.threadId === input.threadId), + batch: batch?.threadId === input.threadId ? batch : null, + now, + includeResolutions: true, + }) +} + +async function winnerOutcome( + transaction: Prisma.TransactionClient, + input: CommitInterruptResolutionsInput, + winner: InterruptBatchRecord, + now: number, +): Promise { + const authoritativeState = await recoveryInTransaction( + transaction, + { + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + knownGeneration: input.expectedGeneration, + }, + now, + ) + if ( + winner.threadId === input.threadId && + winner.generation === input.expectedGeneration && + hasExactInterruptIds( + winner.expectedInterruptIds, + input.expectedInterruptIds, + ) && + winner.fingerprint === input.fingerprint && + winner.canonicalResolutions === input.canonicalResolutions + ) { + return { + status: 'replayed', + continuationRunId: winner.continuationRunId, + } + } + return { + status: 'conflict', + authoritativeState, + } +} + +function errorCode(error: unknown): string | undefined { + return isRecord(error) && typeof error.code === 'string' + ? error.code + : undefined +} + +function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error) +} + +function isReloadableCommitRace(error: unknown): boolean { + if (errorCode(error) === 'P2002') return true + if (!['P1008', 'P2028', 'P2034', undefined].includes(errorCode(error))) { + return false + } + return /database is locked|database is busy|write conflict|deadlock|timed out|transaction.*closed/i.test( + errorMessage(error), + ) +} + +export function createInterruptStore( + prisma: PrismaClient, + clock: () => number = Date.now, +): InterruptStore { + async function readRows( + where: Prisma.InterruptWhereInput, + ): Promise> { + return prisma.$transaction(async (transaction) => { + const rows = await transaction.interrupt.findMany({ + where, + orderBy: [{ requestedAt: 'asc' }, { interruptId: 'asc' }], + }) + const upgraded = await upgradeLegacyPendingRows(transaction, rows) + return upgraded.map(mapInterrupt) + }) + } + + async function reloadWinner( + input: CommitInterruptResolutionsInput, + ): Promise { + return prisma.$transaction(async (transaction) => { + const row = await transaction.interruptBatch.findUnique({ + where: { interruptedRunId: input.interruptedRunId }, + }) + if (!row) return null + return winnerOutcome( + transaction, + input, + decodeInterruptBatch(row, input.interruptedRunId), + clock(), + ) + }) + } + + const store: InterruptStore = { async create(record) { + const generation = record.generation ?? 0 + const binding = + record.binding ?? + ({ + kind: 'generic', + interruptId: record.interruptId, + interruptedRunId: record.runId, + generation, + responseSchemaHash: 'legacy:unknown', + } as const) await prisma.interrupt.upsert({ where: { interruptId: record.interruptId }, create: { interruptId: record.interruptId, runId: record.runId, threadId: record.threadId, + generation, status: record.status, requestedAt: BigInt(record.requestedAt), - payloadJson: JSON.stringify(record.payload), + payloadJson: canonicalInterruptJson(record.payload), + bindingJson: canonicalInterruptJson(binding), + schemaHash: binding.responseSchemaHash, responseJson: - record.response == null ? null : JSON.stringify(record.response), + record.response === undefined + ? null + : canonicalInterruptJson(record.response), }, update: {}, }) @@ -242,50 +692,274 @@ export function createInterruptStore(prisma: PrismaClient): InterruptStore { where: { interruptId }, data: { status: 'resolved', - resolvedAt: BigInt(Date.now()), - responseJson: response == null ? null : JSON.stringify(response), + resolvedAt: BigInt(clock()), + ...(response === undefined + ? {} + : { responseJson: canonicalInterruptJson(response) }), }, }) }, async cancel(interruptId) { await prisma.interrupt.updateMany({ where: { interruptId }, - data: { status: 'cancelled', resolvedAt: BigInt(Date.now()) }, + data: { status: 'cancelled', resolvedAt: BigInt(clock()) }, }) }, async get(interruptId) { - const row = await prisma.interrupt.findUnique({ where: { interruptId } }) - return row ? mapInterrupt(row) : null + return prisma.$transaction(async (transaction) => { + const row = await transaction.interrupt.findUnique({ + where: { interruptId }, + }) + if (!row) return null + const [upgraded] = await upgradeLegacyPendingRows(transaction, [row]) + return upgraded ? mapInterrupt(upgraded) : null + }) }, async list(threadId) { - const rows = await prisma.interrupt.findMany({ - where: { threadId }, - orderBy: { requestedAt: 'asc' }, - }) - return rows.map(mapInterrupt) + return readRows({ threadId }) }, async listPending(threadId) { - const rows = await prisma.interrupt.findMany({ - where: { threadId, status: 'pending' }, - orderBy: { requestedAt: 'asc' }, - }) - return rows.map(mapInterrupt) + return readRows({ threadId, status: 'pending' }) }, async listByRun(runId) { - const rows = await prisma.interrupt.findMany({ - where: { runId }, - orderBy: { requestedAt: 'asc' }, - }) - return rows.map(mapInterrupt) + return readRows({ runId }) }, async listPendingByRun(runId) { - const rows = await prisma.interrupt.findMany({ - where: { runId, status: 'pending' }, - orderBy: { requestedAt: 'asc' }, + return readRows({ runId, status: 'pending' }) + }, + async openInterruptBatch(input) { + const descriptors = cloneAndDeepFreezeJson(input.descriptors) + const bindings = cloneAndDeepFreezeJson(input.bindings) + const descriptorIds = descriptors.map((descriptor) => descriptor.id) + const bindingIds = bindings.map((binding) => binding.interruptId) + if ( + descriptorIds.length === 0 || + !hasExactInterruptIds(descriptorIds, bindingIds) + ) { + throw new TypeError( + 'Interrupt descriptors and bindings must contain the same exact nonempty IDs.', + ) + } + + return prisma.$transaction(async (transaction) => { + const committed = await transaction.interruptBatch.findUnique({ + where: { interruptedRunId: input.interruptedRunId }, + }) + if (committed) { + decodeInterruptBatch(committed, input.interruptedRunId) + throw new Error('Cannot reopen a committed interrupt batch.') + } + + const existing = await loadRowsForRun( + transaction, + input.interruptedRunId, + ) + if (existing.length > 0) { + const generation = existing[0]?.generation + const existingDescriptors = [...existing] + .sort((left, right) => + left.interruptId.localeCompare(right.interruptId), + ) + .map((record) => record.payload) + const requestedDescriptors = [...descriptors].sort((left, right) => + left.id.localeCompare(right.id), + ) + const existingBindings = [...existing] + .map((record) => record.binding) + .sort((left, right) => + left.interruptId.localeCompare(right.interruptId), + ) + const requestedBindings = bindings + .map((binding) => ({ + ...binding, + interruptedRunId: input.interruptedRunId, + generation, + })) + .sort((left, right) => + left.interruptId.localeCompare(right.interruptId), + ) + if ( + generation !== undefined && + hasExactInterruptIds( + descriptorIds, + existing.map((record) => record.interruptId), + ) && + canonicalInterruptJson(existingDescriptors) === + canonicalInterruptJson(requestedDescriptors) && + canonicalInterruptJson(existingBindings) === + canonicalInterruptJson(requestedBindings) && + existing.every( + (record) => + record.status === 'pending' && + record.threadId === input.threadId && + record.generation === generation, + ) + ) { + return { generation, descriptors } + } + throw new Error('An incompatible interrupt batch is already open.') + } + + const generation = 1 + const bindingsById = new Map( + bindings.map((binding) => [binding.interruptId, binding]), + ) + const requestedAt = BigInt(clock()) + for (const descriptor of descriptors) { + const unopened = bindingsById.get(descriptor.id) + if (!unopened) { + throw new Error(`Missing binding for interrupt: ${descriptor.id}`) + } + const binding = cloneAndDeepFreezeJson({ + ...unopened, + interruptedRunId: input.interruptedRunId, + generation, + }) + await transaction.interrupt.create({ + data: { + interruptId: descriptor.id, + runId: input.interruptedRunId, + threadId: input.threadId, + generation, + status: 'pending', + requestedAt, + payloadJson: canonicalInterruptJson(descriptor), + bindingJson: canonicalInterruptJson(binding), + schemaHash: binding.responseSchemaHash, + }, + }) + } + return { generation, descriptors } }) - return rows.map(mapInterrupt) + }, + async commitInterruptResolutions(input) { + const candidate = canonicalizeInterruptResolutions(input.resolutions) + if ( + candidate.fingerprint !== input.fingerprint || + candidate.canonicalResolutions !== input.canonicalResolutions + ) { + throw new TypeError( + 'Interrupt batch identity does not match its resolutions.', + ) + } + + try { + return await prisma.$transaction(async (transaction) => { + const winnerRow = await transaction.interruptBatch.findUnique({ + where: { interruptedRunId: input.interruptedRunId }, + }) + if (winnerRow) { + return winnerOutcome( + transaction, + input, + decodeInterruptBatch(winnerRow, input.interruptedRunId), + clock(), + ) + } + + const rows = await loadRowsForRun(transaction, input.interruptedRunId) + const pending = rows.filter((row) => row.status === 'pending') + const resolutionIds = candidate.resolutions.map( + (resolution) => resolution.interruptId, + ) + const now = clock() + if ( + !hasExactInterruptIds( + input.expectedInterruptIds, + pending.map((record) => record.interruptId), + ) || + !hasExactInterruptIds(input.expectedInterruptIds, resolutionIds) || + pending.some((record) => record.threadId !== input.threadId) || + pending.some( + (record) => record.generation !== input.expectedGeneration, + ) || + pending.some( + (record) => + record.binding.expiresAt !== undefined && + Date.parse(record.binding.expiresAt) <= now, + ) + ) { + return { + status: 'conflict', + authoritativeState: await recoveryInTransaction( + transaction, + { + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + knownGeneration: input.expectedGeneration, + }, + now, + ), + } + } + + await transaction.interruptBatch.create({ + data: { + interruptedRunId: input.interruptedRunId, + threadId: input.threadId, + generation: input.expectedGeneration, + expectedInterruptIdsJson: canonicalInterruptJson( + [...input.expectedInterruptIds].sort(), + ), + fingerprint: candidate.fingerprint, + canonicalResolutions: candidate.canonicalResolutions, + resolutionsJson: candidate.canonicalResolutions, + continuationRunId: input.continuationRunId, + committedAt: BigInt(now), + }, + }) + + for (const resolution of candidate.resolutions) { + const transitioned = await transaction.interrupt.updateMany({ + where: { + interruptId: resolution.interruptId, + runId: input.interruptedRunId, + threadId: input.threadId, + generation: input.expectedGeneration, + status: 'pending', + }, + data: { + status: + resolution.status === 'cancelled' ? 'cancelled' : 'resolved', + responseJson: canonicalInterruptJson(resolution), + resolvedAt: BigInt(now), + }, + }) + if (transitioned.count !== 1) { + throw new Error( + `Pending interrupt changed: ${resolution.interruptId}`, + ) + } + } + + return { + status: 'committed', + continuationRunId: input.continuationRunId, + } + }) + } catch (error) { + if (!isReloadableCommitRace(error)) throw error + for (let attempt = 0; attempt < 4; attempt++) { + try { + const winner = await reloadWinner(input) + if (winner) return winner + } catch (reloadError) { + if (!isReloadableCommitRace(reloadError)) throw reloadError + } + await new Promise((resolve) => { + setTimeout(resolve, 10 * (attempt + 1)) + }) + } + throw error + } + }, + async getInterruptRecoveryState(input) { + return prisma.$transaction((transaction) => + recoveryInTransaction(transaction, input, clock()), + ) }, } + return store } export function createMetadataStore(prisma: PrismaClient): MetadataStore { diff --git a/packages/ai-persistence-prisma/tests/api-types.test-d.ts b/packages/ai-persistence-prisma/tests/api-types.test-d.ts index ceadef10f..4738815c1 100644 --- a/packages/ai-persistence-prisma/tests/api-types.test-d.ts +++ b/packages/ai-persistence-prisma/tests/api-types.test-d.ts @@ -17,6 +17,15 @@ const persistence = prismaPersistence(prisma) expectTypeOf(persistence.stores.messages).toEqualTypeOf() expectTypeOf(persistence.stores.runs).toEqualTypeOf() expectTypeOf(persistence.stores.interrupts).toEqualTypeOf() +expectTypeOf(persistence.stores.interrupts.openInterruptBatch).toEqualTypeOf< + InterruptStore['openInterruptBatch'] +>() +expectTypeOf( + persistence.stores.interrupts.commitInterruptResolutions, +).toEqualTypeOf() +expectTypeOf( + persistence.stores.interrupts.getInterruptRecoveryState, +).toEqualTypeOf() expectTypeOf(persistence.stores.metadata).toEqualTypeOf() expectTypeOf(persistence.stores.artifacts).toEqualTypeOf() expectTypeOf(persistence.stores.blobs).toEqualTypeOf() diff --git a/packages/ai-persistence-prisma/tests/models-cli.test.ts b/packages/ai-persistence-prisma/tests/models-cli.test.ts index 36a997568..ea07ebfb4 100644 --- a/packages/ai-persistence-prisma/tests/models-cli.test.ts +++ b/packages/ai-persistence-prisma/tests/models-cli.test.ts @@ -30,6 +30,8 @@ describe('Prisma models CLI', () => { }, }) expect(output).toBe(`${prismaModels.trimEnd()}\n`) + expect(output).toContain('model InterruptBatch') + expect(output).toContain('bindingJson') }) it('copies the fragment without overwriting divergent files by default', async () => { diff --git a/packages/ai-persistence-prisma/tests/models.test.ts b/packages/ai-persistence-prisma/tests/models.test.ts index 06844786d..d81058e3b 100644 --- a/packages/ai-persistence-prisma/tests/models.test.ts +++ b/packages/ai-persistence-prisma/tests/models.test.ts @@ -29,7 +29,17 @@ describe('Prisma models asset', () => { it('is a models-only fragment with no provider or client generator', () => { expect(prismaModelsFilename).toBe('tanstack-ai.prisma') expect(prismaModels).toContain('model Message') + expect(prismaModels).toContain('model InterruptBatch') expect(prismaModels).toContain('model Blob') + expect(prismaModels).toMatch(/generation\s+Int\s+@default\(0\)/) + expect(prismaModels).toMatch( + /bindingJson\s+String\?\s+@map\("binding_json"\)/, + ) + expect(prismaModels).toMatch( + /schemaHash\s+String\?\s+@map\("schema_hash"\)/, + ) + expect(prismaModels).toContain('@@index([threadId, status])') + expect(prismaModels).toContain('@@index([runId, generation, status])') expect(prismaModels).not.toMatch(/\bgenerator\s+\w+\s*{/) expect(prismaModels).not.toMatch(/\bdatasource\s+\w+\s*{/) }) diff --git a/packages/ai-persistence-prisma/tests/prisma.conformance.test.ts b/packages/ai-persistence-prisma/tests/prisma.conformance.test.ts index a73190329..2a848bc14 100644 --- a/packages/ai-persistence-prisma/tests/prisma.conformance.test.ts +++ b/packages/ai-persistence-prisma/tests/prisma.conformance.test.ts @@ -3,10 +3,17 @@ import { mkdtempSync } from 'node:fs' import { tmpdir } from 'node:os' import { join } from 'node:path' import { DatabaseSync } from 'node:sqlite' -import { afterAll } from 'vitest' +import { afterAll, describe, expect, it } from 'vitest' import { PrismaClient } from '@prisma/client' -import { runPersistenceConformance } from '@tanstack/ai-persistence/testkit' +import { canonicalizeInterruptResolutions } from '@tanstack/ai' +import { + runInterruptStoreConformance, + runPersistenceConformance, + type InterruptConformanceHarness, +} from '@tanstack/ai-persistence/testkit' import { prismaPersistence } from '../src/index' +import { createInterruptStore } from '../src/stores' +import type { InterruptStore } from '@tanstack/ai-persistence' const clients: Array = [] const temporaryDirectories: Array = [] @@ -29,12 +36,32 @@ const sqliteTestSchema = ` interrupt_id TEXT NOT NULL PRIMARY KEY, run_id TEXT NOT NULL, thread_id TEXT NOT NULL, + generation INTEGER NOT NULL DEFAULT 0, status TEXT NOT NULL, requested_at BIGINT NOT NULL, resolved_at BIGINT, payload_json TEXT NOT NULL, + binding_json TEXT, + schema_hash TEXT, response_json TEXT ); + CREATE INDEX interrupts_thread_status_idx + ON interrupts(thread_id, status); + CREATE INDEX interrupts_run_generation_status_idx + ON interrupts(run_id, generation, status); + CREATE TABLE interrupt_batches ( + interrupted_run_id TEXT NOT NULL PRIMARY KEY, + thread_id TEXT NOT NULL, + generation INTEGER NOT NULL, + expected_interrupt_ids_json TEXT NOT NULL, + fingerprint TEXT NOT NULL, + canonical_resolutions TEXT NOT NULL, + resolutions_json TEXT NOT NULL, + continuation_run_id TEXT NOT NULL, + committed_at BIGINT NOT NULL + ); + CREATE INDEX interrupt_batches_thread_idx + ON interrupt_batches(thread_id); CREATE TABLE metadata ( scope TEXT NOT NULL, key TEXT NOT NULL, @@ -85,6 +112,40 @@ async function makeTestClient(): Promise { return prisma } +async function makeSharedTestClients(): Promise<{ + clients: readonly [PrismaClient, PrismaClient] + dbPath: string +}> { + const dir = mkdtempSync(join(tmpdir(), 'tanstack-ai-prisma-interrupts-')) + temporaryDirectories.push(dir) + const dbPath = join(dir, 'state.db').replace(/\\/g, '/') + initializeSqliteTestDatabase(dbPath) + const first = new PrismaClient({ + datasources: { db: { url: `file:${dbPath}` } }, + }) + const second = new PrismaClient({ + datasources: { db: { url: `file:${dbPath}` } }, + }) + clients.push(first, second) + return { clients: [first, second], dbPath } +} + +function withTransitionFailure(dbPath: string, interruptId: string): void { + const database = new DatabaseSync(dbPath) + try { + database.exec(` + CREATE TRIGGER fail_next_interrupt_transition + BEFORE UPDATE ON interrupts + WHEN OLD.status = 'pending' AND NEW.interrupt_id = '${interruptId}' + BEGIN + SELECT RAISE(ABORT, 'Injected transition failure'); + END; + `) + } finally { + database.close() + } +} + afterAll(async () => { await Promise.all(clients.map((client) => client.$disconnect())) await Promise.all( @@ -97,3 +158,234 @@ afterAll(async () => { runPersistenceConformance('prisma', async () => prismaPersistence(await makeTestClient()), ) + +runInterruptStoreConformance(async (): Promise => { + let now = Date.parse('2026-07-13T10:00:00.000Z') + const shared = await makeSharedTestClients() + const first = createInterruptStore(shared.clients[0], () => now) + const second = createInterruptStore(shared.clients[1], () => now) + let commitCount = 0 + const store: InterruptStore = { + create: (record) => first.create(record), + resolve: (interruptId, response) => first.resolve(interruptId, response), + cancel: (interruptId) => first.cancel(interruptId), + get: (interruptId) => first.get(interruptId), + list: (threadId) => first.list(threadId), + listPending: (threadId) => first.listPending(threadId), + listByRun: (runId) => first.listByRun(runId), + listPendingByRun: (runId) => first.listPendingByRun(runId), + openInterruptBatch: (input) => first.openInterruptBatch(input), + commitInterruptResolutions: (input) => { + const target = commitCount++ % 2 === 0 ? first : second + return target.commitInterruptResolutions(input) + }, + getInterruptRecoveryState: (input) => + first.getInterruptRecoveryState(input), + } + + return { + getStore: () => store, + advanceBy: (milliseconds) => { + now += milliseconds + }, + inspect: async (interruptedRunId) => { + const database = new DatabaseSync(shared.dbPath, { readOnly: true }) + try { + const statuses = database + .prepare( + 'SELECT status FROM interrupts WHERE run_id = ? ORDER BY interrupt_id', + ) + .all(interruptedRunId) + .map((row) => { + if (typeof row.status !== 'string') { + throw new TypeError('Expected a string interrupt status.') + } + return row.status + }) + const batch = database + .prepare( + 'SELECT COUNT(*) AS count FROM interrupt_batches WHERE interrupted_run_id = ?', + ) + .get(interruptedRunId) + if (typeof batch?.count !== 'number') { + throw new TypeError('Expected a numeric interrupt batch count.') + } + return { statuses, batchCount: batch.count } + } finally { + database.close() + } + }, + failTransitionOnce: (interruptId) => { + withTransitionFailure(shared.dbPath, interruptId) + }, + reopen: async () => { + const reopened = new PrismaClient({ + datasources: { db: { url: `file:${shared.dbPath}` } }, + }) + clients.push(reopened) + return createInterruptStore(reopened, () => now) + }, + } +}) + +describe('Prisma interrupt persistence hardening', () => { + it('upgrades a pending legacy row to a safe generic binding', async () => { + const shared = await makeSharedTestClients() + const database = new DatabaseSync(shared.dbPath) + try { + database + .prepare( + `INSERT INTO interrupts ( + interrupt_id, run_id, thread_id, generation, status, requested_at, + payload_json, binding_json, schema_hash + ) VALUES (?, ?, ?, ?, ?, ?, ?, NULL, NULL)`, + ) + .run( + 'legacy-pending', + 'legacy-run', + 'legacy-thread', + 0, + 'pending', + 1, + JSON.stringify({ + id: 'legacy-pending', + reason: 'confirmation', + toolName: 'must-not-be-guessed', + }), + ) + } finally { + database.close() + } + + const store = createInterruptStore(shared.clients[0], () => 2) + await expect( + store.getInterruptRecoveryState({ + threadId: 'legacy-thread', + interruptedRunId: 'legacy-run', + knownGeneration: 0, + }), + ).resolves.toMatchObject({ state: 'pending' }) + + const verification = new DatabaseSync(shared.dbPath, { readOnly: true }) + try { + const row = verification + .prepare( + 'SELECT binding_json AS bindingJson, schema_hash AS schemaHash FROM interrupts WHERE interrupt_id = ?', + ) + .get('legacy-pending') + expect(row?.schemaHash).toBe('legacy:unknown') + expect(JSON.parse(String(row?.bindingJson))).toEqual({ + generation: 0, + interruptId: 'legacy-pending', + interruptedRunId: 'legacy-run', + kind: 'generic', + responseSchemaHash: 'legacy:unknown', + }) + } finally { + verification.close() + } + }) + + it('rejects persisted rows whose payload ID does not match the row ID', async () => { + const shared = await makeSharedTestClients() + const store = createInterruptStore(shared.clients[0], () => 1) + await store.create({ + interruptId: 'row-id', + runId: 'row-run', + threadId: 'row-thread', + status: 'pending', + requestedAt: 1, + payload: { id: 'different-id', reason: 'confirmation' }, + }) + await expect(store.get('row-id')).rejects.toThrow(/corrupt|match/i) + }) + + it('requires replay correlation to match the stored winner exactly', async () => { + const shared = await makeSharedTestClients() + const store = createInterruptStore(shared.clients[0], () => 1) + const opened = await store.openInterruptBatch({ + threadId: 'replay-thread', + interruptedRunId: 'replay-run', + descriptors: [{ id: 'replay-int', reason: 'confirmation' }], + bindings: [ + { + interruptId: 'replay-int', + kind: 'generic', + responseSchemaHash: 'sha256:replay', + }, + ], + }) + const candidate = canonicalizeInterruptResolutions([ + { interruptId: 'replay-int', status: 'resolved', payload: true }, + ]) + const input = { + threadId: 'replay-thread', + interruptedRunId: 'replay-run', + continuationRunId: 'continuation-winner', + expectedGeneration: opened.generation, + expectedInterruptIds: ['replay-int'], + resolutions: candidate.resolutions, + fingerprint: candidate.fingerprint, + canonicalResolutions: candidate.canonicalResolutions, + } + await expect( + store.commitInterruptResolutions(input), + ).resolves.toMatchObject({ status: 'committed' }) + await expect( + store.commitInterruptResolutions({ + ...input, + continuationRunId: 'continuation-loser', + expectedGeneration: opened.generation + 1, + }), + ).resolves.toMatchObject({ status: 'conflict' }) + }) + + it('rejects a stored batch whose generation diverges from its rows', async () => { + const shared = await makeSharedTestClients() + const store = createInterruptStore(shared.clients[0], () => 1) + const opened = await store.openInterruptBatch({ + threadId: 'corrupt-thread', + interruptedRunId: 'corrupt-run', + descriptors: [{ id: 'corrupt-int', reason: 'confirmation' }], + bindings: [ + { + interruptId: 'corrupt-int', + kind: 'generic', + responseSchemaHash: 'sha256:corrupt', + }, + ], + }) + const candidate = canonicalizeInterruptResolutions([ + { interruptId: 'corrupt-int', status: 'cancelled' }, + ]) + await store.commitInterruptResolutions({ + threadId: 'corrupt-thread', + interruptedRunId: 'corrupt-run', + continuationRunId: 'corrupt-continuation', + expectedGeneration: opened.generation, + expectedInterruptIds: ['corrupt-int'], + resolutions: candidate.resolutions, + fingerprint: candidate.fingerprint, + canonicalResolutions: candidate.canonicalResolutions, + }) + + const database = new DatabaseSync(shared.dbPath) + try { + database + .prepare( + 'UPDATE interrupt_batches SET generation = ? WHERE interrupted_run_id = ?', + ) + .run(opened.generation + 1, 'corrupt-run') + } finally { + database.close() + } + + await expect( + store.getInterruptRecoveryState({ + threadId: 'corrupt-thread', + interruptedRunId: 'corrupt-run', + knownGeneration: opened.generation, + }), + ).rejects.toThrow(/corrupt|correlation/i) + }) +}) diff --git a/packages/ai-persistence/src/capabilities.ts b/packages/ai-persistence/src/capabilities.ts index b1673a4ab..324d63247 100644 --- a/packages/ai-persistence/src/capabilities.ts +++ b/packages/ai-persistence/src/capabilities.ts @@ -5,18 +5,42 @@ * can read durable state. `LocksCapability` is re-exported from core * (`@tanstack/ai`) — a shared, single-owner token owned with the sandbox layer. */ -import { createCapability } from '@tanstack/ai' -import type { AIPersistence, InterruptStore } from './types' +import { + InterruptPersistenceCapability, + createCapability, + getInterruptPersistence, + provideInterruptPersistence, +} from '@tanstack/ai' +import { validatePersistenceStoreKeys } from './types' +import type { CapabilityContext } from '@tanstack/ai' +import type { AIPersistence } from './types' export const PersistenceCapability = createCapability()('persistence') -export const InterruptsCapability = createCapability()( - 'persistence.interrupts', -) +export const [getPersistence, providePersistenceUnchecked] = + PersistenceCapability -export const [getPersistence, providePersistence] = PersistenceCapability -export const [getInterrupts, provideInterrupts] = InterruptsCapability +export function providePersistence( + ctx: CapabilityContext, + persistence: AIPersistence, +): void { + validatePersistenceStoreKeys(persistence) + providePersistenceUnchecked(ctx, persistence) +} + +export { + InterruptPersistenceCapability, + getInterruptPersistence, + provideInterruptPersistence, +} + +/** @deprecated Use InterruptPersistenceCapability. Removed at 1.0. */ +export const InterruptsCapability = InterruptPersistenceCapability +/** @deprecated Use getInterruptPersistence. Removed at 1.0. */ +export const getInterrupts = getInterruptPersistence +/** @deprecated Use provideInterruptPersistence. Removed at 1.0. */ +export const provideInterrupts = provideInterruptPersistence // Shared, single-owner tokens live in core; re-export so consumers import // everything persistence-related from this package. diff --git a/packages/ai-persistence/src/index.ts b/packages/ai-persistence/src/index.ts index d8df76e68..7821a6899 100644 --- a/packages/ai-persistence/src/index.ts +++ b/packages/ai-persistence/src/index.ts @@ -6,7 +6,9 @@ export type { RunRecord, RunStore, InterruptRecord, + InterruptBatchRecord, InterruptStore, + LegacyInterruptRecordInput, MetadataStore, ArtifactRecord, ArtifactStore, @@ -24,14 +26,33 @@ export type { } from './types' // Middleware -export { withChatPersistence, withGenerationPersistence } from './middleware' +export { + InterruptReplaySignal, + InterruptResumeValidationError, + validateInterruptResumeBatch, + withChatPersistence, + withGenerationPersistence, +} from './middleware' export type { GenerationArtifactDescriptor, GenerationArtifactExtractionInput, GenerationArtifactNameInput, + ValidateInterruptResumeBatchInput, + ValidatedInterruptResumeBatch, WithPersistenceOptions, } from './middleware' +// Authenticated recovery helpers (route registration remains application-owned) +export { + createInterruptRecoveryHandler, + getInterruptRecoveryState, +} from './recovery' +export type { + GetInterruptRecoveryStateOptions, + InterruptRecoveryAuthorization, + InterruptRecoveryHandlerOptions, +} from './recovery' + export type { PersistedArtifactActivity, PersistedArtifactRef, @@ -42,17 +63,25 @@ export type { export { memoryPersistence } from './memory' // Interrupt controller -export { createInterruptController } from './interrupts' +export { + createInterruptController, + hasExactInterruptIds, + InterruptStoreCorruptionError, + projectInterruptRecovery, +} from './interrupts' export type { InterruptController } from './interrupts' // Capabilities (incl. re-exported core Locks token) export { PersistenceCapability, InterruptsCapability, + InterruptPersistenceCapability, getPersistence, providePersistence, getInterrupts, provideInterrupts, + getInterruptPersistence, + provideInterruptPersistence, LocksCapability, getLocks, provideLocks, diff --git a/packages/ai-persistence/src/interrupts.ts b/packages/ai-persistence/src/interrupts.ts index fdc39c084..28ab0f6d7 100644 --- a/packages/ai-persistence/src/interrupts.ts +++ b/packages/ai-persistence/src/interrupts.ts @@ -1,13 +1,102 @@ -import type { InterruptRecord, InterruptStore } from './types' +import { cloneAndDeepFreezeJson } from '@tanstack/ai' +import type { + CommitInterruptResolutionsInput, + Interrupt, + InterruptRecoveryQuery, + InterruptRecoveryStateV1, + OpenInterruptBatchInput, +} from '@tanstack/ai' +import type { + InterruptBatchRecord, + InterruptRecord, + InterruptStore, +} from './types' + +export class InterruptStoreCorruptionError extends Error { + constructor(message: string, options?: ErrorOptions) { + super(message, options) + this.name = 'InterruptStoreCorruptionError' + } +} + +export function hasExactInterruptIds( + expected: ReadonlyArray, + actual: ReadonlyArray, +): boolean { + const left = [...new Set(expected)].sort() + const right = [...new Set(actual)].sort() + return ( + left.length === expected.length && + right.length === actual.length && + left.length === right.length && + left.every((id, index) => id === right[index]) + ) +} + +export function projectInterruptRecovery(input: { + query: InterruptRecoveryQuery + rows: ReadonlyArray + batch: InterruptBatchRecord | null + now: number + includeResolutions: boolean +}): InterruptRecoveryStateV1 { + const correlation = { + schemaVersion: 1 as const, + threadId: input.query.threadId, + interruptedRunId: input.query.interruptedRunId, + generation: + input.batch?.generation ?? + input.rows[0]?.generation ?? + input.query.knownGeneration, + pendingInterrupts: [] as ReadonlyArray, + } + + if (input.batch) { + return cloneAndDeepFreezeJson({ + ...correlation, + state: 'committed', + committed: { + fingerprint: input.batch.fingerprint, + ...(input.includeResolutions && { + resolutions: input.batch.resolutions, + }), + continuationRunId: input.batch.continuationRunId, + committedAt: new Date(input.batch.committedAt).toISOString(), + }, + }) + } + if (input.rows.length === 0) { + return cloneAndDeepFreezeJson({ ...correlation, state: 'missing' }) + } + + const pending = input.rows.filter((row) => row.status === 'pending') + if (pending.length > 0) { + const expired = pending.some( + (row) => + row.binding.expiresAt !== undefined && + Date.parse(row.binding.expiresAt) <= input.now, + ) + return cloneAndDeepFreezeJson({ + ...correlation, + state: expired ? 'expired' : 'pending', + pendingInterrupts: expired + ? [] + : pending.map((row) => row.payload as Interrupt), + }) + } + + return cloneAndDeepFreezeJson({ + ...correlation, + state: 'legacy-committed', + }) +} export interface InterruptController { - resolve: (interruptId: string, response?: unknown) => Promise - cancel: (interruptId: string) => Promise - request: ( - record: Omit, - ) => Promise - listPending: (threadId: string) => Promise> - listPendingByRun: (runId: string) => Promise> + openInterruptBatch: InterruptStore['openInterruptBatch'] + commitInterruptResolutions: InterruptStore['commitInterruptResolutions'] + getInterruptRecoveryState: InterruptStore['getInterruptRecoveryState'] + listPending: InterruptStore['listPending'] + listPendingByRun: InterruptStore['listPendingByRun'] } export function createInterruptController(opts: { @@ -15,13 +104,12 @@ export function createInterruptController(opts: { }): InterruptController { const { store } = opts return { - resolve: (interruptId, response) => store.resolve(interruptId, response), - cancel: (interruptId) => store.cancel(interruptId), - request: (record) => - store.create({ - ...record, - status: 'pending', - }), + openInterruptBatch: (input: OpenInterruptBatchInput) => + store.openInterruptBatch(input), + commitInterruptResolutions: (input: CommitInterruptResolutionsInput) => + store.commitInterruptResolutions(input), + getInterruptRecoveryState: (input: InterruptRecoveryQuery) => + store.getInterruptRecoveryState(input), listPending: (threadId) => store.listPending(threadId), listPendingByRun: (runId) => store.listPendingByRun(runId), } diff --git a/packages/ai-persistence/src/memory.ts b/packages/ai-persistence/src/memory.ts index 3729c0983..e09d2965b 100644 --- a/packages/ai-persistence/src/memory.ts +++ b/packages/ai-persistence/src/memory.ts @@ -1,6 +1,20 @@ -import { InMemoryLockStore } from '@tanstack/ai' +import { + InMemoryLockStore, + canonicalInterruptJson, + canonicalizeInterruptResolutions, + cloneAndDeepFreezeJson, +} from '@tanstack/ai' import { defineAIPersistence } from './types' -import type { LockStore, ModelMessage } from '@tanstack/ai' +import { hasExactInterruptIds, projectInterruptRecovery } from './interrupts' +import type { + CommitInterruptResolutionsInput, + InterruptCommitResult, + InterruptRecoveryQuery, + InterruptRecoveryStateV1, + LockStore, + ModelMessage, + OpenInterruptBatchInput, +} from '@tanstack/ai' import type { ArtifactRecord, ArtifactStore, @@ -9,8 +23,10 @@ import type { BlobObject, BlobRecord, BlobStore, + InterruptBatchRecord, InterruptRecord, InterruptStore, + LegacyInterruptRecordInput, MessageStore, MetadataStore, RunRecord, @@ -63,34 +79,351 @@ class MemoryRunStore implements RunStore { } class MemoryInterruptStore implements InterruptStore { - private readonly interrupts = new Map() - create(record: Omit): Promise { - this.interrupts.set(record.interruptId, { ...record }) + private interrupts = new Map() + private batches = new Map() + private readonly interruptLocks = new Map>() + + constructor(private readonly clock: () => number) {} + + private async withInterruptLock( + interruptedRunId: string, + operation: () => Promise | T, + ): Promise { + const previous = + this.interruptLocks.get(interruptedRunId) ?? Promise.resolve() + let release = (): void => undefined + const current = new Promise((resolve) => { + release = resolve + }) + const queued = previous.then(() => current) + this.interruptLocks.set(interruptedRunId, queued) + await previous + try { + return await operation() + } finally { + release() + if (this.interruptLocks.get(interruptedRunId) === queued) { + this.interruptLocks.delete(interruptedRunId) + } + } + } + + private rowsForRun(interruptedRunId: string): Array { + return [...this.interrupts.values()].filter( + (record) => record.runId === interruptedRunId, + ) + } + + private recoveryForRun(input: { + threadId: string + interruptedRunId: string + expectedGeneration: number + }): InterruptRecoveryStateV1 { + const rows = this.rowsForRun(input.interruptedRunId).filter( + (record) => record.threadId === input.threadId, + ) + const batch = this.batches.get(input.interruptedRunId) + return projectInterruptRecovery({ + query: { + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + knownGeneration: input.expectedGeneration, + }, + rows, + batch: batch?.threadId === input.threadId ? batch : null, + now: this.clock(), + includeResolutions: true, + }) + } + + create(record: LegacyInterruptRecordInput): Promise { + const generation = record.generation ?? 0 + const binding = + record.binding ?? + ({ + kind: 'generic', + interruptId: record.interruptId, + interruptedRunId: record.runId, + generation, + responseSchemaHash: 'legacy:unknown', + } as const) + this.interrupts.set( + record.interruptId, + cloneAndDeepFreezeJson({ + ...record, + generation, + binding, + }), + ) return Promise.resolve() } + resolve(interruptId: string, response?: unknown): Promise { const existing = this.interrupts.get(interruptId) if (existing) { - this.interrupts.set(interruptId, { - ...existing, - status: 'resolved', - resolvedAt: Date.now(), - response, - }) + this.interrupts.set( + interruptId, + cloneAndDeepFreezeJson({ + ...existing, + status: 'resolved', + resolvedAt: this.clock(), + ...(response !== undefined && { response }), + }), + ) } return Promise.resolve() } + cancel(interruptId: string): Promise { const existing = this.interrupts.get(interruptId) if (existing) { - this.interrupts.set(interruptId, { - ...existing, - status: 'cancelled', - resolvedAt: Date.now(), - }) + this.interrupts.set( + interruptId, + cloneAndDeepFreezeJson({ + ...existing, + status: 'cancelled', + resolvedAt: this.clock(), + }), + ) } return Promise.resolve() } + + openInterruptBatch(input: OpenInterruptBatchInput): Promise<{ + generation: number + descriptors: OpenInterruptBatchInput['descriptors'] + }> { + return this.withInterruptLock(input.interruptedRunId, () => { + const descriptorIds = input.descriptors.map((descriptor) => descriptor.id) + const bindingIds = input.bindings.map((binding) => binding.interruptId) + if ( + descriptorIds.length === 0 || + !hasExactInterruptIds(descriptorIds, bindingIds) + ) { + throw new TypeError( + 'Interrupt descriptors and bindings must contain the same exact nonempty IDs.', + ) + } + if (this.batches.has(input.interruptedRunId)) { + throw new Error('Cannot reopen a committed interrupt batch.') + } + + const existing = this.rowsForRun(input.interruptedRunId) + if (existing.length > 0) { + const generation = existing[0]?.generation + const existingDescriptors = [...existing] + .sort((left, right) => + left.interruptId.localeCompare(right.interruptId), + ) + .map((record) => record.payload) + const requestedDescriptors = [...input.descriptors].sort( + (left, right) => left.id.localeCompare(right.id), + ) + const existingBindings = [...existing] + .map((record) => record.binding) + .sort((left, right) => + left.interruptId.localeCompare(right.interruptId), + ) + const requestedBindings = [...input.bindings] + .map((binding) => ({ + ...binding, + interruptedRunId: input.interruptedRunId, + generation, + })) + .sort((left, right) => + left.interruptId.localeCompare(right.interruptId), + ) + const sameDescriptors = + hasExactInterruptIds( + descriptorIds, + existing.map((record) => record.interruptId), + ) && + canonicalInterruptJson(existingDescriptors) === + canonicalInterruptJson(requestedDescriptors) + const sameBindings = + canonicalInterruptJson(existingBindings) === + canonicalInterruptJson(requestedBindings) + if ( + sameDescriptors && + sameBindings && + generation !== undefined && + existing.every( + (record) => + record.status === 'pending' && + record.threadId === input.threadId && + record.generation === generation, + ) + ) { + return { + generation, + descriptors: cloneAndDeepFreezeJson(input.descriptors), + } + } + throw new Error('An incompatible interrupt batch is already open.') + } + + const generation = 1 + const bindings = new Map( + input.bindings.map((binding) => [binding.interruptId, binding]), + ) + const nextInterrupts = new Map(this.interrupts) + const descriptors = cloneAndDeepFreezeJson(input.descriptors) + for (const descriptor of descriptors) { + if (nextInterrupts.has(descriptor.id)) { + throw new Error(`Interrupt ID already exists: ${descriptor.id}`) + } + const unopened = bindings.get(descriptor.id) + if (!unopened) { + throw new Error(`Missing binding for interrupt: ${descriptor.id}`) + } + const binding = cloneAndDeepFreezeJson({ + ...unopened, + interruptedRunId: input.interruptedRunId, + generation, + }) + nextInterrupts.set( + descriptor.id, + cloneAndDeepFreezeJson({ + interruptId: descriptor.id, + runId: input.interruptedRunId, + threadId: input.threadId, + generation, + status: 'pending', + requestedAt: this.clock(), + payload: descriptor, + binding, + }), + ) + } + this.interrupts = nextInterrupts + return { generation, descriptors } + }) + } + + commitInterruptResolutions( + input: CommitInterruptResolutionsInput, + ): Promise { + return this.withInterruptLock(input.interruptedRunId, () => { + const candidate = canonicalizeInterruptResolutions(input.resolutions) + if ( + candidate.fingerprint !== input.fingerprint || + candidate.canonicalResolutions !== input.canonicalResolutions + ) { + throw new TypeError( + 'Interrupt batch identity does not match its resolutions.', + ) + } + + const winner = this.batches.get(input.interruptedRunId) + if (winner) { + return winner.threadId === input.threadId && + winner.fingerprint === input.fingerprint && + winner.canonicalResolutions === input.canonicalResolutions + ? { + status: 'replayed', + continuationRunId: winner.continuationRunId, + } + : { + status: 'conflict', + authoritativeState: this.recoveryForRun(input), + } + } + + const pending = this.rowsForRun(input.interruptedRunId).filter( + (record) => record.status === 'pending', + ) + const resolutionIds = candidate.resolutions.map( + (resolution) => resolution.interruptId, + ) + if ( + !hasExactInterruptIds( + input.expectedInterruptIds, + pending.map((record) => record.interruptId), + ) || + !hasExactInterruptIds(input.expectedInterruptIds, resolutionIds) || + pending.some((record) => record.threadId !== input.threadId) || + pending.some( + (record) => record.generation !== input.expectedGeneration, + ) || + pending.some( + (record) => + record.binding.expiresAt !== undefined && + Date.parse(record.binding.expiresAt) <= this.clock(), + ) + ) { + return { + status: 'conflict', + authoritativeState: this.recoveryForRun(input), + } + } + + const committedAt = this.clock() + const nextInterrupts = new Map(this.interrupts) + for (const resolution of candidate.resolutions) { + const current = nextInterrupts.get(resolution.interruptId) + if ( + !current || + current.runId !== input.interruptedRunId || + current.status !== 'pending' || + current.generation !== input.expectedGeneration + ) { + throw new Error( + `Pending interrupt changed: ${resolution.interruptId}`, + ) + } + nextInterrupts.set( + resolution.interruptId, + cloneAndDeepFreezeJson({ + ...current, + status: + resolution.status === 'cancelled' ? 'cancelled' : 'resolved', + response: resolution, + resolvedAt: committedAt, + }), + ) + } + + const immutableWinner = cloneAndDeepFreezeJson({ + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + generation: input.expectedGeneration, + expectedInterruptIds: [...input.expectedInterruptIds], + fingerprint: candidate.fingerprint, + canonicalResolutions: candidate.canonicalResolutions, + resolutions: candidate.resolutions, + continuationRunId: input.continuationRunId, + committedAt, + }) + const nextBatches = new Map(this.batches) + nextBatches.set(input.interruptedRunId, immutableWinner) + + this.interrupts = nextInterrupts + this.batches = nextBatches + return { + status: 'committed', + continuationRunId: input.continuationRunId, + } + }) + } + + getInterruptRecoveryState( + input: InterruptRecoveryQuery, + ): Promise { + return this.withInterruptLock(input.interruptedRunId, () => { + const rows = this.rowsForRun(input.interruptedRunId).filter( + (record) => record.threadId === input.threadId, + ) + const batch = this.batches.get(input.interruptedRunId) + return projectInterruptRecovery({ + query: input, + rows, + batch: batch?.threadId === input.threadId ? batch : null, + now: this.clock(), + includeResolutions: true, + }) + }) + } + get(interruptId: string): Promise { return Promise.resolve(this.interrupts.get(interruptId) ?? null) } @@ -348,11 +681,11 @@ class MemoryBlobStore implements BlobStore { } } -export function memoryPersistence() { +export function memoryPersistence(options?: { clock?: () => number }) { const stores: MemoryPersistenceStores = { messages: new MemoryMessageStore(), runs: new MemoryRunStore(), - interrupts: new MemoryInterruptStore(), + interrupts: new MemoryInterruptStore(options?.clock ?? Date.now), metadata: new MemoryMetadataStore(), artifacts: new MemoryArtifactStore(), blobs: new MemoryBlobStore(), diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index b56efe27c..1ea2288ec 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -1,4 +1,14 @@ -import { defineChatMiddleware } from '@tanstack/ai' +import { + canonicalInterruptJson, + canonicalizeInterruptResolutions, + compileJsonSchema202012, + defineChatMiddleware, + digestInterruptJson, + hashSchemaInput, + isStandardSchema, + normalizeApprovalSchema, + validateWithStandardSchema, +} from '@tanstack/ai' import { base64ToUint8Array } from '@tanstack/ai-utils' import { InterruptsCapability, @@ -25,12 +35,19 @@ import type { GenerationFinishInfo, GenerationMiddleware, GenerationMiddlewareContext, + Interrupt, + InterruptBinding, + InterruptRecoveryStateV1, + InterruptSubmissionError, + ItemInterruptErrorCode, PersistedArtifactActivity, PersistedArtifactRef, PersistedArtifactRole, RunAgentResumeItem, StreamChunk, TokenUsage, + ToolApprovalResolution, + UnopenedInterruptBinding, } from '@tanstack/ai' import type { AIPersistence, @@ -88,69 +105,42 @@ const runState = new WeakMap< { merged: boolean; interrupted: boolean } >() -const validResumeStatuses = new Set(['resolved', 'cancelled']) +const interruptBindingMetadataKey = 'tanstack:interruptBinding' -function validatePendingResumes( - pending: Array, - resume: Array | undefined, -): Map { - const pendingInterruptIds = new Set( - pending.map((interrupt) => interrupt.interruptId), - ) - const resumeByInterruptId = new Map( - (resume ?? []).map((entry) => [entry.interruptId, entry]), - ) - if (pending.length === 0) { - const staleEntry = resume?.[0] - if (staleEntry) { - throw new Error( - `Resume entry references non-pending interrupt ${staleEntry.interruptId}.`, - ) - } - return resumeByInterruptId - } - if (!resume || resume.length === 0) { - throw new Error( - `Thread has pending interrupts; resume is required before accepting new input.`, - ) - } +export interface ValidateInterruptResumeBatchInput { + threadId: string + interruptedRunId: string + generation: number + pending: ReadonlyArray + resume?: ReadonlyArray + tools: ChatMiddlewareConfig['tools'] + now?: number +} - for (const interrupt of pending) { - const entry = resumeByInterruptId.get(interrupt.interruptId) - if (!entry) { - throw new Error( - `Missing resume entry for pending interrupt ${interrupt.interruptId}.`, - ) - } - if (!validResumeStatuses.has(entry.status)) { - throw new Error( - `Invalid resume status for pending interrupt ${interrupt.interruptId}: ${entry.status}.`, - ) - } - } - for (const entry of resume) { - if (!pendingInterruptIds.has(entry.interruptId)) { - throw new Error( - `Resume entry references non-pending interrupt ${entry.interruptId}.`, - ) - } +export interface ValidatedInterruptResumeBatch { + errors: ReadonlyArray + resolutions?: ReadonlyArray + canonicalResolutions?: string + fingerprint?: string + resumeToolState?: ChatResumeToolState +} + +export class InterruptResumeValidationError extends Error { + override readonly name = 'InterruptResumeValidationError' + + constructor( + readonly errors: ReadonlyArray, + readonly recovery?: InterruptRecoveryStateV1, + ) { + super(errors.map((error) => error.message).join(' ')) } - return resumeByInterruptId } -async function applyPendingResumes( - pending: Array, - resumeByInterruptId: Map, - interrupts: NonNullable, -): Promise { - for (const interrupt of pending) { - const entry = resumeByInterruptId.get(interrupt.interruptId) - if (!entry) continue - if (entry.status === 'resolved') { - await interrupts.resolve(interrupt.interruptId, entry.payload) - } else { - await interrupts.cancel(interrupt.interruptId) - } +export class InterruptReplaySignal extends Error { + override readonly name = 'InterruptReplaySignal' + + constructor(readonly continuationRunId: string) { + super(`Interrupt resolutions already committed by ${continuationRunId}.`) } } @@ -167,54 +157,681 @@ function stringField( return typeof value[key] === 'string' ? value[key] : undefined } -function interruptKind(interrupt: InterruptRecord): string | undefined { - const metadata = objectValue(interrupt.payload.metadata) - return metadata ? stringField(metadata, 'kind') : undefined +function normalizeIssuePath( + path: ReadonlyArray | undefined, +): ReadonlyArray | undefined { + if (!path) return undefined + return path.map((segment) => { + if (typeof segment === 'string' || typeof segment === 'number') { + return segment + } + const record = objectValue(segment) + const key = record?.key + return typeof key === 'number' ? key : String(key ?? segment) + }) +} + +function itemError( + input: ValidateInterruptResumeBatchInput, + interruptId: string, + code: ItemInterruptErrorCode, + message: string, + options?: { + path?: ReadonlyArray + source?: 'client' | 'server' + retryable?: boolean + }, +): InterruptSubmissionError { + return { + scope: 'item', + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + generation: input.generation, + interruptId, + code, + message, + source: options?.source ?? 'client', + retryable: options?.retryable ?? false, + ...(options?.path ? { path: options.path } : {}), + } +} + +async function validateSchemaValue(input: { + schema: unknown + value: unknown + onIssue: (message: string, path?: ReadonlyArray) => void +}): Promise { + if (isStandardSchema(input.schema)) { + const result = await validateWithStandardSchema( + input.schema, + input.value, + ) + if (!result.success) { + for (const issue of result.issues) { + input.onIssue(issue.message, normalizeIssuePath(issue.path)) + } + } + return + } + + const validate = compileJsonSchema202012(input.schema) + for (const issue of validate(input.value)) { + input.onIssue(issue.message, issue.path) + } } -function resolvedApprovalDecision(entry: RunAgentResumeItem): boolean { - if (entry.status === 'cancelled') return false - const payload = objectValue(entry.payload) - return typeof payload?.approved === 'boolean' ? payload.approved : true +type RuntimeTool = ChatMiddlewareConfig['tools'][number] & { + approvalSchema?: Parameters[0] } -function resumeToolStateFromPending( - pending: Array, - resumeByInterruptId: Map, -): ChatResumeToolState | undefined { - const approvals = new Map() - const clientToolResults = new Map() +function runtimeTool( + tools: ChatMiddlewareConfig['tools'], + name: string, +): RuntimeTool | undefined { + return tools.find((tool) => tool.name === name) as RuntimeTool | undefined +} - for (const interrupt of pending) { - const entry = resumeByInterruptId.get(interrupt.interruptId) +function descriptorResponseSchema(record: InterruptRecord): unknown { + return objectValue(record.payload)?.responseSchema +} + +function schemaHash(schema: unknown): string { + return digestInterruptJson(canonicalInterruptJson(schema)) +} + +async function pushSchemaIssues(input: { + request: ValidateInterruptResumeBatchInput + errors: Array + interruptId: string + schema: unknown + value: unknown + code: ItemInterruptErrorCode + label: string +}): Promise { + try { + await validateSchemaValue({ + schema: input.schema, + value: input.value, + onIssue: (message, path) => { + input.errors.push( + itemError( + input.request, + input.interruptId, + input.code, + `${input.label}: ${message}`, + { path }, + ), + ) + }, + }) + } catch (error) { + input.errors.push( + itemError( + input.request, + input.interruptId, + 'invalid-response-schema', + `${input.label} could not be validated: ${error instanceof Error ? error.message : String(error)}`, + { source: 'server' }, + ), + ) + } +} + +function validateDescriptorSchema( + input: ValidateInterruptResumeBatchInput, + record: InterruptRecord, + binding: InterruptBinding, + errors: Array, +): unknown { + const schema = descriptorResponseSchema(record) + if ( + schema === undefined || + schemaHash(schema) !== binding.responseSchemaHash + ) { + errors.push( + itemError( + input, + record.interruptId, + 'invalid-response-schema', + `Interrupt ${record.interruptId} response schema no longer matches its binding.`, + { source: 'server' }, + ), + ) + } + return schema +} + +export async function validateInterruptResumeBatch( + input: ValidateInterruptResumeBatchInput, +): Promise { + const grouped = new Map>() + const batchErrors: Array = [] + const group = (interruptId: string): Array => { + const existing = grouped.get(interruptId) + if (existing) return existing + const created: Array = [] + grouped.set(interruptId, created) + return created + } + const pendingById = new Map( + input.pending.map((record) => [record.interruptId, record]), + ) + const resumeById = new Map() + const counts = new Map() + for (const entry of input.resume ?? []) { + counts.set(entry.interruptId, (counts.get(entry.interruptId) ?? 0) + 1) + if (!resumeById.has(entry.interruptId)) { + resumeById.set(entry.interruptId, entry) + } + } + + for (const [interruptId, count] of counts) { + if (count > 1) { + group(interruptId).push( + itemError( + input, + interruptId, + 'conflict', + `Interrupt ${interruptId} has duplicate resume entries.`, + ), + ) + } + } + + let incomplete = false + for (const record of input.pending) { + const errors = group(record.interruptId) + const entry = resumeById.get(record.interruptId) + const binding = record.binding + if (!entry) { + incomplete = true + errors.push( + itemError( + input, + record.interruptId, + 'unknown-interrupt', + `Missing resume entry for interrupt ${record.interruptId}.`, + ), + ) + } + if ( + binding.interruptedRunId !== input.interruptedRunId || + binding.generation !== input.generation || + binding.interruptId !== record.interruptId + ) { + errors.push( + itemError( + input, + record.interruptId, + 'stale', + `Interrupt ${record.interruptId} has stale correlation metadata.`, + { source: 'server' }, + ), + ) + } + if ( + binding.expiresAt !== undefined && + Date.parse(binding.expiresAt) <= (input.now ?? Date.now()) + ) { + errors.push( + itemError( + input, + record.interruptId, + 'expired', + `Interrupt ${record.interruptId} has expired.`, + { source: 'server' }, + ), + ) + } + + const responseSchema = validateDescriptorSchema( + input, + record, + binding, + errors, + ) if (!entry) continue + const entryStatus: unknown = entry.status + if (entryStatus !== 'resolved' && entryStatus !== 'cancelled') { + errors.push( + itemError( + input, + record.interruptId, + 'invalid-payload', + `Interrupt ${record.interruptId} has invalid status ${String(entryStatus)}.`, + ), + ) + continue + } + if (binding.kind === 'generic') { + if (entry.status === 'cancelled') { + if (entry.payload !== undefined) { + errors.push( + itemError( + input, + record.interruptId, + 'invalid-payload', + `Cancelled interrupt ${record.interruptId} must not include a payload.`, + ), + ) + } + continue + } + if (responseSchema !== undefined) { + await pushSchemaIssues({ + request: input, + errors, + interruptId: record.interruptId, + schema: responseSchema, + value: entry.payload, + code: 'invalid-payload', + label: `Interrupt ${record.interruptId} payload is invalid`, + }) + } + continue + } - const kind = interruptKind(interrupt) - const reason = stringField(interrupt.payload, 'reason') - const toolCallId = stringField(interrupt.payload, 'toolCallId') + const tool = runtimeTool(input.tools, binding.toolName) + if (!tool) { + errors.push( + itemError( + input, + record.interruptId, + 'stale', + `Tool ${binding.toolName} is unavailable for interrupt ${record.interruptId}.`, + { source: 'server' }, + ), + ) + continue + } + + let approval: ReturnType | undefined + let schemaDrifted = false + if (binding.kind === 'client-tool-execution') { + if (hashSchemaInput(tool.outputSchema) !== binding.outputSchemaHash) { + errors.push( + itemError( + input, + record.interruptId, + 'stale', + `Tool ${binding.toolName} output schema has changed.`, + { source: 'server' }, + ), + ) + schemaDrifted = true + } + } else { + try { + approval = normalizeApprovalSchema( + tool.approvalSchema, + tool.inputSchema, + ) + } catch { + errors.push( + itemError( + input, + record.interruptId, + 'stale', + `Tool ${binding.toolName} approval schema is unavailable.`, + { source: 'server' }, + ), + ) + schemaDrifted = true + } + if ( + approval !== undefined && + (hashSchemaInput(tool.inputSchema) !== binding.inputSchemaHash || + approval.approvalSchemaHash !== binding.approvalSchemaHash || + approval.responseSchemaHash !== binding.responseSchemaHash) + ) { + errors.push( + itemError( + input, + record.interruptId, + 'stale', + `Tool ${binding.toolName} approval schema has changed.`, + { source: 'server' }, + ), + ) + schemaDrifted = true + } + } + + if (entry.status === 'cancelled') { + if (entry.payload !== undefined) { + errors.push( + itemError( + input, + record.interruptId, + 'invalid-payload', + `Cancelled interrupt ${record.interruptId} must not include a payload.`, + ), + ) + } + continue + } + if (schemaDrifted) continue + + if (binding.kind === 'client-tool-execution') { + if (responseSchema !== undefined) { + await pushSchemaIssues({ + request: input, + errors, + interruptId: record.interruptId, + schema: responseSchema, + value: entry.payload, + code: 'invalid-tool-output', + label: `Tool ${binding.toolName} output is invalid`, + }) + } + if (tool.outputSchema !== undefined) { + await pushSchemaIssues({ + request: input, + errors, + interruptId: record.interruptId, + schema: tool.outputSchema, + value: entry.payload, + code: 'invalid-tool-output', + label: `Tool ${binding.toolName} output is invalid`, + }) + } + continue + } - if (kind === 'approval' || reason === 'approval_required') { - approvals.set(interrupt.interruptId, resolvedApprovalDecision(entry)) + if (approval === undefined) continue + const envelope = objectValue(entry.payload) + const approved = + typeof entry.payload === 'boolean' + ? entry.payload + : typeof envelope?.approved === 'boolean' + ? envelope.approved + : undefined + if (approved === undefined) { + errors.push( + itemError( + input, + record.interruptId, + 'invalid-payload', + `Approval ${record.interruptId} must be a boolean or decision envelope.`, + ), + ) continue } + if (envelope) { + await pushSchemaIssues({ + request: input, + errors, + interruptId: record.interruptId, + schema: approval.responseSchema, + value: entry.payload, + code: 'invalid-payload', + label: `Approval ${record.interruptId} envelope is invalid`, + }) + } + if (approved && envelope?.editedArgs !== undefined) { + if (tool.inputSchema === undefined) { + errors.push( + itemError( + input, + record.interruptId, + 'invalid-edited-args', + `Approval ${record.interruptId} cannot edit arguments without an input schema.`, + ), + ) + } else { + await pushSchemaIssues({ + request: input, + errors, + interruptId: record.interruptId, + schema: tool.inputSchema, + value: envelope.editedArgs, + code: 'invalid-edited-args', + label: `Approval ${record.interruptId} edited arguments are invalid`, + }) + } + } + const branch = approved + ? approval.branches.approve + : approval.branches.reject + if (branch && envelope) { + await pushSchemaIssues({ + request: input, + errors, + interruptId: record.interruptId, + schema: branch.source, + value: envelope.payload, + code: 'invalid-payload', + label: `Approval ${record.interruptId} payload is invalid`, + }) + } + } + + for (const entry of input.resume ?? []) { + if (!pendingById.has(entry.interruptId)) { + incomplete = true + group(entry.interruptId).push( + itemError( + input, + entry.interruptId, + 'unknown-interrupt', + `Resume entry references unknown interrupt ${entry.interruptId}.`, + ), + ) + } + } + if (incomplete) { + batchErrors.push({ + scope: 'batch', + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + generation: input.generation, + code: 'incomplete-batch', + message: + 'Resume entries must resolve or cancel the complete interrupt batch.', + source: 'client', + retryable: false, + interruptIds: input.pending.map((record) => record.interruptId), + }) + } + + const itemErrors = [...grouped.entries()] + .sort(([left], [right]) => left.localeCompare(right)) + .flatMap(([, errors]) => errors) + if (itemErrors.length > 0) { + batchErrors.push({ + scope: 'batch', + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + generation: input.generation, + code: 'item-validation-failed', + message: 'One or more interrupt resolutions are invalid.', + source: 'client', + retryable: false, + interruptIds: input.pending.map((record) => record.interruptId), + }) + return { errors: [...itemErrors, ...batchErrors] } + } + + const canonical = canonicalizeInterruptResolutions(input.resume ?? []) + const approvals = new Map() + const clientToolResults = new Map() + const genericInterrupts = new Map< + string, + | { interruptId: string; status: 'resolved'; payload: unknown } + | { + interruptId: string + status: 'cancelled' + } + >() + const deniedToolResults = new Map() + const cancelledToolCallIds = new Set() + + for (const record of input.pending) { + const entry = resumeById.get(record.interruptId) + if (!entry) continue + const binding = record.binding + if (binding.kind === 'generic') { + genericInterrupts.set( + record.interruptId, + entry.status === 'resolved' + ? { + interruptId: record.interruptId, + status: 'resolved', + payload: entry.payload, + } + : { interruptId: record.interruptId, status: 'cancelled' }, + ) + continue + } + if (entry.status === 'cancelled') { + cancelledToolCallIds.add(binding.toolCallId) + continue + } + if (binding.kind === 'client-tool-execution') { + clientToolResults.set(binding.toolCallId, entry.payload) + continue + } + const envelope = objectValue(entry.payload) + const resolution: ToolApprovalResolution = + typeof entry.payload === 'boolean' + ? entry.payload + : envelope?.approved === true + ? { + approved: true, + ...(envelope.editedArgs !== undefined + ? { editedArgs: envelope.editedArgs } + : {}), + ...(envelope.payload !== undefined + ? { payload: envelope.payload } + : {}), + } + : { + approved: false, + ...(envelope?.payload !== undefined + ? { payload: envelope.payload } + : {}), + } + approvals.set(binding.toolCallId, resolution) if ( - entry.status === 'resolved' && - toolCallId && - (kind === 'client_tool' || reason === 'client_tool_input') + resolution === false || + (typeof resolution === 'object' && !resolution.approved) ) { - clientToolResults.set(toolCallId, entry.payload) + deniedToolResults.set( + binding.toolCallId, + typeof resolution === 'object' ? resolution.payload : undefined, + ) + } + } + + return { + errors: [], + resolutions: canonical.resolutions, + canonicalResolutions: canonical.canonicalResolutions, + fingerprint: canonical.fingerprint, + resumeToolState: { + approvals, + clientToolResults, + genericInterrupts, + deniedToolResults, + cancelledToolCallIds, + }, + } +} + +function readUnopenedInterruptBinding( + descriptor: Interrupt, +): UnopenedInterruptBinding | undefined { + const metadata = objectValue(descriptor.metadata) + const raw = metadata + ? objectValue(metadata[interruptBindingMetadataKey]) + : null + if (!raw || stringField(raw, 'interruptId') !== descriptor.id) { + return undefined + } + const kind = stringField(raw, 'kind') + const interruptId = stringField(raw, 'interruptId') + const responseSchemaHash = stringField(raw, 'responseSchemaHash') + const expiresAt = stringField(raw, 'expiresAt') + if (!interruptId || !responseSchemaHash) return undefined + if (kind === 'generic') { + return { + kind, + interruptId, + responseSchemaHash, + ...(expiresAt ? { expiresAt } : {}), + } + } + const toolName = stringField(raw, 'toolName') + const toolCallId = stringField(raw, 'toolCallId') + if (!toolName || !toolCallId) return undefined + if (kind === 'client-tool-execution') { + const outputSchemaHash = stringField(raw, 'outputSchemaHash') + if (!outputSchemaHash) return undefined + return { + kind, + interruptId, + toolName, + toolCallId, + outputSchemaHash, + responseSchemaHash, + ...(expiresAt ? { expiresAt } : {}), + } + } + if (kind === 'tool-approval') { + const inputSchemaHash = stringField(raw, 'inputSchemaHash') + const approvalSchemaHash = stringField(raw, 'approvalSchemaHash') + if (!inputSchemaHash || !approvalSchemaHash) return undefined + return { + kind, + interruptId, + toolName, + toolCallId, + originalArgs: raw.originalArgs, + inputSchemaHash, + approvalSchemaHash, + responseSchemaHash, + ...(expiresAt ? { expiresAt } : {}), } } + return undefined +} - if (approvals.size === 0 && clientToolResults.size === 0) return undefined - return { approvals, clientToolResults } +function withoutInterruptBinding(descriptor: Interrupt): Interrupt { + const metadata = objectValue(descriptor.metadata) + if (!metadata || !(interruptBindingMetadataKey in metadata)) { + return descriptor + } + const publicMetadata = { ...metadata } + delete publicMetadata[interruptBindingMetadataKey] + return { ...descriptor, metadata: publicMetadata } } -function interruptPayload(interrupt: unknown): Record { - return interrupt && typeof interrupt === 'object' - ? { ...(interrupt as Record) } - : { value: interrupt } +function batchError(input: { + threadId: string + interruptedRunId: string + generation: number + interruptIds: ReadonlyArray + code: + | 'stale' + | 'conflict' + | 'persistence-required' + | 'server' + | 'invalid-response-schema' + message: string + retryable?: boolean +}): InterruptSubmissionError { + return { + scope: 'batch', + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + generation: input.generation, + interruptIds: input.interruptIds, + code: input.code, + message: input.message, + source: 'server', + retryable: input.retryable ?? false, + } } function isArtifactRef(value: unknown): value is PersistedArtifactRef { @@ -779,7 +1396,8 @@ async function interruptRun( * Chat-only persistence middleware. Provides durable **state** for `chat()`: * thread messages, run records, interrupts, and locks. Delivery durability * (replaying a disconnected/reloaded stream) lives on the transport layer via - * `StreamDurability`, not here — this middleware never mutates the chunk stream. + * `StreamDurability`, not here. Interrupt terminals are enriched only with the + * store-issued run/generation correlation required by the public binding. */ export function withChatPersistence( persistence: AIPersistence & ValidChatPersistence, @@ -822,22 +1440,89 @@ export function withChatPersistence( let resumeToolState: ChatResumeToolState | undefined if (wantsInterrupts && persistence.stores.interrupts) { - const pending = await persistence.stores.interrupts.listPending( - ctx.threadId, - ) - const resumeByInterruptId = validatePendingResumes( - pending, - config.resume, - ) - resumeToolState = resumeToolStateFromPending( - pending, - resumeByInterruptId, - ) - await applyPendingResumes( - pending, - resumeByInterruptId, - persistence.stores.interrupts, - ) + const pending = ctx.parentRunId + ? await persistence.stores.interrupts.listByRun(ctx.parentRunId) + : await persistence.stores.interrupts.listPending(ctx.threadId) + const interruptedRunId = + ctx.parentRunId ?? pending[0]?.runId ?? ctx.runId + const generation = pending[0]?.generation ?? 0 + if (pending.length > 0 || (config.resume?.length ?? 0) > 0) { + if (!ctx.parentRunId) { + throw new InterruptResumeValidationError([ + batchError({ + threadId: ctx.threadId, + interruptedRunId, + generation, + interruptIds: pending.map((record) => record.interruptId), + code: 'stale', + message: + 'Interrupt continuation requires parentRunId to identify the interrupted run.', + }), + ]) + } + const validated = await validateInterruptResumeBatch({ + threadId: ctx.threadId, + interruptedRunId, + generation, + pending, + resume: config.resume, + tools: config.tools, + }) + if (validated.errors.length > 0) { + throw new InterruptResumeValidationError(validated.errors) + } + if ( + !validated.resolutions || + validated.canonicalResolutions === undefined || + validated.fingerprint === undefined + ) { + throw new InterruptResumeValidationError([ + batchError({ + threadId: ctx.threadId, + interruptedRunId, + generation, + interruptIds: pending.map((record) => record.interruptId), + code: 'server', + message: + 'Validated interrupt batch is missing canonical commit data.', + }), + ]) + } + const expectedInterruptIds = pending + .map((record) => record.interruptId) + .sort((left, right) => left.localeCompare(right)) + const commit = + await persistence.stores.interrupts.commitInterruptResolutions({ + threadId: ctx.threadId, + interruptedRunId, + continuationRunId: ctx.runId, + expectedGeneration: generation, + expectedInterruptIds, + resolutions: validated.resolutions, + canonicalResolutions: validated.canonicalResolutions, + fingerprint: validated.fingerprint, + }) + if (commit.status === 'replayed') { + throw new InterruptReplaySignal(commit.continuationRunId) + } + if (commit.status === 'conflict') { + throw new InterruptResumeValidationError( + [ + batchError({ + threadId: ctx.threadId, + interruptedRunId, + generation, + interruptIds: expectedInterruptIds, + code: 'conflict', + message: + 'Interrupt resolutions conflict with the authoritative committed batch.', + }), + ], + commit.authoritativeState, + ) + } + resumeToolState = validated.resumeToolState + } } await createOrResumeRun(runs, ctx.runId, ctx.threadId) @@ -860,9 +1545,9 @@ export function withChatPersistence( }, async onChunk(ctx: ChatMiddlewareContext, chunk: StreamChunk) { - // State-only: react to the interrupt boundary (create interrupt records, - // mark the run interrupted, snapshot thread messages). The chunk stream is - // never mutated — delivery durability is a transport-layer concern. + // Persist the interrupt boundary before returning it. The only stream + // projection added here is the store-issued run/generation correlation; + // delivery durability remains a transport-layer concern. if ( chunk.type !== 'RUN_FINISHED' || chunk.outcome?.type !== 'interrupt' @@ -872,17 +1557,74 @@ export function withChatPersistence( const state = runState.get(ctx) if (!state) return - if (wantsInterrupts && persistence.stores.interrupts) { - for (const interrupt of chunk.outcome.interrupts) { - await persistence.stores.interrupts.create({ - interruptId: interrupt.id, - runId: ctx.runId, + const interruptIds = chunk.outcome.interrupts.map( + (interrupt) => interrupt.id, + ) + if (!wantsInterrupts || !persistence.stores.interrupts) { + throw new InterruptResumeValidationError([ + batchError({ threadId: ctx.threadId, - status: 'pending', - requestedAt: Date.now(), - payload: interruptPayload(interrupt), - }) + interruptedRunId: ctx.runId, + generation: 0, + interruptIds, + code: 'persistence-required', + message: + 'Interrupt terminal events require an atomic interrupt persistence store.', + }), + ]) + } + const bindings: Array = [] + const descriptors: Array = [] + const bindingErrors: Array = [] + for (const interrupt of chunk.outcome.interrupts) { + const binding = readUnopenedInterruptBinding(interrupt) + if (!binding) { + bindingErrors.push( + itemError( + { + threadId: ctx.threadId, + interruptedRunId: ctx.runId, + generation: 0, + pending: [], + tools: [], + }, + interrupt.id, + 'invalid-response-schema', + `Interrupt ${interrupt.id} is missing a valid server binding.`, + { source: 'server' }, + ), + ) + continue } + bindings.push(binding) + descriptors.push(withoutInterruptBinding(interrupt)) + } + if (bindingErrors.length > 0) { + throw new InterruptResumeValidationError(bindingErrors) + } + let opened: Awaited< + ReturnType + > + try { + opened = await persistence.stores.interrupts.openInterruptBatch({ + threadId: ctx.threadId, + interruptedRunId: ctx.runId, + descriptors, + bindings, + }) + } catch (error) { + if (error instanceof InterruptResumeValidationError) throw error + throw new InterruptResumeValidationError([ + batchError({ + threadId: ctx.threadId, + interruptedRunId: ctx.runId, + generation: 0, + interruptIds, + code: 'server', + message: `Failed to persist interrupt batch: ${error instanceof Error ? error.message : String(error)}`, + retryable: true, + }), + ]) } await interruptRun(runs, ctx.runId) if (wantsMessages && persistence.stores.messages) { @@ -891,6 +1633,34 @@ export function withChatPersistence( ]) } state.interrupted = true + const openedBindings = new Map( + bindings.map((binding) => [ + binding.interruptId, + { + ...binding, + interruptedRunId: ctx.runId, + generation: opened.generation, + } satisfies InterruptBinding, + ]), + ) + return { + ...chunk, + outcome: { + ...chunk.outcome, + interrupts: chunk.outcome.interrupts.map((interrupt) => { + const binding = openedBindings.get(interrupt.id) + return binding + ? { + ...interrupt, + metadata: { + ...interrupt.metadata, + [interruptBindingMetadataKey]: binding, + }, + } + : interrupt + }), + }, + } }, async onFinish(ctx: ChatMiddlewareContext, info: FinishInfo) { diff --git a/packages/ai-persistence/src/recovery.ts b/packages/ai-persistence/src/recovery.ts new file mode 100644 index 000000000..fdf085967 --- /dev/null +++ b/packages/ai-persistence/src/recovery.ts @@ -0,0 +1,179 @@ +import { cloneAndDeepFreezeJson } from '@tanstack/ai' +import type { + InterruptPersistenceGateway, + InterruptRecoveryQuery, + InterruptRecoveryStateV1, + InterruptSubmissionError, +} from '@tanstack/ai' + +export type InterruptRecoveryAuthorization = + | { authorized: false } + | { authorized: true; includeResolutions: boolean } + +export interface InterruptRecoveryHandlerOptions { + gateway: InterruptPersistenceGateway + authorize: ( + request: Request, + input: InterruptRecoveryQuery, + ) => InterruptRecoveryAuthorization | Promise +} + +export interface GetInterruptRecoveryStateOptions { + includeResolutions?: boolean +} + +function redactCommittedResolutions( + state: InterruptRecoveryStateV1, +): InterruptRecoveryStateV1 { + if (!state.committed || state.committed.resolutions === undefined) { + return state + } + return cloneAndDeepFreezeJson({ + ...state, + committed: { + fingerprint: state.committed.fingerprint, + ...(state.committed.continuationRunId + ? { continuationRunId: state.committed.continuationRunId } + : {}), + committedAt: state.committed.committedAt, + }, + }) +} + +export async function getInterruptRecoveryState( + gateway: InterruptPersistenceGateway, + query: InterruptRecoveryQuery, + options?: GetInterruptRecoveryStateOptions, +): Promise { + const state = await gateway.getInterruptRecoveryState(query) + return options?.includeResolutions === true + ? state + : redactCommittedResolutions(state) +} + +function recoveryError(input: { + threadId: string + interruptedRunId: string + generation: number + code: 'recovery-unavailable' | 'protocol' | 'server' + message: string + source?: 'client' | 'server' + retryable?: boolean +}): InterruptSubmissionError { + return { + scope: 'batch', + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + generation: input.generation, + interruptIds: [], + code: input.code, + message: input.message, + source: input.source ?? 'server', + retryable: input.retryable ?? false, + } +} + +function jsonResponse(body: unknown, status: number): Response { + return new Response(JSON.stringify(body), { + status, + headers: { 'content-type': 'application/json' }, + }) +} + +export function createInterruptRecoveryHandler( + options: InterruptRecoveryHandlerOptions, +): (request: Request) => Promise { + return async (request) => { + const url = new URL(request.url) + const threadId = url.searchParams.get('threadId') ?? '' + const interruptedRunId = url.searchParams.get('interruptedRunId') ?? '' + const knownGenerationText = url.searchParams.get('knownGeneration') + const knownGeneration = + knownGenerationText === null ? Number.NaN : Number(knownGenerationText) + if ( + !threadId || + !interruptedRunId || + !Number.isSafeInteger(knownGeneration) || + knownGeneration < 0 + ) { + return jsonResponse( + { + errors: [ + recoveryError({ + threadId, + interruptedRunId, + generation: Number.isSafeInteger(knownGeneration) + ? knownGeneration + : 0, + code: 'protocol', + message: + 'Recovery requires threadId, interruptedRunId, and a non-negative integer knownGeneration.', + source: 'client', + }), + ], + }, + 400, + ) + } + const input = { threadId, interruptedRunId, knownGeneration } + let authorization: InterruptRecoveryAuthorization + try { + authorization = await options.authorize(request, input) + } catch (error) { + return jsonResponse( + { + errors: [ + recoveryError({ + threadId: '', + interruptedRunId: '', + generation: 0, + code: 'server', + message: `Interrupt recovery authorization failed: ${error instanceof Error ? error.message : String(error)}`, + retryable: true, + }), + ], + }, + 500, + ) + } + if (!authorization.authorized) { + return jsonResponse( + { + errors: [ + recoveryError({ + threadId: '', + interruptedRunId: '', + generation: 0, + code: 'recovery-unavailable', + message: 'Interrupt recovery is not authorized.', + }), + ], + }, + 401, + ) + } + + try { + const state = await getInterruptRecoveryState(options.gateway, input, { + includeResolutions: authorization.includeResolutions, + }) + return jsonResponse(state, 200) + } catch (error) { + return jsonResponse( + { + errors: [ + recoveryError({ + threadId, + interruptedRunId, + generation: knownGeneration, + code: 'recovery-unavailable', + message: `Interrupt recovery failed: ${error instanceof Error ? error.message : String(error)}`, + retryable: true, + }), + ], + }, + 503, + ) + } + } +} diff --git a/packages/ai-persistence/src/testkit/conformance.ts b/packages/ai-persistence/src/testkit/conformance.ts index e2e80edfc..cb62aac97 100644 --- a/packages/ai-persistence/src/testkit/conformance.ts +++ b/packages/ai-persistence/src/testkit/conformance.ts @@ -11,7 +11,9 @@ * hold identically for byte-storing (memory) and reference-only (SQL) backends. */ import { beforeAll, describe, expect, it } from 'vitest' -import type { AIPersistence } from '../types' +import { canonicalizeInterruptResolutions } from '@tanstack/ai' +import type { CommitInterruptResolutionsInput } from '@tanstack/ai' +import type { AIPersistence, InterruptStore } from '../types' type MakePersistence = () => Promise | AIPersistence @@ -352,3 +354,267 @@ export function runPersistenceConformance( }) }) } + +export interface InterruptConformanceHarness { + getStore: () => InterruptStore | undefined + advanceBy: (milliseconds: number) => void + inspect: (interruptedRunId: string) => Promise<{ + statuses: ReadonlyArray + batchCount: number + }> + failTransitionOnce: (interruptId: string) => void + reopen?: () => Promise +} + +export async function openTwoInterrupts(store: InterruptStore) { + return store.openInterruptBatch({ + threadId: 'thread-cas', + interruptedRunId: 'run-interrupted', + descriptors: [ + { id: 'int-a', reason: 'confirmation' }, + { + id: 'int-b', + reason: 'input_required', + expiresAt: '2026-07-13T10:01:00.000Z', + }, + ], + bindings: [ + { + interruptId: 'int-a', + kind: 'generic', + responseSchemaHash: 'sha256:a', + }, + { + interruptId: 'int-b', + kind: 'generic', + responseSchemaHash: 'sha256:b', + expiresAt: '2026-07-13T10:01:00.000Z', + }, + ], + }) +} + +export function validTwoItemCommit( + generation: number, + continuationRunId = 'run-continuation-a', +): CommitInterruptResolutionsInput { + const candidate = canonicalizeInterruptResolutions([ + { interruptId: 'int-b', status: 'cancelled' }, + { + interruptId: 'int-a', + status: 'resolved', + payload: { ok: true }, + }, + ]) + return { + threadId: 'thread-cas', + interruptedRunId: 'run-interrupted', + continuationRunId, + expectedGeneration: generation, + expectedInterruptIds: ['int-b', 'int-a'], + resolutions: candidate.resolutions, + fingerprint: candidate.fingerprint, + canonicalResolutions: candidate.canonicalResolutions, + } +} + +export function runInterruptStoreConformance( + createHarness: () => Promise, +): void { + describe('atomic interrupt store conformance', () => { + it('requires the atomic interrupt capability', async () => { + const harness = await createHarness() + expect( + harness.getStore(), + 'interrupt conformance requires stores.interrupts', + ).toBeDefined() + if (!harness.getStore()) throw new Error('interrupt store missing') + }) + + it('keeps canonical order stable and rejects malformed exact sets', async () => { + const harness = await createHarness() + const store = harness.getStore() + if (!store) throw new Error('interrupt store missing') + const opened = await openTwoInterrupts(store) + const ordered = validTwoItemCommit(opened.generation) + const reordered = canonicalizeInterruptResolutions([ + { + interruptId: 'int-a', + status: 'resolved', + payload: { ok: true }, + }, + { interruptId: 'int-b', status: 'cancelled' }, + ]) + expect(reordered.fingerprint).toBe(ordered.fingerprint) + expect(reordered.canonicalResolutions).toBe(ordered.canonicalResolutions) + + for (const expectedInterruptIds of [ + ['int-a'], + ['int-a', 'int-b', 'int-c'], + ['int-a', 'int-a'], + ]) { + await expect( + store.commitInterruptResolutions({ + ...ordered, + expectedInterruptIds, + }), + ).resolves.toMatchObject({ status: 'conflict' }) + } + await expect( + store.commitInterruptResolutions({ + ...ordered, + expectedGeneration: opened.generation + 1, + }), + ).resolves.toMatchObject({ status: 'conflict' }) + await expect( + store.commitInterruptResolutions({ + ...ordered, + threadId: 'thread-other', + }), + ).resolves.toMatchObject({ + status: 'conflict', + authoritativeState: { state: 'missing' }, + }) + }) + + it('checks canonical bytes and digest before any transition', async () => { + const harness = await createHarness() + const store = harness.getStore() + if (!store) throw new Error('interrupt store missing') + const opened = await openTwoInterrupts(store) + const input = validTwoItemCommit(opened.generation) + await expect( + store.commitInterruptResolutions({ + ...input, + canonicalResolutions: `${input.canonicalResolutions} `, + }), + ).rejects.toThrow(/identity/i) + await expect(harness.inspect('run-interrupted')).resolves.toEqual({ + statuses: ['pending', 'pending'], + batchCount: 0, + }) + }) + + it('rolls back every row when a transition fails', async () => { + const harness = await createHarness() + const store = harness.getStore() + if (!store) throw new Error('interrupt store missing') + const opened = await openTwoInterrupts(store) + harness.failTransitionOnce('int-b') + await expect( + store.commitInterruptResolutions(validTwoItemCommit(opened.generation)), + ).rejects.toThrow() + await expect(harness.inspect('run-interrupted')).resolves.toEqual({ + statuses: ['pending', 'pending'], + batchCount: 0, + }) + }) + + it('projects pending, committed replay, restart, and one concurrent winner', async () => { + const harness = await createHarness() + const store = harness.getStore() + if (!store) throw new Error('interrupt store missing') + const opened = await openTwoInterrupts(store) + await expect( + store.getInterruptRecoveryState({ + threadId: 'thread-cas', + interruptedRunId: 'run-interrupted', + knownGeneration: opened.generation, + }), + ).resolves.toMatchObject({ + state: 'pending', + generation: opened.generation, + }) + + const first = validTwoItemCommit(opened.generation, 'run-continuation-a') + const secondCandidate = canonicalizeInterruptResolutions([ + { interruptId: 'int-a', status: 'cancelled' }, + { interruptId: 'int-b', status: 'cancelled' }, + ]) + const second = { + ...first, + continuationRunId: 'run-continuation-b', + resolutions: secondCandidate.resolutions, + fingerprint: secondCandidate.fingerprint, + canonicalResolutions: secondCandidate.canonicalResolutions, + } + const results = await Promise.all([ + store.commitInterruptResolutions(first), + store.commitInterruptResolutions(second), + ]) + expect(results.map((result) => result.status).sort()).toEqual([ + 'committed', + 'conflict', + ]) + const winningRunId = results.find( + (result) => result.status === 'committed', + )?.continuationRunId + expect(winningRunId).toBeDefined() + const replayInput = winningRunId === 'run-continuation-a' ? first : second + await expect( + store.commitInterruptResolutions({ + ...replayInput, + continuationRunId: 'ignored-replay-run', + }), + ).resolves.toEqual({ + status: 'replayed', + continuationRunId: winningRunId, + }) + + const recoveredStore = harness.reopen ? await harness.reopen() : store + await expect( + recoveredStore.getInterruptRecoveryState({ + threadId: 'thread-cas', + interruptedRunId: 'run-interrupted', + knownGeneration: opened.generation, + }), + ).resolves.toMatchObject({ + state: 'committed', + committed: { continuationRunId: winningRunId }, + }) + }) + + it('projects missing and legacy-committed recovery states', async () => { + const harness = await createHarness() + const store = harness.getStore() + if (!store) throw new Error('interrupt store missing') + await expect( + store.getInterruptRecoveryState({ + threadId: 'thread-cas', + interruptedRunId: 'run-missing', + knownGeneration: 3, + }), + ).resolves.toMatchObject({ state: 'missing', generation: 3 }) + + await store.create({ + interruptId: 'legacy-int', + runId: 'legacy-run', + threadId: 'thread-cas', + status: 'resolved', + requestedAt: 1, + payload: { id: 'legacy-int', reason: 'confirmation' }, + }) + await expect( + store.getInterruptRecoveryState({ + threadId: 'thread-cas', + interruptedRunId: 'legacy-run', + knownGeneration: 0, + }), + ).resolves.toMatchObject({ state: 'legacy-committed' }) + }) + + it('marks an uncommitted expired batch expired', async () => { + const harness = await createHarness() + const store = harness.getStore() + if (!store) throw new Error('interrupt store missing') + const opened = await openTwoInterrupts(store) + harness.advanceBy(60_001) + await expect( + store.commitInterruptResolutions(validTwoItemCommit(opened.generation)), + ).resolves.toMatchObject({ + status: 'conflict', + authoritativeState: { state: 'expired' }, + }) + }) + }) +} diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index bd3734da3..61a4c3071 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -1,4 +1,11 @@ -import type { LockStore, ModelMessage, TokenUsage } from '@tanstack/ai' +import type { + InterruptBinding, + InterruptPersistenceGateway, + LockStore, + ModelMessage, + RunAgentResumeItem, + TokenUsage, +} from '@tanstack/ai' export interface MessageStore { loadThread: (threadId: string) => Promise> @@ -37,16 +44,41 @@ export interface InterruptRecord { interruptId: string runId: string threadId: string + generation: number status: 'pending' | 'resolved' | 'cancelled' requestedAt: number resolvedAt?: number - payload: Record + payload: unknown + binding: InterruptBinding response?: unknown } -export interface InterruptStore { - create: (record: Omit) => Promise +export interface InterruptBatchRecord { + threadId: string + interruptedRunId: string + generation: number + expectedInterruptIds: ReadonlyArray + fingerprint: string + canonicalResolutions: string + resolutions: ReadonlyArray + continuationRunId: string + committedAt: number +} + +export type LegacyInterruptRecordInput = Omit< + InterruptRecord, + 'binding' | 'generation' | 'resolvedAt' +> & { + binding?: InterruptBinding + generation?: number +} + +export interface InterruptStore extends InterruptPersistenceGateway { + /** @deprecated Use openInterruptBatch for native interrupt writes. */ + create: (record: LegacyInterruptRecordInput) => Promise + /** @deprecated Use commitInterruptResolutions for native interrupt writes. */ resolve: (interruptId: string, response?: unknown) => Promise + /** @deprecated Use commitInterruptResolutions for native interrupt writes. */ cancel: (interruptId: string) => Promise get: (interruptId: string) => Promise list: (threadId: string) => Promise> @@ -271,6 +303,17 @@ function assertKnownStoreKeys(stores: object, location: string): void { export function validatePersistenceStoreKeys(persistence: AIPersistence): void { assertKnownStoreKeys(persistence.stores, 'store') + const interrupts = persistence.stores.interrupts + if ( + interrupts !== undefined && + (typeof interrupts.openInterruptBatch !== 'function' || + typeof interrupts.commitInterruptResolutions !== 'function' || + typeof interrupts.getInterruptRecoveryState !== 'function') + ) { + throw new Error( + 'atomic-commit-unsupported: stores.interrupts must implement the atomic interrupt gateway.', + ) + } } export function validateChatPersistenceStores( diff --git a/packages/ai-persistence/tests/interrupts.test.ts b/packages/ai-persistence/tests/interrupts.test.ts index f900b180d..4f4e3e26c 100644 --- a/packages/ai-persistence/tests/interrupts.test.ts +++ b/packages/ai-persistence/tests/interrupts.test.ts @@ -1,610 +1,639 @@ import { describe, expect, it, vi } from 'vitest' -import { EventType, chat } from '@tanstack/ai' -import type { AnyTextAdapter, StreamChunk, Tool } from '@tanstack/ai' +import { + EventType, + canonicalInterruptJson, + digestInterruptJson, + chat, + hashSchemaInput, + normalizeApprovalSchema, + toolDefinition, +} from '@tanstack/ai' +import { + createInterruptRecoveryHandler, + getInterruptRecoveryState, +} from '../src' import { memoryPersistence } from '../src/memory' -import { withChatPersistence } from '../src/middleware' +import { + validateInterruptResumeBatch, + withChatPersistence, +} from '../src/middleware' +import type { + AnyTextAdapter, + InterruptBinding, + SchemaInput, + StreamChunk, +} from '@tanstack/ai' +import type { InterruptRecord } from '../src' + +const bindingMetadataKey = 'tanstack:interruptBinding' function mockAdapter(iterations: Array>) { const calls: Array = [] - let i = 0 - const adapter = { + let index = 0 + const adapter: AnyTextAdapter = { kind: 'text', name: 'mock', model: 'test-model', - '~types': {}, - chatStream: (opts: unknown) => { - calls.push(opts) - const chunks = iterations[i] ?? [] - i++ + '~types': { + providerOptions: {}, + inputModalities: ['text'], + messageMetadataByModality: { + text: undefined, + image: undefined, + audio: undefined, + video: undefined, + document: undefined, + }, + toolCapabilities: [], + toolCallMetadata: undefined, + systemPromptMetadata: undefined, + }, + chatStream: (options) => { + calls.push(options) + const chunks = iterations[index] ?? [] + index++ return (async function* () { - for (const c of chunks) yield c + for (const chunk of chunks) yield chunk })() }, structuredOutput: async () => ({ data: {}, rawText: '{}' }), - } as unknown as AnyTextAdapter + } return { adapter, calls } } async function collect(stream: AsyncIterable) { - const out: Array = [] - for await (const c of stream) out.push(c) - return out + const chunks: Array = [] + for await (const chunk of stream) chunks.push(chunk) + return chunks } -const interruptFinished = (): StreamChunk => ({ - type: EventType.RUN_FINISHED, - runId: 'r1', - threadId: 't1', - finishReason: 'tool_calls', - timestamp: 1, - outcome: { - type: 'interrupt', - interrupts: [ - { - id: 'interrupt-1', - reason: 'tool_call', - message: 'Approve the tool call?', - toolCallId: 'tool-call-1', - }, - ], - }, -}) - -const runStarted = (): StreamChunk => ({ - type: EventType.RUN_STARTED, - runId: 'r1', - threadId: 't1', - timestamp: 1, -}) - -const toolStart = (): StreamChunk => ({ - type: EventType.TOOL_CALL_START, - toolCallId: 'tool-call-1', - toolCallName: 'clientSearch', - toolName: 'clientSearch', - timestamp: 1, -}) - -const toolArgs = (): StreamChunk => ({ - type: EventType.TOOL_CALL_ARGS, - toolCallId: 'tool-call-1', - delta: '{"query":"test"}', - timestamp: 1, -}) +const responseSchema = { + type: 'object', + properties: { value: { type: 'number' } }, + required: ['value'], + additionalProperties: false, +} -const text = (delta: string): StreamChunk => ({ - type: EventType.TEXT_MESSAGE_CONTENT, - messageId: 'm1', - delta, - timestamp: 1, -}) +function responseSchemaHash(): string { + return digestInterruptJson(canonicalInterruptJson(responseSchema)) +} -const runFinished = (runId = 'r1'): StreamChunk => ({ - type: EventType.RUN_FINISHED, - runId, - threadId: 't1', - finishReason: 'stop', - timestamp: 1, -}) +function interruptFinished(options?: { + includeBinding?: boolean +}): StreamChunk { + return { + type: EventType.RUN_FINISHED, + runId: 'interrupted-run', + threadId: 'thread-1', + finishReason: 'tool_calls', + timestamp: 1, + outcome: { + type: 'interrupt', + interrupts: [ + { + id: 'interrupt-1', + reason: 'generic', + responseSchema, + ...(options?.includeBinding === false + ? {} + : { + metadata: { + [bindingMetadataKey]: { + kind: 'generic', + interruptId: 'interrupt-1', + responseSchemaHash: responseSchemaHash(), + }, + }, + }), + }, + ], + }, + } +} -const clientTool = (name: string): Tool => ({ - name, - description: `${name} client tool`, -}) +function successfulRun(runId: string): Array { + return [ + { + type: EventType.RUN_STARTED, + runId, + threadId: 'thread-1', + timestamp: 1, + }, + { + type: EventType.RUN_FINISHED, + runId, + threadId: 'thread-1', + finishReason: 'stop', + timestamp: 2, + }, + ] +} -const approvalClientTool = (name: string): Tool => ({ - ...clientTool(name), - needsApproval: true, -}) +function pendingRecord(input: { + interruptId: string + binding: InterruptBinding + payload?: unknown +}): InterruptRecord { + return { + interruptId: input.interruptId, + runId: 'interrupted-run', + threadId: 'thread-1', + generation: 1, + status: 'pending', + requestedAt: 1, + payload: + input.payload ?? + ({ + id: input.interruptId, + reason: 'generic', + responseSchema, + } satisfies Record), + binding: input.binding, + } +} -describe('interrupt persistence', () => { - it('persists RUN_FINISHED interrupt outcomes as pending interrupt records', async () => { +describe('authoritative interrupt persistence', () => { + it('opens one atomic batch before exposing an interrupt terminal', async () => { const persistence = memoryPersistence() - const { adapter } = mockAdapter([[runStarted(), interruptFinished()]]) + const open = vi.spyOn(persistence.stores.interrupts!, 'openInterruptBatch') + const legacyCreate = vi.spyOn(persistence.stores.interrupts!, 'create') + const { adapter } = mockAdapter([ + [interruptFinished({ includeBinding: true })], + ]) const chunks = await collect( chat({ adapter, - messages: [{ role: 'user', content: 'hi' }], - runId: 'r1', - threadId: 't1', + messages: [{ role: 'user', content: 'pause' }], + runId: 'interrupted-run', + threadId: 'thread-1', middleware: [withChatPersistence(persistence)], }) as AsyncIterable, ) - const pending = await persistence.stores.interrupts!.listPending('t1') - expect(pending).toHaveLength(1) - expect(pending[0]?.interruptId).toBe('interrupt-1') - expect((await persistence.stores.runs!.get('r1'))?.status).toBe( - 'interrupted', - ) - // Persistence is state-only: it never stamps delivery cursors on the stream. - expect(chunks.every((chunk) => !('cursor' in chunk))).toBe(true) + expect(open).toHaveBeenCalledTimes(1) + expect(open).toHaveBeenCalledWith({ + threadId: 'thread-1', + interruptedRunId: 'interrupted-run', + descriptors: expect.arrayContaining([ + expect.objectContaining({ id: 'interrupt-1' }), + ]), + bindings: [ + { + kind: 'generic', + interruptId: 'interrupt-1', + responseSchemaHash: responseSchemaHash(), + }, + ], + }) + expect(legacyCreate).not.toHaveBeenCalled() + expect(chunks.at(-1)).toMatchObject({ + type: EventType.RUN_FINISHED, + outcome: { type: 'interrupt' }, + }) }) - it('saves thread messages when a messages-enabled run pauses on an interrupt', async () => { + it('reports a terminal with no reserved server binding before visibility', async () => { const persistence = memoryPersistence() - const { adapter } = mockAdapter([[runStarted(), interruptFinished()]]) + const { adapter } = mockAdapter([ + [interruptFinished({ includeBinding: false })], + ]) - await collect( + const chunks = await collect( chat({ adapter, - messages: [{ role: 'user', content: 'hi' }], - runId: 'r1', - threadId: 't1', + messages: [{ role: 'user', content: 'pause' }], + runId: 'interrupted-run', + threadId: 'thread-1', middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, + stream: true, + }), ) - expect(await persistence.stores.messages!.loadThread('t1')).toEqual([ - { role: 'user', content: 'hi' }, + expect(chunks).toEqual([ + expect.objectContaining({ + type: EventType.RUN_ERROR, + 'tanstack:interruptErrors': expect.arrayContaining([ + expect.objectContaining({ + scope: 'item', + interruptId: 'interrupt-1', + code: 'invalid-response-schema', + }), + ]), + }), ]) + expect( + await persistence.stores.interrupts!.listPendingByRun('interrupted-run'), + ).toEqual([]) }) - it('does not persist duplicate records before terminal interrupt outcome', async () => { - const persistence = memoryPersistence() - const create = vi.spyOn(persistence.stores.interrupts!, 'create') - const { adapter } = mockAdapter([ - [ - runStarted(), - toolStart(), - toolArgs(), + it('reports duplicate, extra, missing, expiry, and payload failures together', async () => { + const expiredBinding: InterruptBinding = { + kind: 'generic', + interruptId: 'interrupt-a', + interruptedRunId: 'interrupted-run', + generation: 1, + responseSchemaHash: responseSchemaHash(), + expiresAt: '2026-07-13T10:00:00.000Z', + } + const missingBinding: InterruptBinding = { + kind: 'generic', + interruptId: 'interrupt-b', + interruptedRunId: 'interrupted-run', + generation: 1, + responseSchemaHash: responseSchemaHash(), + } + + const result = await validateInterruptResumeBatch({ + threadId: 'thread-1', + interruptedRunId: 'interrupted-run', + generation: 1, + pending: [ + pendingRecord({ interruptId: 'interrupt-a', binding: expiredBinding }), + pendingRecord({ interruptId: 'interrupt-b', binding: missingBinding }), + ], + resume: [ { - type: EventType.RUN_FINISHED, - runId: 'r1', - threadId: 't1', - finishReason: 'tool_calls', - timestamp: 1, + interruptId: 'interrupt-a', + status: 'resolved', + payload: { value: 'not-a-number' }, }, + { interruptId: 'interrupt-a', status: 'cancelled', payload: true }, + { interruptId: 'interrupt-extra', status: 'resolved' }, ], - ]) - - await collect( - chat({ - adapter, - messages: [{ role: 'user', content: 'hi' }], - tools: [approvalClientTool('clientSearch')], - runId: 'r1', - threadId: 't1', - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ) + tools: [], + now: Date.parse('2026-07-13T10:00:00.001Z'), + }) - expect(create).toHaveBeenCalledTimes(1) - expect(await persistence.stores.interrupts!.listPending('t1')).toHaveLength( - 1, + expect(result.errors).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + scope: 'item', + interruptId: 'interrupt-a', + code: 'conflict', + }), + expect.objectContaining({ + scope: 'item', + interruptId: 'interrupt-a', + code: 'expired', + }), + expect.objectContaining({ + scope: 'item', + interruptId: 'interrupt-a', + code: 'invalid-payload', + }), + expect.objectContaining({ + scope: 'item', + interruptId: 'interrupt-b', + code: 'unknown-interrupt', + }), + expect.objectContaining({ + scope: 'item', + interruptId: 'interrupt-extra', + code: 'unknown-interrupt', + }), + expect.objectContaining({ + scope: 'batch', + code: 'incomplete-batch', + }), + ]), ) }) - it('blocks normal new input while a thread has pending interrupts', async () => { - const persistence = memoryPersistence() - await persistence.stores.interrupts!.create({ - interruptId: 'interrupt-1', - runId: 'old-run', - threadId: 't1', - status: 'pending', - requestedAt: 1, - payload: {}, + it('rejects cancelled tool interrupts after removal or schema drift', async () => { + const original = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema: { + type: 'object', + properties: { cents: { type: 'number' } }, + required: ['cents'], + additionalProperties: false, + }, }) - const { adapter } = mockAdapter([[interruptFinished()]]) - - await expect( - collect( - chat({ - adapter, - messages: [{ role: 'user', content: 'new input' }], - runId: 'r2', - threadId: 't1', - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ), - ).rejects.toThrow(/pending interrupt/i) - - expect(await persistence.stores.runs!.get('r2')).toBeNull() - }) - - it('treats resume entries as interrupt continuation on the same run', async () => { - const persistence = memoryPersistence() - const first = mockAdapter([[runStarted(), interruptFinished()]]) - await collect( - chat({ - adapter: first.adapter, - messages: [{ role: 'user', content: 'hi' }], - runId: 'r1', - threadId: 't1', - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ) - expect(await persistence.stores.interrupts!.listPending('t1')).toHaveLength( - 1, + const drifted = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema: { + type: 'object', + properties: { amount: { type: 'number' } }, + required: ['amount'], + additionalProperties: false, + }, + }) + const approval = normalizeApprovalSchema( + original.approvalSchema, + original.inputSchema, ) - - const continuation = mockAdapter([ - [runStarted(), text('continued'), runFinished('r1')], - ]) - const chunks = await collect( - chat({ - adapter: continuation.adapter, - messages: [], - runId: 'r1', - threadId: 't1', - resume: [ - { - interruptId: 'interrupt-1', - status: 'resolved', - payload: { approved: true }, + const toolBinding = ( + interruptId: string, + toolName: string, + ): InterruptBinding => ({ + kind: 'tool-approval', + interruptId, + interruptedRunId: 'interrupted-run', + generation: 1, + toolName, + toolCallId: `call-${interruptId}`, + originalArgs: { cents: 1 }, + inputSchemaHash: hashSchemaInput(original.inputSchema), + approvalSchemaHash: approval.approvalSchemaHash, + responseSchemaHash: approval.responseSchemaHash, + }) + const result = await validateInterruptResumeBatch({ + threadId: 'thread-1', + interruptedRunId: 'interrupted-run', + generation: 1, + pending: [ + pendingRecord({ + interruptId: 'removed', + binding: toolBinding('removed', 'removed-tool'), + payload: { + id: 'removed', + reason: 'tool_call', + responseSchema: approval.responseSchema, }, - ], - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ) + }), + pendingRecord({ + interruptId: 'drifted', + binding: toolBinding('drifted', 'transfer'), + payload: { + id: 'drifted', + reason: 'tool_call', + responseSchema: approval.responseSchema, + }, + }), + ], + resume: [ + { interruptId: 'removed', status: 'cancelled' }, + { interruptId: 'drifted', status: 'cancelled' }, + ], + tools: [drifted], + }) - expect(continuation.calls).toHaveLength(1) - expect(chunks).toContainEqual( - expect.objectContaining({ delta: 'continued' }), + expect(result.errors).toEqual( + expect.arrayContaining([ + expect.objectContaining({ interruptId: 'removed', code: 'stale' }), + expect.objectContaining({ interruptId: 'drifted', code: 'stale' }), + ]), ) - expect(await persistence.stores.interrupts!.listPending('t1')).toEqual([]) - expect( - (await persistence.stores.interrupts!.get('interrupt-1'))?.status, - ).toBe('resolved') + expect(result.resolutions).toBeUndefined() }) - it('applies persisted approval and client-tool resume decisions with empty client messages', async () => { - const persistence = memoryPersistence() - const toolCallChunks = () => [ - runStarted(), - toolStart(), - toolArgs(), - { - type: EventType.RUN_FINISHED, - runId: 'r1', - threadId: 't1', - finishReason: 'tool_calls', - timestamp: 1, - } as StreamChunk, - ] - const first = mockAdapter([toolCallChunks()]) - await collect( - chat({ - adapter: first.adapter, - messages: [{ role: 'user', content: 'hi' }], - tools: [approvalClientTool('clientSearch')], - runId: 'r1', - threadId: 't1', - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ) - - const approvalInterrupt = await persistence.stores.interrupts!.get( - 'approval_tool-call-1', - ) - expect(approvalInterrupt?.status).toBe('pending') - - const afterApproval = mockAdapter([toolCallChunks()]) - const approvalChunks = await collect( - chat({ - adapter: afterApproval.adapter, - messages: [], - tools: [approvalClientTool('clientSearch')], - runId: 'r1', - threadId: 't1', - resume: [ - { - interruptId: 'approval_tool-call-1', - status: 'resolved', - payload: { approved: true }, - }, - ], - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, + it('awaits async Standard Schema approval validation', async () => { + const asyncPayloadSchema = { + '~standard': { + version: 1, + vendor: 'interrupt-test', + validate: async (value: unknown) => { + await Promise.resolve() + if ( + value !== null && + typeof value === 'object' && + 'note' in value && + typeof value.note === 'string' + ) { + return { value: { note: value.note } } + } + return { + issues: [{ message: 'Expected a note.', path: ['note'] }], + } + }, + }, + } satisfies SchemaInput + const tool = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + approvalSchema: { approve: asyncPayloadSchema }, + }) + const approval = normalizeApprovalSchema( + tool.approvalSchema, + tool.inputSchema, ) - - expect(afterApproval.calls).toHaveLength(1) - expect( - ( - afterApproval.calls[0] as { approvals?: ReadonlyMap } - ).approvals?.get('approval_tool-call-1'), - ).toBe(true) - expect( - approvalChunks.find( - (chunk) => - chunk.type === EventType.RUN_FINISHED && - chunk.outcome?.type === 'interrupt', - ), - ).toMatchObject({ - outcome: { - interrupts: [ - { - id: 'client_tool_tool-call-1', - reason: 'client_tool_input', - toolCallId: 'tool-call-1', + const binding: InterruptBinding = { + kind: 'tool-approval', + interruptId: 'approval', + interruptedRunId: 'interrupted-run', + generation: 1, + toolName: tool.name, + toolCallId: 'call-approval', + originalArgs: {}, + inputSchemaHash: hashSchemaInput(tool.inputSchema), + approvalSchemaHash: approval.approvalSchemaHash, + responseSchemaHash: approval.responseSchemaHash, + } + + const result = await validateInterruptResumeBatch({ + threadId: 'thread-1', + interruptedRunId: 'interrupted-run', + generation: 1, + pending: [ + pendingRecord({ + interruptId: 'approval', + binding, + payload: { + id: 'approval', + reason: 'tool_call', + responseSchema: approval.responseSchema, }, - ], - }, + }), + ], + resume: [ + { + interruptId: 'approval', + status: 'resolved', + payload: { approved: true, payload: {} }, + }, + ], + tools: [tool], }) - expect( - (await persistence.stores.interrupts!.get('approval_tool-call-1')) - ?.status, - ).toBe('resolved') - expect( - (await persistence.stores.interrupts!.get('client_tool_tool-call-1')) - ?.status, - ).toBe('pending') - const afterClientTool = mockAdapter([ - toolCallChunks(), - [runStarted(), text('done'), runFinished('r1')], + expect(result.errors).toEqual([ + expect.objectContaining({ + scope: 'item', + interruptId: 'approval', + code: 'invalid-payload', + path: ['note'], + }), + expect.objectContaining({ + scope: 'batch', + code: 'item-validation-failed', + }), ]) - const finalChunks = await collect( + }) + + it('commits one CAS batch with parent and continuation run correlation', async () => { + const persistence = memoryPersistence() + await persistence.stores.interrupts!.openInterruptBatch({ + threadId: 'thread-1', + interruptedRunId: 'interrupted-run', + descriptors: [{ id: 'interrupt-1', reason: 'generic', responseSchema }], + bindings: [ + { + kind: 'generic', + interruptId: 'interrupt-1', + responseSchemaHash: responseSchemaHash(), + }, + ], + }) + const commit = vi.spyOn( + persistence.stores.interrupts!, + 'commitInterruptResolutions', + ) + const { adapter } = mockAdapter([successfulRun('continuation-run')]) + + await collect( chat({ - adapter: afterClientTool.adapter, + adapter, messages: [], - tools: [clientTool('clientSearch')], - runId: 'r1', - threadId: 't1', + runId: 'continuation-run', + parentRunId: 'interrupted-run', + threadId: 'thread-1', resume: [ { - interruptId: 'client_tool_tool-call-1', + interruptId: 'interrupt-1', status: 'resolved', - payload: { answer: 42 }, + payload: { value: 42 }, }, ], middleware: [withChatPersistence(persistence)], }) as AsyncIterable, ) - expect(afterClientTool.calls).toHaveLength(2) - expect(finalChunks).toContainEqual( + expect(commit).toHaveBeenCalledTimes(1) + expect(commit).toHaveBeenCalledWith( expect.objectContaining({ - type: EventType.TOOL_CALL_RESULT, - toolCallId: 'tool-call-1', - content: JSON.stringify({ answer: 42 }), + threadId: 'thread-1', + interruptedRunId: 'interrupted-run', + continuationRunId: 'continuation-run', + expectedGeneration: 1, + expectedInterruptIds: ['interrupt-1'], }), ) - expect(finalChunks).toContainEqual( - expect.objectContaining({ delta: 'done' }), - ) - expect(await persistence.stores.interrupts!.listPending('t1')).toEqual([]) }) +}) - it('rejects invalid resume entries against pending interrupts', async () => { +describe('interrupt recovery', () => { + it('requires authorization and can redact committed resolutions', async () => { const persistence = memoryPersistence() - const first = mockAdapter([[runStarted(), interruptFinished()]]) - await collect( - chat({ - adapter: first.adapter, - messages: [{ role: 'user', content: 'hi' }], - runId: 'r1', - threadId: 't1', - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ) - - const continuation = mockAdapter([[text('SHOULD NOT RUN')]]) - await expect( - collect( - chat({ - adapter: continuation.adapter, - messages: [], - runId: 'r1', - threadId: 't1', - resume: [], - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ), - ).rejects.toThrow(/pending interrupts.*resume is required/i) - - await expect( - collect( - chat({ - adapter: continuation.adapter, - messages: [], - runId: 'r1', - threadId: 't1', - resume: [{ interruptId: 'stale-interrupt', status: 'resolved' }], - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ), - ).rejects.toThrow(/missing resume entry.*interrupt-1/i) - - expect(continuation.calls).toHaveLength(0) - expect(await persistence.stores.interrupts!.listPending('t1')).toHaveLength( - 1, - ) - }) + await persistence.stores.interrupts!.openInterruptBatch({ + threadId: 'thread-1', + interruptedRunId: 'interrupted-run', + descriptors: [{ id: 'interrupt-1', reason: 'generic', responseSchema }], + bindings: [ + { + kind: 'generic', + interruptId: 'interrupt-1', + responseSchemaHash: responseSchemaHash(), + }, + ], + }) + const resolution = { + interruptId: 'interrupt-1', + status: 'resolved' as const, + payload: { value: 42 }, + } + const canonical = canonicalInterruptJson([resolution]) + await persistence.stores.interrupts!.commitInterruptResolutions({ + threadId: 'thread-1', + interruptedRunId: 'interrupted-run', + continuationRunId: 'continuation-run', + expectedGeneration: 1, + expectedInterruptIds: ['interrupt-1'], + resolutions: [resolution], + canonicalResolutions: canonical, + fingerprint: digestInterruptJson(canonical), + }) - it('rejects stale resume entries when a thread has no pending interrupts', async () => { - const persistence = memoryPersistence() - const first = mockAdapter([[runStarted(), text('done'), runFinished()]]) - await collect( - chat({ - adapter: first.adapter, - messages: [{ role: 'user', content: 'hi' }], - runId: 'r1', - threadId: 't1', - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, + const redacted = await getInterruptRecoveryState( + persistence.stores.interrupts!, + { + threadId: 'thread-1', + interruptedRunId: 'interrupted-run', + knownGeneration: 1, + }, ) - expect(await persistence.stores.interrupts!.listPending('t1')).toEqual([]) - - const continuation = mockAdapter([[text('SHOULD NOT RUN')]]) - await expect( - collect( - chat({ - adapter: continuation.adapter, - messages: [], - runId: 'r1', - threadId: 't1', - resume: [{ interruptId: 'stale-interrupt', status: 'resolved' }], - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ), - ).rejects.toThrow(/non-pending interrupt stale-interrupt/i) + expect(redacted.committed).not.toHaveProperty('resolutions') - expect(continuation.calls).toHaveLength(0) - }) - - it('accepts resume only when every pending interrupt has a valid matching entry', async () => { - const persistence = memoryPersistence() - await persistence.stores.interrupts!.create({ - interruptId: 'interrupt-1', - runId: 'old-run', - threadId: 't1', - status: 'pending', - requestedAt: 1, - payload: {}, + const handler = createInterruptRecoveryHandler({ + gateway: persistence.stores.interrupts!, + authorize: async (_request, _input) => ({ + authorized: true, + includeResolutions: false, + }), }) - const bad = mockAdapter([[runStarted(), interruptFinished()]]) - - await expect( - collect( - chat({ - adapter: bad.adapter, - messages: [{ role: 'user', content: 'new input' }], - runId: 'r2', - threadId: 't1', - resume: [{ interruptId: 'different', status: 'resolved' }], - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, + const response = await handler( + new Request( + 'https://example.test/recovery?threadId=thread-1&interruptedRunId=interrupted-run&knownGeneration=1', ), - ).rejects.toThrow(/missing resume entry.*interrupt-1/i) - - const good = mockAdapter([ - [runStarted(), { ...interruptFinished(), runId: 'r2' }], - ]) - await collect( - chat({ - adapter: good.adapter, - messages: [{ role: 'user', content: 'new input' }], - runId: 'r2', - threadId: 't1', - resume: [{ interruptId: 'interrupt-1', status: 'resolved' }], - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, ) - - expect(good.calls).toHaveLength(1) + expect(response.status).toBe(200) + expect(handler).toBeTypeOf('function') + expect(await response.json()).toMatchObject({ + state: 'committed', + committed: { continuationRunId: 'continuation-run' }, + }) }) - it('rejects extra stale resume entries when pending interrupts are satisfied', async () => { + it('validates the recovery query before resource authorization', async () => { const persistence = memoryPersistence() - await persistence.stores.interrupts!.create({ - interruptId: 'interrupt-1', - runId: 'old-run', - threadId: 't1', - status: 'pending', - requestedAt: 1, - payload: {}, + const authorize = vi.fn(() => ({ + authorized: true, + includeResolutions: false, + })) + const handler = createInterruptRecoveryHandler({ + gateway: persistence.stores.interrupts!, + authorize, }) - const run = mockAdapter([[text('SHOULD NOT RUN')]]) - - await expect( - collect( - chat({ - adapter: run.adapter, - messages: [{ role: 'user', content: 'new input' }], - runId: 'r2', - threadId: 't1', - resume: [ - { interruptId: 'interrupt-1', status: 'resolved' }, - { interruptId: 'stale-interrupt', status: 'resolved' }, - ], - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ), - ).rejects.toThrow(/non-pending interrupt stale-interrupt/i) - - expect(run.calls).toHaveLength(0) - expect(await persistence.stores.interrupts!.listPending('t1')).toHaveLength( - 1, + const request = new Request( + 'https://example.test/recovery?threadId=thread-1&knownGeneration=1', ) + + const response = await handler(request) + + expect(response.status).toBe(400) + expect(authorize).not.toHaveBeenCalled() }) - it('applies valid resume entries and allows later normal input', async () => { + it('passes the validated resource identity to authorization', async () => { const persistence = memoryPersistence() - await persistence.stores.interrupts!.create({ - interruptId: 'resolve-me', - runId: 'old-run', - threadId: 't1', - status: 'pending', - requestedAt: 1, - payload: {}, - }) - await persistence.stores.interrupts!.create({ - interruptId: 'cancel-me', - runId: 'old-run', - threadId: 't1', - status: 'pending', - requestedAt: 1, - payload: {}, + const authorize = vi.fn(() => ({ authorized: false }) as const) + const handler = createInterruptRecoveryHandler({ + gateway: persistence.stores.interrupts!, + authorize, }) - - const resumeRun = mockAdapter([ - [runStarted(), text('ok'), runFinished('r2')], - ]) - await collect( - chat({ - adapter: resumeRun.adapter, - messages: [{ role: 'user', content: 'resume' }], - runId: 'r2', - threadId: 't1', - resume: [ - { - interruptId: 'resolve-me', - status: 'resolved', - payload: { approved: true }, - }, - { interruptId: 'cancel-me', status: 'cancelled' }, - ], - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, + const request = new Request( + 'https://example.test/recovery?threadId=thread-1&interruptedRunId=interrupted-run&knownGeneration=2', ) - expect(await persistence.stores.interrupts!.listPending('t1')).toEqual([]) - expect( - (await persistence.stores.interrupts!.get('resolve-me'))?.status, - ).toBe('resolved') - expect( - (await persistence.stores.interrupts!.get('resolve-me'))?.response, - ).toEqual({ approved: true }) - expect( - (await persistence.stores.interrupts!.get('cancel-me'))?.status, - ).toBe('cancelled') + await handler(request) - const nextRun = mockAdapter([ - [runStarted(), text('next'), runFinished('r3')], - ]) - await collect( - chat({ - adapter: nextRun.adapter, - messages: [{ role: 'user', content: 'next' }], - runId: 'r3', - threadId: 't1', - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ) - expect(nextRun.calls).toHaveLength(1) + expect(authorize).toHaveBeenCalledWith(request, { + threadId: 'thread-1', + interruptedRunId: 'interrupted-run', + knownGeneration: 2, + }) }) - it('marks terminal interrupt outcomes as interrupted', async () => { + it('does not query recovery state when authorization is denied', async () => { const persistence = memoryPersistence() - const { adapter } = mockAdapter([[runStarted(), interruptFinished()]]) - - await collect( - chat({ - adapter, - messages: [{ role: 'user', content: 'hi' }], - runId: 'r1', - threadId: 't1', - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ) - - expect((await persistence.stores.runs!.get('r1'))?.status).toBe( - 'interrupted', + const recovery = vi.spyOn( + persistence.stores.interrupts!, + 'getInterruptRecoveryState', ) - expect(await persistence.stores.interrupts!.listPending('t1')).toHaveLength( - 1, + const handler = createInterruptRecoveryHandler({ + gateway: persistence.stores.interrupts!, + authorize: () => ({ authorized: false }), + }) + const response = await handler( + new Request( + 'https://example.test/recovery?threadId=thread-1&interruptedRunId=interrupted-run&knownGeneration=1', + ), ) + expect(response.status).toBe(401) + expect(recovery).not.toHaveBeenCalled() }) }) diff --git a/packages/ai-persistence/tests/memory.conformance.test.ts b/packages/ai-persistence/tests/memory.conformance.test.ts index 0b382b6a0..d2d2a0238 100644 --- a/packages/ai-persistence/tests/memory.conformance.test.ts +++ b/packages/ai-persistence/tests/memory.conformance.test.ts @@ -1,4 +1,57 @@ -import { runPersistenceConformance } from '../src/testkit/conformance' +import { + runInterruptStoreConformance, + runPersistenceConformance, + type InterruptConformanceHarness, +} from '../src/testkit/conformance' import { memoryPersistence } from '../src/memory' +import type { InterruptStore } from '../src/types' runPersistenceConformance('memory', () => memoryPersistence()) + +runInterruptStoreConformance(async (): Promise => { + let now = Date.parse('2026-07-13T10:00:00.000Z') + let failNextTransition = false + const persistence = memoryPersistence({ clock: () => now }) + const base = persistence.stores.interrupts + const store: InterruptStore = { + create: (record) => base.create(record), + resolve: (interruptId, response) => base.resolve(interruptId, response), + cancel: (interruptId) => base.cancel(interruptId), + get: (interruptId) => base.get(interruptId), + list: (threadId) => base.list(threadId), + listPending: (threadId) => base.listPending(threadId), + listByRun: (runId) => base.listByRun(runId), + listPendingByRun: (runId) => base.listPendingByRun(runId), + openInterruptBatch: (input) => base.openInterruptBatch(input), + commitInterruptResolutions: (input) => { + if (failNextTransition) { + failNextTransition = false + return Promise.reject(new Error('Injected transition failure')) + } + return base.commitInterruptResolutions(input) + }, + getInterruptRecoveryState: (input) => base.getInterruptRecoveryState(input), + } + + return { + getStore: () => store, + advanceBy: (milliseconds) => { + now += milliseconds + }, + inspect: async (interruptedRunId) => { + const rows = await base.listByRun(interruptedRunId) + const recovery = await base.getInterruptRecoveryState({ + threadId: rows[0]?.threadId ?? 'thread-cas', + interruptedRunId, + knownGeneration: rows[0]?.generation ?? 1, + }) + return { + statuses: rows.map((row) => row.status).sort(), + batchCount: recovery.state === 'committed' ? 1 : 0, + } + }, + failTransitionOnce: () => { + failNextTransition = true + }, + } +}) diff --git a/packages/ai-persistence/tests/persistence-fixtures.ts b/packages/ai-persistence/tests/persistence-fixtures.ts index 84e2a8e16..d1ff82de6 100644 --- a/packages/ai-persistence/tests/persistence-fixtures.ts +++ b/packages/ai-persistence/tests/persistence-fixtures.ts @@ -46,6 +46,22 @@ export function createRunStore(): RunStore { export function createInterruptStore(): InterruptStore { return { + openInterruptBatch: (input) => + Promise.resolve({ generation: 1, descriptors: input.descriptors }), + commitInterruptResolutions: (input) => + Promise.resolve({ + status: 'committed', + continuationRunId: input.continuationRunId, + }), + getInterruptRecoveryState: (input) => + Promise.resolve({ + schemaVersion: 1, + state: 'missing', + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + generation: input.knownGeneration, + pendingInterrupts: [], + }), create: () => Promise.resolve(), resolve: () => Promise.resolve(), cancel: () => Promise.resolve(), diff --git a/packages/ai-persistence/tests/persistence-validation.test.ts b/packages/ai-persistence/tests/persistence-validation.test.ts index 5f03fc201..1bc961611 100644 --- a/packages/ai-persistence/tests/persistence-validation.test.ts +++ b/packages/ai-persistence/tests/persistence-validation.test.ts @@ -4,6 +4,14 @@ import { withGenerationPersistence, } from '../src/middleware' import type { AIPersistence } from '../src' +import { + InterruptPersistenceCapability, + InterruptsCapability, + getInterruptPersistence, + getInterrupts, + provideInterruptPersistence, + provideInterrupts, +} from '../src' import { createArtifactStore, createBlobStore, @@ -13,6 +21,11 @@ import { } from './persistence-fixtures' describe('persistence store dependency validation', () => { + it('keeps deprecated interrupt capability names as exact core aliases', () => { + expect(InterruptsCapability).toBe(InterruptPersistenceCapability) + expect(getInterrupts).toBe(getInterruptPersistence) + expect(provideInterrupts).toBe(provideInterruptPersistence) + }) it('rejects a dynamic chat persistence with interrupts but no runs', () => { const persistence: AIPersistence = { stores: { interrupts: createInterruptStore() }, @@ -44,6 +57,18 @@ describe('persistence store dependency validation', () => { expect(() => withChatPersistence(persistence)).not.toThrow() }) + it('rejects a sequential interrupt store without the atomic gateway', () => { + const interrupts = createInterruptStore() + Reflect.deleteProperty(interrupts, 'commitInterruptResolutions') + const persistence: AIPersistence = { + stores: { runs: createRunStore(), interrupts }, + } + + expect(() => withChatPersistence(persistence)).toThrow( + /atomic-commit-unsupported/, + ) + }) + it.each([ ['artifacts', { artifacts: createArtifactStore() }], ['blobs', { blobs: createBlobStore() }], diff --git a/packages/ai-persistence/tests/with-persistence.test.ts b/packages/ai-persistence/tests/with-persistence.test.ts index 3065ae8ba..dec163c14 100644 --- a/packages/ai-persistence/tests/with-persistence.test.ts +++ b/packages/ai-persistence/tests/with-persistence.test.ts @@ -1,21 +1,40 @@ import { describe, expect, it } from 'vitest' -import { EventType, chat } from '@tanstack/ai' -import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import { + EventType, + canonicalInterruptJson, + chat, + defineChatMiddleware, + digestInterruptJson, +} from '@tanstack/ai' import { memoryPersistence } from '../src/memory' import { withChatPersistence } from '../src/middleware' import { defineAIPersistence } from '../src/types' +import type { AnyTextAdapter, StreamChunk, Tool } from '@tanstack/ai' // --- minimal mock text adapter --------------------------------------------- function mockAdapter(iterations: Array>) { const calls: Array = [] let i = 0 - const adapter = { + const adapter: AnyTextAdapter = { kind: 'text', name: 'mock', model: 'test-model', - '~types': {}, - chatStream: (opts: unknown) => { + '~types': { + providerOptions: {}, + inputModalities: ['text'], + messageMetadataByModality: { + text: undefined, + image: undefined, + audio: undefined, + video: undefined, + document: undefined, + }, + toolCapabilities: [], + toolCallMetadata: undefined, + systemPromptMetadata: undefined, + }, + chatStream: (opts) => { calls.push(opts) const chunks = iterations[i] ?? [] i++ @@ -24,7 +43,7 @@ function mockAdapter(iterations: Array>) { })() }, structuredOutput: async () => ({ data: {}, rawText: '{}' }), - } as unknown as AnyTextAdapter + } return { adapter, calls } } @@ -52,16 +71,34 @@ const ev = { type: EventType.RUN_FINISHED, runId: 'r1', threadId: 't1', - finishReason: 'stop', + finishReason: 'tool_calls', timestamp: 1, outcome: { type: 'interrupt', interrupts: [ { id: interruptId, - reason: 'approval_required', - toolCallId: 'tool-1', - metadata: { kind: 'approval' }, + reason: 'generic', + responseSchema: { + type: 'object', + properties: { value: { type: 'number' } }, + required: ['value'], + additionalProperties: false, + }, + metadata: { + 'tanstack:interruptBinding': { + kind: 'generic', + interruptId, + responseSchemaHash: digestInterruptJson( + canonicalInterruptJson({ + type: 'object', + properties: { value: { type: 'number' } }, + required: ['value'], + additionalProperties: false, + }), + ), + }, + }, }, ], }, @@ -74,11 +111,17 @@ async function collect(stream: AsyncIterable) { return out } -async function expectCollectRejects( +async function expectCollectInterruptError( stream: AsyncIterable, pattern: RegExp, ) { - await expect(collect(stream)).rejects.toThrow(pattern) + expect(await collect(stream)).toEqual([ + expect.objectContaining({ + type: EventType.RUN_ERROR, + message: expect.stringMatching(pattern), + 'tanstack:interruptErrors': expect.any(Array), + }), + ]) } describe('withChatPersistence (state-only)', () => { @@ -131,6 +174,196 @@ describe('withChatPersistence (state-only)', () => { ) }) + it('emits only safe correlated bindings while persisting descriptors separately', async () => { + const persistence = memoryPersistence() + const approvalTool: Tool = { + name: 'dangerousAction', + description: 'Perform a dangerous action', + needsApproval: true, + inputSchema: { + type: 'object', + properties: { action: { type: 'string' } }, + required: ['action'], + additionalProperties: false, + }, + execute: ({ action }) => ({ action }), + } + const browserTool: Tool = { + name: 'browserAction', + description: 'Perform an action in the browser', + outputSchema: { + type: 'object', + properties: { ok: { type: 'boolean' } }, + required: ['ok'], + additionalProperties: false, + }, + } + const addPrivateBindingFields = defineChatMiddleware({ + name: 'add-private-binding-fields', + onChunk(_ctx, chunk) { + if ( + chunk.type !== EventType.RUN_FINISHED || + chunk.outcome?.type !== 'interrupt' + ) { + return + } + return { + ...chunk, + outcome: { + ...chunk.outcome, + interrupts: chunk.outcome.interrupts.map((interrupt) => ({ + ...interrupt, + metadata: { + ...interrupt.metadata, + 'tanstack:interruptBinding': { + ...interrupt.metadata?.['tanstack:interruptBinding'], + authorizationToken: 'must-not-leak', + internal: { serverOnly: true }, + }, + }, + })), + }, + } + }, + }) + + async function runInterrupt(input: { + runId: string + threadId: string + toolCallId: string + tool: Tool + args: string + }) { + const { adapter } = mockAdapter([ + [ + ev.runStarted(input.runId, input.threadId), + { + type: EventType.TOOL_CALL_START, + toolCallId: input.toolCallId, + toolCallName: input.tool.name, + toolName: input.tool.name, + timestamp: 1, + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: input.toolCallId, + delta: input.args, + timestamp: 1, + }, + { + type: EventType.RUN_FINISHED, + runId: input.runId, + threadId: input.threadId, + finishReason: 'tool_calls', + timestamp: 1, + }, + ], + ]) + const chunks = await collect( + chat({ + adapter, + messages: [{ role: 'user', content: 'pause' }], + runId: input.runId, + threadId: input.threadId, + tools: [input.tool], + middleware: [ + withChatPersistence(persistence), + addPrivateBindingFields, + ], + }) as AsyncIterable, + ) + const terminal = chunks.find( + (chunk) => + chunk.type === EventType.RUN_FINISHED && + chunk.outcome?.type === 'interrupt', + ) + if ( + terminal?.type !== EventType.RUN_FINISHED || + terminal.outcome?.type !== 'interrupt' + ) { + throw new Error('Expected one public interrupt terminal.') + } + return terminal + } + + const approvalTerminal = await runInterrupt({ + runId: 'approval-run', + threadId: 'approval-thread', + toolCallId: 'approval-call-1', + tool: approvalTool, + args: '{"action":"delete"}', + }) + const browserTerminal = await runInterrupt({ + runId: 'client-run', + threadId: 'client-thread', + toolCallId: 'client-call-1', + tool: browserTool, + args: '{}', + }) + expect(approvalTerminal).toMatchObject({ + outcome: { + interrupts: [ + { + id: 'approval_approval-call-1', + reason: 'tool_call', + toolCallId: 'approval-call-1', + }, + ], + }, + }) + expect(browserTerminal).toMatchObject({ + outcome: { + interrupts: [ + { + id: 'client_tool_client-call-1', + reason: 'tanstack:client_tool_execution', + toolCallId: 'client-call-1', + }, + ], + }, + }) + + const stored = [ + ...(await persistence.stores.interrupts.listPendingByRun('approval-run')), + ...(await persistence.stores.interrupts.listPendingByRun('client-run')), + ] + expect(stored).toHaveLength(2) + for (const record of stored) { + if ( + record.payload === null || + typeof record.payload !== 'object' || + Array.isArray(record.payload) + ) { + throw new Error('Expected a persisted interrupt descriptor.') + } + const payload: Record = Object.fromEntries( + Object.entries(record.payload), + ) + expect(payload.metadata).not.toHaveProperty('tanstack:interruptBinding') + expect(record.binding).toMatchObject({ + interruptId: record.interruptId, + interruptedRunId: record.runId, + generation: 1, + }) + expect(record.binding).not.toHaveProperty('authorizationToken') + expect(record.binding).not.toHaveProperty('internal') + } + + const publicBindings = [approvalTerminal, browserTerminal].flatMap( + (terminal) => + terminal.outcome?.type === 'interrupt' + ? terminal.outcome.interrupts.map( + (interrupt) => interrupt.metadata?.['tanstack:interruptBinding'], + ) + : [], + ) + expect(publicBindings).toEqual(stored.map((record) => record.binding)) + for (const binding of publicBindings) { + expect(binding).not.toHaveProperty('authorizationToken') + expect(binding).not.toHaveProperty('internal') + } + }) + it('blocks normal new input while a thread has pending interrupts', async () => { const persistence = memoryPersistence() const first = mockAdapter([[ev.runStarted(), ev.interrupted()]]) @@ -146,15 +379,16 @@ describe('withChatPersistence (state-only)', () => { ) const next = mockAdapter([[ev.text('SHOULD NOT RUN')]]) - await expectCollectRejects( + await expectCollectInterruptError( chat({ adapter: next.adapter, messages: [{ role: 'user', content: 'new input' }], runId: 'r2', threadId: 't1', middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - /pending interrupts.*resume is required/i, + stream: true, + }), + /parentRunId/i, ) expect(next.calls.length).toBe(0) }) @@ -174,16 +408,18 @@ describe('withChatPersistence (state-only)', () => { ) const next = mockAdapter([[ev.text('SHOULD NOT RUN')]]) - await expectCollectRejects( + await expectCollectInterruptError( chat({ adapter: next.adapter, messages: [{ role: 'user', content: 'new input' }], runId: 'r2', + parentRunId: 'r1', threadId: 't1', resume: [{ interruptId: 'other-interrupt', status: 'resolved' }], middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - /missing resume entry for pending interrupt interrupt-1/i, + stream: true, + }), + /missing resume entry for interrupt interrupt-1/i, ) expect(next.calls.length).toBe(0) }) @@ -208,12 +444,13 @@ describe('withChatPersistence (state-only)', () => { adapter: next.adapter, messages: [{ role: 'user', content: 'new input' }], runId: 'r2', + parentRunId: 'r1', threadId: 't1', resume: [ { interruptId: 'interrupt-1', status: 'resolved', - payload: { approved: true }, + payload: { value: 42 }, }, ], middleware: [withChatPersistence(persistence)], diff --git a/packages/ai-preact/src/types.ts b/packages/ai-preact/src/types.ts index 91aee1804..dd59ef413 100644 --- a/packages/ai-preact/src/types.ts +++ b/packages/ai-preact/src/types.ts @@ -6,9 +6,11 @@ import type { } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, + BoundInterrupts, ChatClientOptions, ChatClientState, - ChatPendingInterrupt, + ChatInterrupt, + ChatInterruptState, ChatRequestBody, ChatResumeState, ClientContextOptionFromTools, @@ -105,7 +107,22 @@ export interface UseChatReturn< }) => Promise resumeState: ChatResumeState | null - pendingInterrupts: Array + interrupts: BoundInterrupts + /** @deprecated Use `interrupts`. */ + pendingInterrupts: BoundInterrupts + interruptErrors: ChatInterruptState['interruptErrors'] + resuming: boolean + resolveInterrupts: { + (approved: boolean): void + (resolver: (interrupt: ChatInterrupt) => undefined): void + } + cancelInterrupts: () => void + retryInterrupts: () => void + resumeInterruptsUnsafe: ( + resume: Array, + state?: ChatResumeState, + ) => Promise + /** @deprecated Use bound interrupt methods or `resumeInterruptsUnsafe`. */ resumeInterrupts: ( resume: Array, state?: ChatResumeState, diff --git a/packages/ai-preact/src/use-chat.ts b/packages/ai-preact/src/use-chat.ts index 3f1cc1b08..3341d4277 100644 --- a/packages/ai-preact/src/use-chat.ts +++ b/packages/ai-preact/src/use-chat.ts @@ -10,7 +10,8 @@ import { } from 'preact/hooks' import type { ChatClientState, - ChatPendingInterrupt, + ChatInterrupt, + ChatInterruptState, ChatResumeState, ConnectionStatus, InferredClientContext, @@ -28,6 +29,9 @@ import type { UseChatReturn, } from './types' +const EMPTY_INTERRUPTS = Object.freeze([]) +const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) + export function useChat< const TTools extends ReadonlyArray = any, TContext = InferredClientContext, @@ -48,9 +52,14 @@ export function useChat< const [resumeState, setResumeState] = useState( options.initialResumeSnapshot?.resumeState ?? null, ) - const [pendingInterrupts, setPendingInterrupts] = useState< - Array - >(options.initialResumeSnapshot?.pendingInterrupts ?? []) + const [interruptState, setInterruptState] = useState< + ChatInterruptState + >(() => ({ + interrupts: EMPTY_INTERRUPTS, + pendingInterrupts: EMPTY_INTERRUPTS, + interruptErrors: EMPTY_INTERRUPT_ERRORS, + resuming: false, + })) // Track current messages in a ref to preserve them when client is recreated const messagesRef = useRef>>( @@ -72,7 +81,7 @@ export function useChat< const syncResumeState = useCallback((target: ChatClient | null) => { if (!target) return setResumeState(target.getResumeState()) - setPendingInterrupts(target.getPendingInterrupts()) + setInterruptState(target.getInterruptState()) }, []) useEffect(() => { @@ -92,6 +101,17 @@ export function useChat< ? { connection: initialOptions.connection } : { fetcher: initialOptions.fetcher } + const instanceHolder: { + current: ChatClient | undefined + } = { current: undefined } + const getActiveInstance = () => { + const currentInstance = instanceHolder.current + if (!currentInstance || activeClientRef.current !== currentInstance) { + return undefined + } + return currentInstance + } + const pendingInitializationErrors: Array = [] const instance = new ChatClient({ devtoolsBridgeFactory: createChatDevtoolsBridge, ...transport, @@ -123,23 +143,28 @@ export function useChat< // Capturing the function reference directly would freeze it to whatever // the parent passed on the first render. onResponse: (response) => { - if (activeClientRef.current !== instance) return + if (!getActiveInstance()) return return optionsRef.current.onResponse?.(response) }, onChunk: (chunk) => { - if (activeClientRef.current !== instance) return + if (!getActiveInstance()) return optionsRef.current.onChunk?.(chunk) }, onFinish: (message) => { - if (activeClientRef.current !== instance) return + if (!getActiveInstance()) return optionsRef.current.onFinish?.(message) }, onError: (err) => { - if (activeClientRef.current !== instance) return + const currentInstance = instanceHolder.current + if (!currentInstance) { + pendingInitializationErrors.push(err) + return + } + if (activeClientRef.current !== currentInstance) return optionsRef.current.onError?.(err) }, onCustomEvent: (eventType, data, context) => { - if (activeClientRef.current !== instance) return + if (!getActiveInstance()) return optionsRef.current.onCustomEvent?.(eventType, data, context) }, ...(initialOptions.tools !== undefined && { @@ -149,41 +174,56 @@ export function useChat< streamProcessor: options.streamProcessor, }), onMessagesChange: (newMessages: Array>) => { - if (activeClientRef.current !== instance) return + if (!getActiveInstance()) return setMessages(newMessages) }, onLoadingChange: (newIsLoading: boolean) => { - if (activeClientRef.current !== instance) return + const currentInstance = getActiveInstance() + if (!currentInstance) return setIsLoading(newIsLoading) - syncResumeState(instance) + syncResumeState(currentInstance) }, onStatusChange: (newStatus: ChatClientState) => { - if (activeClientRef.current !== instance) return + if (!getActiveInstance()) return setStatus(newStatus) }, onErrorChange: (newError: Error | undefined) => { - if (activeClientRef.current !== instance) return + if (!getActiveInstance()) return setError(newError) }, onSubscriptionChange: (nextIsSubscribed: boolean) => { - if (activeClientRef.current !== instance) return + if (!getActiveInstance()) return setIsSubscribed(nextIsSubscribed) }, onConnectionStatusChange: (nextStatus: ConnectionStatus) => { - if (activeClientRef.current !== instance) return + if (!getActiveInstance()) return setConnectionStatus(nextStatus) }, onSessionGeneratingChange: (isGenerating: boolean) => { - if (activeClientRef.current !== instance) return + if (!getActiveInstance()) return setSessionGenerating(isGenerating) }, onResumeStateChange: (nextResumeState, nextPendingInterrupts) => { - if (activeClientRef.current !== instance) return + if (!getActiveInstance()) return setResumeState(nextResumeState) - setPendingInterrupts(nextPendingInterrupts) + setInterruptState((current) => ({ + ...current, + interrupts: nextPendingInterrupts, + pendingInterrupts: nextPendingInterrupts, + })) + }, + onInterruptStateChange: (nextInterruptState) => { + if (!getActiveInstance()) return + setInterruptState(nextInterruptState) + optionsRef.current.onInterruptStateChange?.(nextInterruptState) }, }) + instanceHolder.current = instance activeClientRef.current = instance + for (const initializationError of pendingInitializationErrors) { + if (activeClientRef.current !== instance) break + optionsRef.current.onError?.(initializationError) + } return instance }, [clientId, syncResumeState]) @@ -344,6 +384,33 @@ export function useChat< [client, syncResumeState], ) + const resolveInterrupts = useCallback( + ( + resolution: boolean | ((interrupt: ChatInterrupt) => undefined), + ) => { + if (typeof resolution === 'boolean') { + client.resolveInterrupts(resolution) + } else { + client.resolveInterrupts(resolution) + } + }, + [client], + ) + + const cancelInterrupts = useCallback(() => { + client.cancelInterrupts() + }, [client]) + + const retryInterrupts = useCallback(() => { + client.retryInterrupts() + }, [client]) + + const resumeInterruptsUnsafe = useCallback( + (resumeItems: Array, state?: ChatResumeState) => + client.resumeInterruptsUnsafe(resumeItems, state), + [client], + ) + const renderedMessages = client.getMessages() return { @@ -363,7 +430,14 @@ export function useChat< addToolResult, addToolApprovalResponse, resumeState, - pendingInterrupts, + interrupts: interruptState.interrupts, + pendingInterrupts: interruptState.pendingInterrupts, + interruptErrors: interruptState.interruptErrors, + resuming: interruptState.resuming, + resolveInterrupts, + cancelInterrupts, + retryInterrupts, + resumeInterruptsUnsafe, resumeInterrupts, } } diff --git a/packages/ai-preact/tests/test-utils.ts b/packages/ai-preact/tests/test-utils.ts index ac5d85676..ad6f25b6b 100644 --- a/packages/ai-preact/tests/test-utils.ts +++ b/packages/ai-preact/tests/test-utils.ts @@ -4,6 +4,81 @@ import { useChat } from '../src/use-chat' import type { RenderHookResult } from '@testing-library/preact' import type { UseChatOptions, UseChatReturn } from '../src/types' +import type { ChatResumeSnapshotV2 } from '@tanstack/ai-client' + +export function createInterruptResumeSnapshot(): ChatResumeSnapshotV2 { + const pendingInterrupts = [ + { + id: 'staged-interrupt', + reason: 'confirmation', + metadata: { + 'tanstack:interruptBinding': { + kind: 'generic' as const, + interruptId: 'staged-interrupt', + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'none', + }, + }, + }, + { + id: 'invalid-interrupt', + reason: 'confirmation', + metadata: { + 'tanstack:interruptBinding': { + kind: 'generic' as const, + interruptId: 'invalid-interrupt', + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'none', + }, + }, + }, + ] + return { + schemaVersion: 2, + resumeState: { threadId: 'thread-1', runId: 'run-1' }, + pendingInterrupts, + interruptState: { + recoveryState: { + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + pendingInterrupts, + }, + drafts: [ + { + interruptId: 'staged-interrupt', + response: { + interruptId: 'staged-interrupt', + status: 'resolved', + payload: { answer: 42 }, + }, + status: 'staged', + }, + { + interruptId: 'invalid-interrupt', + response: null, + status: 'error', + error: { + scope: 'item', + interruptId: 'invalid-interrupt', + code: 'invalid-payload', + message: 'Invalid persisted response', + source: 'client', + retryable: false, + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + }, + }, + ], + }, + } +} + export { createMockConnectionAdapter, createTextChunks, diff --git a/packages/ai-preact/tests/use-chat-types.test.ts b/packages/ai-preact/tests/use-chat-types.test.ts index 5880305c6..b9367bcd2 100644 --- a/packages/ai-preact/tests/use-chat-types.test.ts +++ b/packages/ai-preact/tests/use-chat-types.test.ts @@ -7,7 +7,16 @@ import { describe, expectTypeOf, it } from 'vitest' import { toolDefinition } from '@tanstack/ai' import { clientTools } from '@tanstack/ai-client' import { useChat } from '../src/use-chat' -import type { UseChatOptions } from '../src/types' +import type { UseChatOptions, UseChatReturn } from '../src/types' + +type TestSchema = { + readonly '~standard': { + readonly version: 1 + readonly vendor: 'test' + readonly types: { readonly input: T; readonly output: T } + readonly validate: (value: unknown) => { readonly value: T } + } +} describe('useChat() return type (preact)', () => { describe('with typed client tool context', () => { @@ -70,3 +79,133 @@ describe('useChat() return type (preact)', () => { }) }) }) + +describe('useChat() interrupt types', () => { + it('preserves approval, generic, and client-tool inference', () => { + const inputSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { cents: 0 }, + output: { cents: 0 }, + }, + validate: (value: unknown) => ({ + value: + value !== null && + typeof value === 'object' && + 'cents' in value && + typeof value.cents === 'number' + ? { cents: value.cents } + : { cents: 0 }, + }), + }, + } + const approveSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { note: '' }, + output: { note: '' }, + }, + validate: () => ({ value: { note: '' } }), + }, + } + const rejectSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { reason: '' }, + output: { reason: '' }, + }, + validate: () => ({ value: { reason: '' } }), + }, + } + const outputSchema: TestSchema<{ accountId: string }> = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { accountId: '' }, + output: { accountId: '' }, + }, + validate: () => ({ value: { accountId: '' } }), + }, + } + const transfer = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema, + approvalSchema: { + approve: approveSchema, + reject: rejectSchema, + }, + }).client() + const confirm = toolDefinition({ + name: 'confirm', + description: 'Confirm without schemas', + needsApproval: true, + }).client() + const lookup = toolDefinition({ + name: 'lookup', + description: 'Lookup account', + outputSchema, + }).client() + const tools = clientTools(transfer, confirm, lookup) + type Interrupt = UseChatReturn['interrupts'][number] + type Transfer = Extract< + Interrupt, + { kind: 'tool-approval'; toolName: 'transfer' } + > + type Confirm = Extract< + Interrupt, + { kind: 'tool-approval'; toolName: 'confirm' } + > + type Lookup = Extract< + Interrupt, + { kind: 'client-tool-execution'; toolName: 'lookup' } + > + type Generic = Extract + + const check = ( + transferInterrupt: Transfer, + confirmInterrupt: Confirm, + lookupInterrupt: Lookup, + genericInterrupt: Generic, + ) => { + transferInterrupt.resolveInterrupt(true, { + editedArgs: { cents: 100 }, + payload: { note: 'approved' }, + }) + transferInterrupt.resolveInterrupt(false, { + payload: { reason: 'declined' }, + }) + // @ts-expect-error rejected approvals cannot edit tool input + transferInterrupt.resolveInterrupt(false, { editedArgs: { cents: 1 } }) + // @ts-expect-error approve payload uses the approve branch + transferInterrupt.resolveInterrupt(true, { + payload: { reason: 'wrong branch' }, + }) + + confirmInterrupt.resolveInterrupt(true) + confirmInterrupt.resolveInterrupt(false) + // @ts-expect-error omitted input schema forbids edited input + confirmInterrupt.resolveInterrupt(true, { editedArgs: { cents: 1 } }) + // @ts-expect-error omitted approval branches forbid payloads + confirmInterrupt.resolveInterrupt(false, { + payload: { reason: 'no branch' }, + }) + + lookupInterrupt.resolveInterrupt({ accountId: 'account-1' }) + // @ts-expect-error client-tool output is inferred + lookupInterrupt.resolveInterrupt({ account: 'account-1' }) + expectTypeOf(genericInterrupt.resolveInterrupt) + .parameter(0) + .toEqualTypeOf() + } + void check + }) +}) diff --git a/packages/ai-preact/tests/use-chat.test.ts b/packages/ai-preact/tests/use-chat.test.ts index f932b1ef2..9d7f4e02f 100644 --- a/packages/ai-preact/tests/use-chat.test.ts +++ b/packages/ai-preact/tests/use-chat.test.ts @@ -1,6 +1,9 @@ import type { ModelMessage, StreamChunk } from '@tanstack/ai' import { EventType } from '@tanstack/ai' -import type { SubscribeConnectionAdapter } from '@tanstack/ai-client' +import { + ChatClient, + type SubscribeConnectionAdapter, +} from '@tanstack/ai-client' import { act, renderHook, waitFor } from '@testing-library/preact' import { StrictMode } from 'preact/compat' import { useState } from 'preact/hooks' @@ -9,6 +12,7 @@ import type { UIMessage } from '../src/types' import { useChat } from '../src/use-chat' import { createMockConnectionAdapter, + createInterruptResumeSnapshot, createTextChunks, createToolCallChunks, renderUseChat, @@ -27,6 +31,85 @@ describe('useChat', () => { return { promise, resolve } } + describe('interrupt state', () => { + it('projects one immutable snapshot with the deprecated pending alias', async () => { + const onInterruptStateChange = vi.fn() + const { result } = renderUseChat({ + connection: createMockConnectionAdapter(), + initialResumeSnapshot: createInterruptResumeSnapshot(), + onInterruptStateChange, + }) + + expect(Object.isFrozen(result.current.interrupts)).toBe(true) + expect(result.current.pendingInterrupts).toBe(result.current.interrupts) + expect(result.current.interrupts[0]).toMatchObject({ + id: 'staged-interrupt', + status: 'staged', + }) + expect(result.current.interrupts[1]).toMatchObject({ + id: 'invalid-interrupt', + status: 'error', + error: { code: 'invalid-payload' }, + }) + expect(result.current.interruptErrors).toEqual([]) + expect(result.current.resuming).toBe(false) + expect(result.current.interrupts[0]).toEqual( + expect.objectContaining({ + resolveInterrupt: expect.any(Function), + cancel: expect.any(Function), + clearResolution: expect.any(Function), + }), + ) + + act(() => result.current.resolveInterrupts(false)) + await waitFor(() => { + expect(result.current.interruptErrors[0]?.code).toBe( + 'unsupported-bulk-operation', + ) + }) + expect(onInterruptStateChange).toHaveBeenLastCalledWith( + expect.objectContaining({ + interrupts: result.current.interrupts, + interruptErrors: result.current.interruptErrors, + }), + ) + }) + + it('delegates every root interrupt control to ChatClient', async () => { + const resolve = vi + .spyOn(ChatClient.prototype, 'resolveInterrupts') + .mockImplementation(() => {}) + const cancel = vi + .spyOn(ChatClient.prototype, 'cancelInterrupts') + .mockImplementation(() => {}) + const retry = vi + .spyOn(ChatClient.prototype, 'retryInterrupts') + .mockImplementation(() => {}) + const unsafe = vi + .spyOn(ChatClient.prototype, 'resumeInterruptsUnsafe') + .mockResolvedValue(true) + const { result } = renderUseChat({ + connection: createMockConnectionAdapter(), + }) + const resolver = () => undefined + const resume = [{ interruptId: 'one', status: 'cancelled' as const }] + + act(() => { + result.current.resolveInterrupts(resolver) + result.current.cancelInterrupts() + result.current.retryInterrupts() + }) + await expect(result.current.resumeInterruptsUnsafe(resume)).resolves.toBe( + true, + ) + + expect(resolve).toHaveBeenCalledWith(resolver) + expect(cancel).toHaveBeenCalledOnce() + expect(retry).toHaveBeenCalledOnce() + expect(unsafe).toHaveBeenCalledWith(resume, undefined) + }) + }) + describe('initialization', () => { it('should initialize with default state', () => { const adapter = createMockConnectionAdapter() diff --git a/packages/ai-react/src/types.ts b/packages/ai-react/src/types.ts index cf38b8a13..7f696c2f6 100644 --- a/packages/ai-react/src/types.ts +++ b/packages/ai-react/src/types.ts @@ -7,9 +7,11 @@ import type { } from '@tanstack/ai/client' import type { AIDevtoolsDisplayOptions, + BoundInterrupts, ChatClientOptions, ChatClientState, - ChatPendingInterrupt, + ChatInterrupt, + ChatInterruptState, ChatRequestBody, ChatResumeState, ClientContextOptionFromTools, @@ -166,7 +168,22 @@ interface BaseUseChatReturn< }) => Promise resumeState: ChatResumeState | null - pendingInterrupts: Array + interrupts: BoundInterrupts + /** @deprecated Use `interrupts`. */ + pendingInterrupts: BoundInterrupts + interruptErrors: ChatInterruptState['interruptErrors'] + resuming: boolean + resolveInterrupts: { + (approved: boolean): void + (resolver: (interrupt: ChatInterrupt) => undefined): void + } + cancelInterrupts: () => void + retryInterrupts: () => void + resumeInterruptsUnsafe: ( + resume: Array, + state?: ChatResumeState, + ) => Promise + /** @deprecated Use bound interrupt methods or `resumeInterruptsUnsafe`. */ resumeInterrupts: ( resume: Array, state?: ChatResumeState, diff --git a/packages/ai-react/src/use-chat.ts b/packages/ai-react/src/use-chat.ts index 30c9ff202..fb65ea808 100644 --- a/packages/ai-react/src/use-chat.ts +++ b/packages/ai-react/src/use-chat.ts @@ -11,7 +11,8 @@ import type { } from '@tanstack/ai/client' import type { ChatClientState, - ChatPendingInterrupt, + ChatInterrupt, + ChatInterruptState, ChatResumeState, ConnectionStatus, InferredClientContext, @@ -26,6 +27,9 @@ import type { UseChatReturn, } from './types' +const EMPTY_INTERRUPTS = Object.freeze([]) +const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) + export function useChat< const TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, @@ -49,9 +53,14 @@ export function useChat< const [resumeState, setResumeState] = useState( options.initialResumeSnapshot?.resumeState ?? null, ) - const [pendingInterrupts, setPendingInterrupts] = useState< - Array - >(options.initialResumeSnapshot?.pendingInterrupts ?? []) + const [interruptState, setInterruptState] = useState< + ChatInterruptState + >(() => ({ + interrupts: EMPTY_INTERRUPTS, + pendingInterrupts: EMPTY_INTERRUPTS, + interruptErrors: EMPTY_INTERRUPT_ERRORS, + resuming: false, + })) type Partial = DeepPartial>> type Final = InferSchemaType> @@ -80,7 +89,7 @@ export function useChat< const syncResumeState = useCallback((target: ChatClient | null) => { if (!target) return setResumeState(target.getResumeState()) - setPendingInterrupts(target.getPendingInterrupts()) + setInterruptState(target.getInterruptState()) }, []) // Create ChatClient instance with callbacks to sync state @@ -199,7 +208,16 @@ export function useChat< onResumeStateChange: (nextResumeState, nextPendingInterrupts) => { if (!getActiveInstance()) return setResumeState(nextResumeState) - setPendingInterrupts(nextPendingInterrupts) + setInterruptState((current) => ({ + ...current, + interrupts: nextPendingInterrupts, + pendingInterrupts: nextPendingInterrupts, + })) + }, + onInterruptStateChange: (nextInterruptState) => { + if (!getActiveInstance()) return + setInterruptState(nextInterruptState) + optionsRef.current.onInterruptStateChange?.(nextInterruptState) }, }) instanceHolder.current = instance @@ -374,6 +392,33 @@ export function useChat< [client, syncResumeState], ) + const resolveInterrupts = useCallback( + ( + resolution: boolean | ((interrupt: ChatInterrupt) => undefined), + ) => { + if (typeof resolution === 'boolean') { + client.resolveInterrupts(resolution) + } else { + client.resolveInterrupts(resolution) + } + }, + [client], + ) + + const cancelInterrupts = useCallback(() => { + client.cancelInterrupts() + }, [client]) + + const retryInterrupts = useCallback(() => { + client.retryInterrupts() + }, [client]) + + const resumeInterruptsUnsafe = useCallback( + (resumeItems: Array, state?: ChatResumeState) => + client.resumeInterruptsUnsafe(resumeItems, state), + [client], + ) + // The "active" structured-output part is the one on the assistant message // that follows the latest user message. No such message exists between // sendMessage() and the first chunk, so partial/final naturally read as @@ -440,7 +485,14 @@ export function useChat< addToolResult, addToolApprovalResponse, resumeState, - pendingInterrupts, + interrupts: interruptState.interrupts, + pendingInterrupts: interruptState.pendingInterrupts, + interruptErrors: interruptState.interruptErrors, + resuming: interruptState.resuming, + resolveInterrupts, + cancelInterrupts, + retryInterrupts, + resumeInterruptsUnsafe, resumeInterrupts, partial, final, diff --git a/packages/ai-react/tests/test-utils.ts b/packages/ai-react/tests/test-utils.ts index 40f0b2ac9..e9e7a6c36 100644 --- a/packages/ai-react/tests/test-utils.ts +++ b/packages/ai-react/tests/test-utils.ts @@ -8,6 +8,80 @@ export { import { renderHook, type RenderHookResult } from '@testing-library/react' import type { UseChatOptions, UseChatReturn } from '../src/types' import { useChat } from '../src/use-chat' +import type { ChatResumeSnapshotV2 } from '@tanstack/ai-client' + +export function createInterruptResumeSnapshot(): ChatResumeSnapshotV2 { + const pendingInterrupts = [ + { + id: 'staged-interrupt', + reason: 'confirmation', + metadata: { + 'tanstack:interruptBinding': { + kind: 'generic' as const, + interruptId: 'staged-interrupt', + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'none', + }, + }, + }, + { + id: 'invalid-interrupt', + reason: 'confirmation', + metadata: { + 'tanstack:interruptBinding': { + kind: 'generic' as const, + interruptId: 'invalid-interrupt', + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'none', + }, + }, + }, + ] + return { + schemaVersion: 2, + resumeState: { threadId: 'thread-1', runId: 'run-1' }, + pendingInterrupts, + interruptState: { + recoveryState: { + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + pendingInterrupts, + }, + drafts: [ + { + interruptId: 'staged-interrupt', + response: { + interruptId: 'staged-interrupt', + status: 'resolved', + payload: { answer: 42 }, + }, + status: 'staged', + }, + { + interruptId: 'invalid-interrupt', + response: null, + status: 'error', + error: { + scope: 'item', + interruptId: 'invalid-interrupt', + code: 'invalid-payload', + message: 'Invalid persisted response', + source: 'client', + retryable: false, + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + }, + }, + ], + }, + } +} /** * Render the useChat hook with testing utilities diff --git a/packages/ai-react/tests/use-chat-types.test.ts b/packages/ai-react/tests/use-chat-types.test.ts index 4edd539fb..69ccc67b1 100644 --- a/packages/ai-react/tests/use-chat-types.test.ts +++ b/packages/ai-react/tests/use-chat-types.test.ts @@ -16,6 +16,14 @@ import type { DeepPartial, UseChatOptions, UseChatReturn } from '../src/types' type Person = { name: string; age: number; email: string } type PersonSchema = StandardJSONSchemaV1 type NoTools = ReadonlyArray +type TestSchema = { + readonly '~standard': { + readonly version: 1 + readonly vendor: 'test' + readonly types: { readonly input: T; readonly output: T } + readonly validate: (value: unknown) => { readonly value: T } + } +} describe('useChat() return type', () => { describe('with outputSchema', () => { @@ -183,3 +191,133 @@ describe('useChat() return type', () => { }) }) }) + +describe('useChat() interrupt types', () => { + it('preserves approval, generic, and client-tool inference', () => { + const inputSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { cents: 0 }, + output: { cents: 0 }, + }, + validate: (value: unknown) => ({ + value: + value !== null && + typeof value === 'object' && + 'cents' in value && + typeof value.cents === 'number' + ? { cents: value.cents } + : { cents: 0 }, + }), + }, + } + const approveSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { note: '' }, + output: { note: '' }, + }, + validate: () => ({ value: { note: '' } }), + }, + } + const rejectSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { reason: '' }, + output: { reason: '' }, + }, + validate: () => ({ value: { reason: '' } }), + }, + } + const outputSchema: TestSchema<{ accountId: string }> = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { accountId: '' }, + output: { accountId: '' }, + }, + validate: () => ({ value: { accountId: '' } }), + }, + } + const transfer = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema, + approvalSchema: { + approve: approveSchema, + reject: rejectSchema, + }, + }).client() + const confirm = toolDefinition({ + name: 'confirm', + description: 'Confirm without schemas', + needsApproval: true, + }).client() + const lookup = toolDefinition({ + name: 'lookup', + description: 'Lookup account', + outputSchema, + }).client() + const tools = clientTools(transfer, confirm, lookup) + type Interrupt = UseChatReturn['interrupts'][number] + type Transfer = Extract< + Interrupt, + { kind: 'tool-approval'; toolName: 'transfer' } + > + type Confirm = Extract< + Interrupt, + { kind: 'tool-approval'; toolName: 'confirm' } + > + type Lookup = Extract< + Interrupt, + { kind: 'client-tool-execution'; toolName: 'lookup' } + > + type Generic = Extract + + const check = ( + transferInterrupt: Transfer, + confirmInterrupt: Confirm, + lookupInterrupt: Lookup, + genericInterrupt: Generic, + ) => { + transferInterrupt.resolveInterrupt(true, { + editedArgs: { cents: 100 }, + payload: { note: 'approved' }, + }) + transferInterrupt.resolveInterrupt(false, { + payload: { reason: 'declined' }, + }) + // @ts-expect-error rejected approvals cannot edit tool input + transferInterrupt.resolveInterrupt(false, { editedArgs: { cents: 1 } }) + // @ts-expect-error approve payload uses the approve branch + transferInterrupt.resolveInterrupt(true, { + payload: { reason: 'wrong branch' }, + }) + + confirmInterrupt.resolveInterrupt(true) + confirmInterrupt.resolveInterrupt(false) + // @ts-expect-error omitted input schema forbids edited input + confirmInterrupt.resolveInterrupt(true, { editedArgs: { cents: 1 } }) + // @ts-expect-error omitted approval branches forbid payloads + confirmInterrupt.resolveInterrupt(false, { + payload: { reason: 'no branch' }, + }) + + lookupInterrupt.resolveInterrupt({ accountId: 'account-1' }) + // @ts-expect-error client-tool output is inferred + lookupInterrupt.resolveInterrupt({ account: 'account-1' }) + expectTypeOf(genericInterrupt.resolveInterrupt) + .parameter(0) + .toEqualTypeOf() + } + void check + }) +}) diff --git a/packages/ai-react/tests/use-chat.test.ts b/packages/ai-react/tests/use-chat.test.ts index 7511c58fe..7e56780cc 100644 --- a/packages/ai-react/tests/use-chat.test.ts +++ b/packages/ai-react/tests/use-chat.test.ts @@ -1,6 +1,9 @@ import type { ModelMessage, StreamChunk } from '@tanstack/ai' import { EventType } from '@tanstack/ai' -import type { SubscribeConnectionAdapter } from '@tanstack/ai-client' +import { + ChatClient, + type SubscribeConnectionAdapter, +} from '@tanstack/ai-client' import { act, renderHook, waitFor } from '@testing-library/react' import { StrictMode, useState } from 'react' import { afterEach, describe, expect, it, vi } from 'vitest' @@ -8,6 +11,7 @@ import type { UIMessage, UseChatOptions } from '../src/types' import { useChat } from '../src/use-chat' import { createMockConnectionAdapter, + createInterruptResumeSnapshot, createTextChunks, createToolCallChunks, renderUseChat, @@ -26,6 +30,85 @@ describe('useChat', () => { return { promise, resolve } } + describe('interrupt state', () => { + it('projects one immutable snapshot with the deprecated pending alias', async () => { + const onInterruptStateChange = vi.fn() + const { result } = renderUseChat({ + connection: createMockConnectionAdapter(), + initialResumeSnapshot: createInterruptResumeSnapshot(), + onInterruptStateChange, + }) + + expect(Object.isFrozen(result.current.interrupts)).toBe(true) + expect(result.current.pendingInterrupts).toBe(result.current.interrupts) + expect(result.current.interrupts[0]).toMatchObject({ + id: 'staged-interrupt', + status: 'staged', + }) + expect(result.current.interrupts[1]).toMatchObject({ + id: 'invalid-interrupt', + status: 'error', + error: { code: 'invalid-payload' }, + }) + expect(result.current.interruptErrors).toEqual([]) + expect(result.current.resuming).toBe(false) + expect(result.current.interrupts[0]).toEqual( + expect.objectContaining({ + resolveInterrupt: expect.any(Function), + cancel: expect.any(Function), + clearResolution: expect.any(Function), + }), + ) + + act(() => result.current.resolveInterrupts(false)) + await waitFor(() => { + expect(result.current.interruptErrors[0]?.code).toBe( + 'unsupported-bulk-operation', + ) + }) + expect(onInterruptStateChange).toHaveBeenLastCalledWith( + expect.objectContaining({ + interrupts: result.current.interrupts, + interruptErrors: result.current.interruptErrors, + }), + ) + }) + + it('delegates every root interrupt control to ChatClient', async () => { + const resolve = vi + .spyOn(ChatClient.prototype, 'resolveInterrupts') + .mockImplementation(() => {}) + const cancel = vi + .spyOn(ChatClient.prototype, 'cancelInterrupts') + .mockImplementation(() => {}) + const retry = vi + .spyOn(ChatClient.prototype, 'retryInterrupts') + .mockImplementation(() => {}) + const unsafe = vi + .spyOn(ChatClient.prototype, 'resumeInterruptsUnsafe') + .mockResolvedValue(true) + const { result } = renderUseChat({ + connection: createMockConnectionAdapter(), + }) + const resolver = () => undefined + const resume = [{ interruptId: 'one', status: 'cancelled' as const }] + + act(() => { + result.current.resolveInterrupts(resolver) + result.current.cancelInterrupts() + result.current.retryInterrupts() + }) + await expect(result.current.resumeInterruptsUnsafe(resume)).resolves.toBe( + true, + ) + + expect(resolve).toHaveBeenCalledWith(resolver) + expect(cancel).toHaveBeenCalledOnce() + expect(retry).toHaveBeenCalledOnce() + expect(unsafe).toHaveBeenCalledWith(resume, undefined) + }) + }) + describe('initialization', () => { it('should initialize with default state', () => { const adapter = createMockConnectionAdapter() diff --git a/packages/ai-solid/src/types.ts b/packages/ai-solid/src/types.ts index 580d0be1c..d1cf04aff 100644 --- a/packages/ai-solid/src/types.ts +++ b/packages/ai-solid/src/types.ts @@ -7,9 +7,11 @@ import type { } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, + BoundInterrupts, ChatClientOptions, ChatClientState, - ChatPendingInterrupt, + ChatInterrupt, + ChatInterruptState, ChatRequestBody, ChatResumeState, ClientContextOptionFromTools, @@ -152,7 +154,22 @@ interface BaseUseChatReturn< }) => Promise resumeState: Accessor - pendingInterrupts: Accessor> + interrupts: Accessor> + /** @deprecated Use `interrupts`. */ + pendingInterrupts: Accessor> + interruptErrors: Accessor['interruptErrors']> + resuming: Accessor + resolveInterrupts: { + (approved: boolean): void + (resolver: (interrupt: ChatInterrupt) => undefined): void + } + cancelInterrupts: () => void + retryInterrupts: () => void + resumeInterruptsUnsafe: ( + resume: Array, + state?: ChatResumeState, + ) => Promise + /** @deprecated Use bound interrupt methods or `resumeInterruptsUnsafe`. */ resumeInterrupts: ( resume: Array, state?: ChatResumeState, diff --git a/packages/ai-solid/src/use-chat.ts b/packages/ai-solid/src/use-chat.ts index 327058d37..f2aed2ce4 100644 --- a/packages/ai-solid/src/use-chat.ts +++ b/packages/ai-solid/src/use-chat.ts @@ -11,7 +11,8 @@ import { ChatClient } from '@tanstack/ai-client' import { createChatDevtoolsBridge } from '@tanstack/ai-client/devtools' import type { ChatClientState, - ChatPendingInterrupt, + ChatInterrupt, + ChatInterruptState, ChatResumeState, ConnectionStatus, InferredClientContext, @@ -33,6 +34,9 @@ import type { UseChatReturn, } from './types' +const EMPTY_INTERRUPTS = Object.freeze([]) +const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) + export function useChat< const TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, @@ -60,13 +64,18 @@ export function useChat< const [resumeState, setResumeState] = createSignal( options.initialResumeSnapshot?.resumeState ?? null, ) - const [pendingInterrupts, setPendingInterrupts] = createSignal< - Array - >(options.initialResumeSnapshot?.pendingInterrupts ?? []) + const [interruptState, setInterruptState] = createSignal< + ChatInterruptState + >({ + interrupts: EMPTY_INTERRUPTS, + pendingInterrupts: EMPTY_INTERRUPTS, + interruptErrors: EMPTY_INTERRUPT_ERRORS, + resuming: false, + }) const syncResumeState = () => { setResumeState(client().getResumeState()) - setPendingInterrupts(client().getPendingInterrupts()) + setInterruptState(client().getInterruptState()) } // Structured-output `partial` / `final` are derived from `messages` — @@ -155,7 +164,15 @@ export function useChat< }, onResumeStateChange: (nextResumeState, nextPendingInterrupts) => { setResumeState(nextResumeState) - setPendingInterrupts(nextPendingInterrupts) + setInterruptState((current) => ({ + ...current, + interrupts: nextPendingInterrupts, + pendingInterrupts: nextPendingInterrupts, + })) + }, + onInterruptStateChange: (nextInterruptState) => { + setInterruptState(nextInterruptState) + options.onInterruptStateChange?.(nextInterruptState) }, }) // Only recreate when clientId changes @@ -163,6 +180,7 @@ export function useChat< }, [clientId]) setMessages(client().getMessages()) + syncResumeState() // Sync body / forwardedProps changes to the client. // Both populate the same wire payload; `forwardedProps` is preferred @@ -279,6 +297,34 @@ export function useChat< return result } + const resolveInterrupts = ( + resolution: boolean | ((interrupt: ChatInterrupt) => undefined), + ) => { + if (typeof resolution === 'boolean') { + client().resolveInterrupts(resolution) + } else { + client().resolveInterrupts(resolution) + } + } + + const cancelInterrupts = () => { + client().cancelInterrupts() + } + + const retryInterrupts = () => { + client().retryInterrupts() + } + + const resumeInterruptsUnsafe = ( + resumeItems: Array, + state?: ChatResumeState, + ) => client().resumeInterruptsUnsafe(resumeItems, state) + + const interrupts = () => interruptState().interrupts + const pendingInterrupts = () => interruptState().pendingInterrupts + const interruptErrors = () => interruptState().interruptErrors + const resuming = () => interruptState().resuming + // The "active" structured-output part is on the assistant message after // the latest user message. When no user message exists yet, return null // rather than scanning history — otherwise a stale `final` from @@ -337,7 +383,14 @@ export function useChat< addToolResult, addToolApprovalResponse, resumeState, + interrupts, pendingInterrupts, + interruptErrors, + resuming, + resolveInterrupts, + cancelInterrupts, + retryInterrupts, + resumeInterruptsUnsafe, resumeInterrupts, partial, final, diff --git a/packages/ai-solid/tests/test-utils.ts b/packages/ai-solid/tests/test-utils.ts index 82bff2ad7..b45cce38f 100644 --- a/packages/ai-solid/tests/test-utils.ts +++ b/packages/ai-solid/tests/test-utils.ts @@ -10,6 +10,81 @@ import { renderHook } from '@solidjs/testing-library' import type { UseChatOptions } from '../src/types' import { useChat } from '../src/use-chat' +import type { ChatResumeSnapshotV2 } from '@tanstack/ai-client' + +export function createInterruptResumeSnapshot(): ChatResumeSnapshotV2 { + const pendingInterrupts = [ + { + id: 'staged-interrupt', + reason: 'confirmation', + metadata: { + 'tanstack:interruptBinding': { + kind: 'generic' as const, + interruptId: 'staged-interrupt', + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'none', + }, + }, + }, + { + id: 'invalid-interrupt', + reason: 'confirmation', + metadata: { + 'tanstack:interruptBinding': { + kind: 'generic' as const, + interruptId: 'invalid-interrupt', + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'none', + }, + }, + }, + ] + return { + schemaVersion: 2, + resumeState: { threadId: 'thread-1', runId: 'run-1' }, + pendingInterrupts, + interruptState: { + recoveryState: { + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + pendingInterrupts, + }, + drafts: [ + { + interruptId: 'staged-interrupt', + response: { + interruptId: 'staged-interrupt', + status: 'resolved', + payload: { answer: 42 }, + }, + status: 'staged', + }, + { + interruptId: 'invalid-interrupt', + response: null, + status: 'error', + error: { + scope: 'item', + interruptId: 'invalid-interrupt', + code: 'invalid-payload', + message: 'Invalid persisted response', + source: 'client', + retryable: false, + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + }, + }, + ], + }, + } +} + /** * Render the useChat hook with testing utilities * diff --git a/packages/ai-solid/tests/use-chat-types.test.ts b/packages/ai-solid/tests/use-chat-types.test.ts index aa2e7c878..a300c269e 100644 --- a/packages/ai-solid/tests/use-chat-types.test.ts +++ b/packages/ai-solid/tests/use-chat-types.test.ts @@ -14,6 +14,14 @@ import type { DeepPartial, UseChatOptions, UseChatReturn } from '../src/types' type Person = { name: string; age: number; email: string } type PersonSchema = StandardJSONSchemaV1 type NoTools = ReadonlyArray +type TestSchema = { + readonly '~standard': { + readonly version: 1 + readonly vendor: 'test' + readonly types: { readonly input: T; readonly output: T } + readonly validate: (value: unknown) => { readonly value: T } + } +} describe('useChat() return type (solid)', () => { describe('with outputSchema', () => { @@ -137,3 +145,135 @@ describe('useChat() return type (solid)', () => { }) }) }) + +describe('useChat() interrupt types', () => { + it('preserves approval, generic, and client-tool inference', () => { + const inputSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { cents: 0 }, + output: { cents: 0 }, + }, + validate: (value: unknown) => ({ + value: + value !== null && + typeof value === 'object' && + 'cents' in value && + typeof value.cents === 'number' + ? { cents: value.cents } + : { cents: 0 }, + }), + }, + } + const approveSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { note: '' }, + output: { note: '' }, + }, + validate: () => ({ value: { note: '' } }), + }, + } + const rejectSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { reason: '' }, + output: { reason: '' }, + }, + validate: () => ({ value: { reason: '' } }), + }, + } + const outputSchema: TestSchema<{ accountId: string }> = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { accountId: '' }, + output: { accountId: '' }, + }, + validate: () => ({ value: { accountId: '' } }), + }, + } + const transfer = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema, + approvalSchema: { + approve: approveSchema, + reject: rejectSchema, + }, + }).client() + const confirm = toolDefinition({ + name: 'confirm', + description: 'Confirm without schemas', + needsApproval: true, + }).client() + const lookup = toolDefinition({ + name: 'lookup', + description: 'Lookup account', + outputSchema, + }).client() + const tools = clientTools(transfer, confirm, lookup) + type Interrupt = ReturnType< + UseChatReturn['interrupts'] + >[number] + type Transfer = Extract< + Interrupt, + { kind: 'tool-approval'; toolName: 'transfer' } + > + type Confirm = Extract< + Interrupt, + { kind: 'tool-approval'; toolName: 'confirm' } + > + type Lookup = Extract< + Interrupt, + { kind: 'client-tool-execution'; toolName: 'lookup' } + > + type Generic = Extract + + const check = ( + transferInterrupt: Transfer, + confirmInterrupt: Confirm, + lookupInterrupt: Lookup, + genericInterrupt: Generic, + ) => { + transferInterrupt.resolveInterrupt(true, { + editedArgs: { cents: 100 }, + payload: { note: 'approved' }, + }) + transferInterrupt.resolveInterrupt(false, { + payload: { reason: 'declined' }, + }) + // @ts-expect-error rejected approvals cannot edit tool input + transferInterrupt.resolveInterrupt(false, { editedArgs: { cents: 1 } }) + // @ts-expect-error approve payload uses the approve branch + transferInterrupt.resolveInterrupt(true, { + payload: { reason: 'wrong branch' }, + }) + + confirmInterrupt.resolveInterrupt(true) + confirmInterrupt.resolveInterrupt(false) + // @ts-expect-error omitted input schema forbids edited input + confirmInterrupt.resolveInterrupt(true, { editedArgs: { cents: 1 } }) + // @ts-expect-error omitted approval branches forbid payloads + confirmInterrupt.resolveInterrupt(false, { + payload: { reason: 'no branch' }, + }) + + lookupInterrupt.resolveInterrupt({ accountId: 'account-1' }) + // @ts-expect-error client-tool output is inferred + lookupInterrupt.resolveInterrupt({ account: 'account-1' }) + expectTypeOf(genericInterrupt.resolveInterrupt) + .parameter(0) + .toEqualTypeOf() + } + void check + }) +}) diff --git a/packages/ai-solid/tests/use-chat.test.ts b/packages/ai-solid/tests/use-chat.test.ts index f2c9c0524..23d25a806 100644 --- a/packages/ai-solid/tests/use-chat.test.ts +++ b/packages/ai-solid/tests/use-chat.test.ts @@ -1,15 +1,97 @@ -import { waitFor } from '@solidjs/testing-library' +import { renderHook, waitFor } from '@solidjs/testing-library' import type { ModelMessage } from '@tanstack/ai' +import { ChatClient } from '@tanstack/ai-client' import { describe, expect, it, vi } from 'vitest' import type { UIMessage } from '../src/types' +import { useChat } from '../src/use-chat' import { createMockConnectionAdapter, + createInterruptResumeSnapshot, createTextChunks, createToolCallChunks, renderUseChat, } from './test-utils' describe('useChat', () => { + describe('interrupt state', () => { + it('projects one immutable reactive snapshot with the deprecated pending alias', async () => { + const onInterruptStateChange = vi.fn() + const rendered = renderHook(() => + useChat({ + connection: createMockConnectionAdapter(), + initialResumeSnapshot: createInterruptResumeSnapshot(), + onInterruptStateChange, + }), + ) + const chat = rendered.result + + expect(Object.isFrozen(chat.interrupts())).toBe(true) + expect(chat.pendingInterrupts()).toBe(chat.interrupts()) + expect(chat.interrupts()[0]).toMatchObject({ + id: 'staged-interrupt', + status: 'staged', + }) + expect(chat.interrupts()[1]).toMatchObject({ + id: 'invalid-interrupt', + status: 'error', + error: { code: 'invalid-payload' }, + }) + expect(chat.interruptErrors()).toEqual([]) + expect(chat.resuming()).toBe(false) + expect(chat.interrupts()[0]).toEqual( + expect.objectContaining({ + resolveInterrupt: expect.any(Function), + cancel: expect.any(Function), + clearResolution: expect.any(Function), + }), + ) + + chat.resolveInterrupts(false) + await waitFor(() => { + expect(chat.interruptErrors()[0]?.code).toBe( + 'unsupported-bulk-operation', + ) + }) + expect(onInterruptStateChange).toHaveBeenLastCalledWith( + expect.objectContaining({ + interrupts: chat.interrupts(), + interruptErrors: chat.interruptErrors(), + }), + ) + }) + + it('delegates every root interrupt control to ChatClient', async () => { + const resolve = vi + .spyOn(ChatClient.prototype, 'resolveInterrupts') + .mockImplementation(() => {}) + const cancel = vi + .spyOn(ChatClient.prototype, 'cancelInterrupts') + .mockImplementation(() => {}) + const retry = vi + .spyOn(ChatClient.prototype, 'retryInterrupts') + .mockImplementation(() => {}) + const unsafe = vi + .spyOn(ChatClient.prototype, 'resumeInterruptsUnsafe') + .mockResolvedValue(true) + const rendered = renderHook(() => + useChat({ connection: createMockConnectionAdapter() }), + ) + const chat = rendered.result + const resolver = () => undefined + const resume = [{ interruptId: 'one', status: 'cancelled' as const }] + + chat.resolveInterrupts(resolver) + chat.cancelInterrupts() + chat.retryInterrupts() + await expect(chat.resumeInterruptsUnsafe(resume)).resolves.toBe(true) + + expect(resolve).toHaveBeenCalledWith(resolver) + expect(cancel).toHaveBeenCalledOnce() + expect(retry).toHaveBeenCalledOnce() + expect(unsafe).toHaveBeenCalledWith(resume, undefined) + }) + }) + describe('initialization', () => { it('should initialize with default state', () => { const adapter = createMockConnectionAdapter() diff --git a/packages/ai-svelte/src/create-chat.svelte.ts b/packages/ai-svelte/src/create-chat.svelte.ts index 91db27ceb..ed97a4d8c 100644 --- a/packages/ai-svelte/src/create-chat.svelte.ts +++ b/packages/ai-svelte/src/create-chat.svelte.ts @@ -3,7 +3,8 @@ import { createChatDevtoolsBridge } from '@tanstack/ai-client/devtools' import { onMount } from 'svelte' import type { ChatClientState, - ChatPendingInterrupt, + ChatInterrupt, + ChatInterruptState, ChatResumeState, ConnectionStatus, InferredClientContext, @@ -25,6 +26,9 @@ import type { UIMessage, } from './types' +const EMPTY_INTERRUPTS = Object.freeze([]) +const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) + /** * Creates a reactive chat instance for Svelte 5. * @@ -78,9 +82,12 @@ export function createChat< let resumeState = $state( options.initialResumeSnapshot?.resumeState ?? null, ) - let pendingInterrupts = $state>( - options.initialResumeSnapshot?.pendingInterrupts ?? [], - ) + let interruptState = $state.raw>({ + interrupts: EMPTY_INTERRUPTS, + pendingInterrupts: EMPTY_INTERRUPTS, + interruptErrors: EMPTY_INTERRUPT_ERRORS, + resuming: false, + }) // Structured-output `partial` / `final` are derived from `messages` — // specifically from the structured-output part on the latest assistant @@ -167,18 +174,22 @@ export function createChat< onSessionGeneratingChange: (isGenerating: boolean) => { sessionGenerating = isGenerating }, - onResumeStateChange: (nextResumeState, nextPendingInterrupts) => { + onResumeStateChange: (nextResumeState) => { resumeState = nextResumeState - pendingInterrupts = nextPendingInterrupts + }, + onInterruptStateChange: (nextInterruptState) => { + interruptState = nextInterruptState + options.onInterruptStateChange?.(nextInterruptState) }, }) function syncResumeState() { resumeState = client.getResumeState() - pendingInterrupts = client.getPendingInterrupts() + interruptState = client.getInterruptState() } messages = client.getMessages() + interruptState = client.getInterruptState() if (options.live) { client.subscribe() @@ -273,6 +284,29 @@ export function createChat< return result } + const resolveInterrupts = ( + resolution: boolean | ((interrupt: ChatInterrupt) => undefined), + ) => { + if (typeof resolution === 'boolean') { + client.resolveInterrupts(resolution) + } else { + client.resolveInterrupts(resolution) + } + } + + const cancelInterrupts = () => { + client.cancelInterrupts() + } + + const retryInterrupts = () => { + client.retryInterrupts() + } + + const resumeInterruptsUnsafe = ( + resumeItems: Array, + state?: ChatResumeState, + ) => client.resumeInterruptsUnsafe(resumeItems, state) + /** * @deprecated Use `updateForwardedProps` instead. * Both populate the same wire payload. @@ -355,8 +389,17 @@ export function createChat< get resumeState() { return resumeState }, + get interrupts() { + return interruptState.interrupts + }, get pendingInterrupts() { - return pendingInterrupts + return interruptState.interrupts + }, + get interruptErrors() { + return interruptState.interruptErrors + }, + get resuming() { + return interruptState.resuming }, get partial() { return partial @@ -373,6 +416,10 @@ export function createChat< clear, addToolResult, addToolApprovalResponse, + resolveInterrupts, + cancelInterrupts, + retryInterrupts, + resumeInterruptsUnsafe, resumeInterrupts, updateBody, updateForwardedProps, diff --git a/packages/ai-svelte/src/types.ts b/packages/ai-svelte/src/types.ts index e416fdb15..beb6371dc 100644 --- a/packages/ai-svelte/src/types.ts +++ b/packages/ai-svelte/src/types.ts @@ -7,9 +7,11 @@ import type { } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, + BoundInterrupts, ChatClientOptions, ChatClientState, - ChatPendingInterrupt, + ChatInterrupt, + ChatInterruptState, ChatRequestBody, ChatResumeState, ClientContextOptionFromTools, @@ -157,7 +159,22 @@ interface BaseCreateChatReturn< }) => Promise readonly resumeState: ChatResumeState | null - readonly pendingInterrupts: Array + readonly interrupts: BoundInterrupts + /** @deprecated Use `interrupts`. */ + readonly pendingInterrupts: BoundInterrupts + readonly interruptErrors: ChatInterruptState['interruptErrors'] + readonly resuming: boolean + resolveInterrupts: { + (approved: boolean): void + (resolver: (interrupt: ChatInterrupt) => undefined): void + } + cancelInterrupts: () => void + retryInterrupts: () => void + resumeInterruptsUnsafe: ( + resume: Array, + state?: ChatResumeState, + ) => Promise + /** @deprecated Use bound interrupt methods or `resumeInterruptsUnsafe`. */ resumeInterrupts: ( resume: Array, state?: ChatResumeState, diff --git a/packages/ai-svelte/tests/create-chat-types.test.ts b/packages/ai-svelte/tests/create-chat-types.test.ts index b647e49b1..28ab99a9a 100644 --- a/packages/ai-svelte/tests/create-chat-types.test.ts +++ b/packages/ai-svelte/tests/create-chat-types.test.ts @@ -4,7 +4,10 @@ */ import { describe, expectTypeOf, it } from 'vitest' -import type { StandardJSONSchemaV1 } from '@standard-schema/spec' +import type { + StandardJSONSchemaV1, + StandardSchemaV1, +} from '@standard-schema/spec' import { toolDefinition, type AnyClientTool } from '@tanstack/ai' import { clientTools, type StructuredOutputPart } from '@tanstack/ai-client' import { createChat } from '../src/create-chat.svelte' @@ -136,3 +139,137 @@ describe('createChat() return type (svelte)', () => { }) }) }) + +describe('createChat() interrupt types', () => { + it('preserves approval, generic, and client-tool inference', () => { + const inputSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { cents: 0 }, + output: { cents: 0 }, + }, + validate: (value: unknown) => ({ + value: + value !== null && + typeof value === 'object' && + 'cents' in value && + typeof value.cents === 'number' + ? { cents: value.cents } + : { cents: 0 }, + }), + }, + } + const approveSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { note: '' }, + output: { note: '' }, + }, + validate: () => ({ value: { note: '' } }), + }, + } + const rejectSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { reason: '' }, + output: { reason: '' }, + }, + validate: () => ({ value: { reason: '' } }), + }, + } + const outputSchema: StandardSchemaV1< + { accountId: string }, + { accountId: string } + > = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { accountId: '' }, + output: { accountId: '' }, + }, + validate: () => ({ value: { accountId: '' } }), + }, + } + const transfer = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema, + approvalSchema: { + approve: approveSchema, + reject: rejectSchema, + }, + }).client() + const confirm = toolDefinition({ + name: 'confirm', + description: 'Confirm without schemas', + needsApproval: true, + }).client() + const lookup = toolDefinition({ + name: 'lookup', + description: 'Lookup account', + outputSchema, + }).client(() => ({ accountId: 'account-1' })) + const tools = clientTools(transfer, confirm, lookup) + type Interrupt = CreateChatReturn['interrupts'][number] + type Transfer = Extract< + Interrupt, + { kind: 'tool-approval'; toolName: 'transfer' } + > + type Confirm = Extract< + Interrupt, + { kind: 'tool-approval'; toolName: 'confirm' } + > + type Lookup = Extract< + Interrupt, + { kind: 'client-tool-execution'; toolName: 'lookup' } + > + type Generic = Extract + + const check = ( + transferInterrupt: Transfer, + confirmInterrupt: Confirm, + lookupInterrupt: Lookup, + genericInterrupt: Generic, + ) => { + transferInterrupt.resolveInterrupt(true, { + editedArgs: { cents: 100 }, + payload: { note: 'approved' }, + }) + transferInterrupt.resolveInterrupt(false, { + payload: { reason: 'declined' }, + }) + // @ts-expect-error rejected approvals cannot edit tool input + transferInterrupt.resolveInterrupt(false, { editedArgs: { cents: 1 } }) + // @ts-expect-error approve payload uses the approve branch + transferInterrupt.resolveInterrupt(true, { + payload: { reason: 'wrong branch' }, + }) + + confirmInterrupt.resolveInterrupt(true) + confirmInterrupt.resolveInterrupt(false) + // @ts-expect-error omitted input schema forbids edited input + confirmInterrupt.resolveInterrupt(true, { editedArgs: { cents: 1 } }) + // @ts-expect-error omitted approval branches forbid payloads + confirmInterrupt.resolveInterrupt(false, { + payload: { reason: 'no branch' }, + }) + + lookupInterrupt.resolveInterrupt({ accountId: 'account-1' }) + expectTypeOf(lookupInterrupt.resolveInterrupt) + .parameter(0) + .toEqualTypeOf<{ accountId: string }>() + expectTypeOf(genericInterrupt.resolveInterrupt) + .parameter(0) + .toEqualTypeOf() + } + void check + }) +}) diff --git a/packages/ai-svelte/tests/test-utils.ts b/packages/ai-svelte/tests/test-utils.ts index 2f6fe6e64..ae356b74d 100644 --- a/packages/ai-svelte/tests/test-utils.ts +++ b/packages/ai-svelte/tests/test-utils.ts @@ -4,3 +4,79 @@ export { createTextChunks, createToolCallChunks, } from '../../ai-client/tests/test-utils' + +import type { ChatResumeSnapshotV2 } from '@tanstack/ai-client' + +export function createInterruptResumeSnapshot(): ChatResumeSnapshotV2 { + const pendingInterrupts = [ + { + id: 'staged-interrupt', + reason: 'confirmation', + metadata: { + 'tanstack:interruptBinding': { + kind: 'generic' as const, + interruptId: 'staged-interrupt', + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'none', + }, + }, + }, + { + id: 'invalid-interrupt', + reason: 'confirmation', + metadata: { + 'tanstack:interruptBinding': { + kind: 'generic' as const, + interruptId: 'invalid-interrupt', + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'none', + }, + }, + }, + ] + + return { + schemaVersion: 2, + resumeState: { threadId: 'thread-1', runId: 'run-1' }, + pendingInterrupts, + interruptState: { + recoveryState: { + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + pendingInterrupts, + }, + drafts: [ + { + interruptId: 'staged-interrupt', + response: { + interruptId: 'staged-interrupt', + status: 'resolved', + payload: { answer: 42 }, + }, + status: 'staged', + }, + { + interruptId: 'invalid-interrupt', + response: null, + status: 'error', + error: { + scope: 'item', + interruptId: 'invalid-interrupt', + code: 'invalid-payload', + message: 'Invalid persisted response', + source: 'client', + retryable: false, + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + }, + }, + ], + }, + } +} diff --git a/packages/ai-svelte/tests/use-chat.test.ts b/packages/ai-svelte/tests/use-chat.test.ts index fd6fe11ef..38bf3f4e1 100644 --- a/packages/ai-svelte/tests/use-chat.test.ts +++ b/packages/ai-svelte/tests/use-chat.test.ts @@ -1,12 +1,92 @@ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { ChatClient } from '@tanstack/ai-client' +import { tick } from 'svelte' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { createChat } from '../src/create-chat.svelte' -import { createMockConnectionAdapter, createTextChunks } from './test-utils' +import { + createInterruptResumeSnapshot, + createMockConnectionAdapter, + createTextChunks, +} from './test-utils' describe('createChat', () => { beforeEach(() => { vi.clearAllMocks() }) + afterEach(() => { + vi.restoreAllMocks() + }) + + describe('interrupt state', () => { + it('projects one immutable reactive snapshot with the deprecated pending alias', async () => { + const onInterruptStateChange = vi.fn() + const chat = createChat({ + connection: createMockConnectionAdapter(), + initialResumeSnapshot: createInterruptResumeSnapshot(), + onInterruptStateChange, + }) + + expect(Object.isFrozen(chat.interrupts)).toBe(true) + expect(chat.pendingInterrupts).toBe(chat.interrupts) + expect(chat.interrupts[0]).toMatchObject({ + id: 'staged-interrupt', + status: 'staged', + }) + expect(chat.interrupts[1]).toMatchObject({ + id: 'invalid-interrupt', + status: 'error', + error: { code: 'invalid-payload' }, + }) + expect(chat.interruptErrors).toEqual([]) + expect(chat.resuming).toBe(false) + expect(chat.interrupts[0]).toEqual( + expect.objectContaining({ + resolveInterrupt: expect.any(Function), + cancel: expect.any(Function), + clearResolution: expect.any(Function), + }), + ) + + chat.resolveInterrupts(false) + await tick() + expect(chat.interruptErrors[0]?.code).toBe('unsupported-bulk-operation') + expect(onInterruptStateChange).toHaveBeenLastCalledWith( + expect.objectContaining({ + interrupts: chat.interrupts, + interruptErrors: chat.interruptErrors, + }), + ) + }) + + it('delegates every root interrupt control to ChatClient', async () => { + const resolve = vi + .spyOn(ChatClient.prototype, 'resolveInterrupts') + .mockImplementation(() => {}) + const cancel = vi + .spyOn(ChatClient.prototype, 'cancelInterrupts') + .mockImplementation(() => {}) + const retry = vi + .spyOn(ChatClient.prototype, 'retryInterrupts') + .mockImplementation(() => {}) + const unsafe = vi + .spyOn(ChatClient.prototype, 'resumeInterruptsUnsafe') + .mockResolvedValue(true) + const chat = createChat({ connection: createMockConnectionAdapter() }) + const resolver = () => undefined + const resume = [{ interruptId: 'one', status: 'cancelled' as const }] + + chat.resolveInterrupts(resolver) + chat.cancelInterrupts() + chat.retryInterrupts() + await expect(chat.resumeInterruptsUnsafe(resume)).resolves.toBe(true) + + expect(resolve).toHaveBeenCalledWith(resolver) + expect(cancel).toHaveBeenCalledOnce() + expect(retry).toHaveBeenCalledOnce() + expect(unsafe).toHaveBeenCalledWith(resume, undefined) + }) + }) + it('should initialize with empty messages', () => { const mockConnection = createMockConnectionAdapter({ chunks: [] }) diff --git a/packages/ai-vue/src/types.ts b/packages/ai-vue/src/types.ts index 3c824e7df..38a56bc14 100644 --- a/packages/ai-vue/src/types.ts +++ b/packages/ai-vue/src/types.ts @@ -7,9 +7,11 @@ import type { } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, + BoundInterrupts, ChatClientOptions, ChatClientState, - ChatPendingInterrupt, + ChatInterrupt, + ChatInterruptState, ChatRequestBody, ChatResumeState, ClientContextOptionFromTools, @@ -155,7 +157,24 @@ interface BaseUseChatReturn< }) => Promise resumeState: DeepReadonly> - pendingInterrupts: DeepReadonly>> + interrupts: DeepReadonly>> + /** @deprecated Use `interrupts`. */ + pendingInterrupts: DeepReadonly>> + interruptErrors: DeepReadonly< + ShallowRef['interruptErrors']> + > + resuming: DeepReadonly> + resolveInterrupts: { + (approved: boolean): void + (resolver: (interrupt: ChatInterrupt) => undefined): void + } + cancelInterrupts: () => void + retryInterrupts: () => void + resumeInterruptsUnsafe: ( + resume: Array, + state?: ChatResumeState, + ) => Promise + /** @deprecated Use bound interrupt methods or `resumeInterruptsUnsafe`. */ resumeInterrupts: ( resume: Array, state?: ChatResumeState, diff --git a/packages/ai-vue/src/use-chat.ts b/packages/ai-vue/src/use-chat.ts index e6bb219d9..ec6aef00f 100644 --- a/packages/ai-vue/src/use-chat.ts +++ b/packages/ai-vue/src/use-chat.ts @@ -19,7 +19,8 @@ import type { } from '@tanstack/ai' import type { ChatClientState, - ChatPendingInterrupt, + ChatInterrupt, + ChatInterruptState, ChatResumeState, ConnectionStatus, InferredClientContext, @@ -33,6 +34,9 @@ import type { UseChatReturn, } from './types' +const EMPTY_INTERRUPTS = Object.freeze([]) +const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) + export function useChat< const TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, @@ -59,9 +63,12 @@ export function useChat< const resumeState = shallowRef( options.initialResumeSnapshot?.resumeState ?? null, ) - const pendingInterrupts = shallowRef>( - options.initialResumeSnapshot?.pendingInterrupts ?? [], - ) + const interruptState = shallowRef>({ + interrupts: EMPTY_INTERRUPTS, + pendingInterrupts: EMPTY_INTERRUPTS, + interruptErrors: EMPTY_INTERRUPT_ERRORS, + resuming: false, + }) // Structured-output `partial` / `final` are derived from `messages` — // specifically from the structured-output part on the latest assistant @@ -151,18 +158,22 @@ export function useChat< onSessionGeneratingChange: (isGenerating: boolean) => { sessionGenerating.value = isGenerating }, - onResumeStateChange: (nextResumeState, nextPendingInterrupts) => { + onResumeStateChange: (nextResumeState) => { resumeState.value = nextResumeState - pendingInterrupts.value = nextPendingInterrupts + }, + onInterruptStateChange: (nextInterruptState) => { + interruptState.value = nextInterruptState + options.onInterruptStateChange?.(nextInterruptState) }, }) function syncResumeState() { resumeState.value = client.getResumeState() - pendingInterrupts.value = client.getPendingInterrupts() + interruptState.value = client.getInterruptState() } messages.value = client.getMessages() + interruptState.value = client.getInterruptState() // Sync body / forwardedProps changes to the client. // Both populate the same wire payload; `forwardedProps` is preferred @@ -280,6 +291,34 @@ export function useChat< return result } + const interrupts = computed(() => interruptState.value.interrupts) + const pendingInterrupts = computed(() => interruptState.value.interrupts) + const interruptErrors = computed(() => interruptState.value.interruptErrors) + const resuming = computed(() => interruptState.value.resuming) + + const resolveInterrupts = ( + resolution: boolean | ((interrupt: ChatInterrupt) => undefined), + ) => { + if (typeof resolution === 'boolean') { + client.resolveInterrupts(resolution) + } else { + client.resolveInterrupts(resolution) + } + } + + const cancelInterrupts = () => { + client.cancelInterrupts() + } + + const retryInterrupts = () => { + client.retryInterrupts() + } + + const resumeInterruptsUnsafe = ( + resumeItems: Array, + state?: ChatResumeState, + ) => client.resumeInterruptsUnsafe(resumeItems, state) + // The "active" structured-output part is the one on the assistant message // that follows the latest user message. No such message exists between // sendMessage() and the first chunk, so partial/final naturally read as @@ -342,7 +381,14 @@ export function useChat< addToolResult, addToolApprovalResponse, resumeState: readonly(resumeState), + interrupts: readonly(interrupts), pendingInterrupts: readonly(pendingInterrupts), + interruptErrors: readonly(interruptErrors), + resuming: readonly(resuming), + resolveInterrupts, + cancelInterrupts, + retryInterrupts, + resumeInterruptsUnsafe, resumeInterrupts, partial: readonly(partial), final: readonly(final), diff --git a/packages/ai-vue/tests/test-utils.ts b/packages/ai-vue/tests/test-utils.ts index 3f8c05c46..b0b68951c 100644 --- a/packages/ai-vue/tests/test-utils.ts +++ b/packages/ai-vue/tests/test-utils.ts @@ -1,4 +1,4 @@ -import type { UIMessage } from '@tanstack/ai-client' +import type { ChatResumeSnapshotV2, UIMessage } from '@tanstack/ai-client' import { mount } from '@vue/test-utils' import { defineComponent } from 'vue' import type { UseChatOptions } from '../src/types' @@ -11,6 +11,80 @@ export { createToolCallChunks, } from '../../ai-client/tests/test-utils' +export function createInterruptResumeSnapshot(): ChatResumeSnapshotV2 { + const pendingInterrupts = [ + { + id: 'staged-interrupt', + reason: 'confirmation', + metadata: { + 'tanstack:interruptBinding': { + kind: 'generic' as const, + interruptId: 'staged-interrupt', + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'none', + }, + }, + }, + { + id: 'invalid-interrupt', + reason: 'confirmation', + metadata: { + 'tanstack:interruptBinding': { + kind: 'generic' as const, + interruptId: 'invalid-interrupt', + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'none', + }, + }, + }, + ] + + return { + schemaVersion: 2, + resumeState: { threadId: 'thread-1', runId: 'run-1' }, + pendingInterrupts, + interruptState: { + recoveryState: { + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + pendingInterrupts, + }, + drafts: [ + { + interruptId: 'staged-interrupt', + response: { + interruptId: 'staged-interrupt', + status: 'resolved', + payload: { answer: 42 }, + }, + status: 'staged', + }, + { + interruptId: 'invalid-interrupt', + response: null, + status: 'error', + error: { + scope: 'item', + interruptId: 'invalid-interrupt', + code: 'invalid-payload', + message: 'Invalid persisted response', + source: 'client', + retryable: false, + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + }, + }, + ], + }, + } +} + /** * Render the useChat hook with testing utilities * @@ -54,6 +128,14 @@ export function renderUseChat(options?: UseChatOptions) { setMessages: hook.setMessages, addToolResult: hook.addToolResult, addToolApprovalResponse: hook.addToolApprovalResponse, + interrupts: hook.interrupts, + pendingInterrupts: hook.pendingInterrupts, + interruptErrors: hook.interruptErrors, + resuming: hook.resuming, + resolveInterrupts: hook.resolveInterrupts, + cancelInterrupts: hook.cancelInterrupts, + retryInterrupts: hook.retryInterrupts, + resumeInterruptsUnsafe: hook.resumeInterruptsUnsafe, } } diff --git a/packages/ai-vue/tests/use-chat-types.test.ts b/packages/ai-vue/tests/use-chat-types.test.ts index 518c47b02..220248368 100644 --- a/packages/ai-vue/tests/use-chat-types.test.ts +++ b/packages/ai-vue/tests/use-chat-types.test.ts @@ -4,7 +4,10 @@ */ import { describe, expectTypeOf, it } from 'vitest' -import type { StandardJSONSchemaV1 } from '@standard-schema/spec' +import type { + StandardJSONSchemaV1, + StandardSchemaV1, +} from '@standard-schema/spec' import { toolDefinition, type AnyClientTool } from '@tanstack/ai' import { clientTools, type StructuredOutputPart } from '@tanstack/ai-client' import type { ShallowRef } from 'vue' @@ -131,3 +134,137 @@ describe('useChat() return type (vue)', () => { }) }) }) + +describe('useChat() interrupt types', () => { + it('preserves approval, generic, and client-tool inference', () => { + const inputSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { cents: 0 }, + output: { cents: 0 }, + }, + validate: (value: unknown) => ({ + value: + value !== null && + typeof value === 'object' && + 'cents' in value && + typeof value.cents === 'number' + ? { cents: value.cents } + : { cents: 0 }, + }), + }, + } + const approveSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { note: '' }, + output: { note: '' }, + }, + validate: () => ({ value: { note: '' } }), + }, + } + const rejectSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { reason: '' }, + output: { reason: '' }, + }, + validate: () => ({ value: { reason: '' } }), + }, + } + const outputSchema: StandardSchemaV1< + { accountId: string }, + { accountId: string } + > = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { + input: { accountId: '' }, + output: { accountId: '' }, + }, + validate: () => ({ value: { accountId: '' } }), + }, + } + const transfer = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema, + approvalSchema: { + approve: approveSchema, + reject: rejectSchema, + }, + }).client() + const confirm = toolDefinition({ + name: 'confirm', + description: 'Confirm without schemas', + needsApproval: true, + }).client() + const lookup = toolDefinition({ + name: 'lookup', + description: 'Lookup account', + outputSchema, + }).client(() => ({ accountId: 'account-1' })) + const tools = clientTools(transfer, confirm, lookup) + type Interrupt = UseChatReturn['interrupts']['value'][number] + type Transfer = Extract< + Interrupt, + { kind: 'tool-approval'; toolName: 'transfer' } + > + type Confirm = Extract< + Interrupt, + { kind: 'tool-approval'; toolName: 'confirm' } + > + type Lookup = Extract< + Interrupt, + { kind: 'client-tool-execution'; toolName: 'lookup' } + > + type Generic = Extract + + const check = ( + transferInterrupt: Transfer, + confirmInterrupt: Confirm, + lookupInterrupt: Lookup, + genericInterrupt: Generic, + ) => { + transferInterrupt.resolveInterrupt(true, { + editedArgs: { cents: 100 }, + payload: { note: 'approved' }, + }) + transferInterrupt.resolveInterrupt(false, { + payload: { reason: 'declined' }, + }) + // @ts-expect-error rejected approvals cannot edit tool input + transferInterrupt.resolveInterrupt(false, { editedArgs: { cents: 1 } }) + // @ts-expect-error approve payload uses the approve branch + transferInterrupt.resolveInterrupt(true, { + payload: { reason: 'wrong branch' }, + }) + + confirmInterrupt.resolveInterrupt(true) + confirmInterrupt.resolveInterrupt(false) + // @ts-expect-error omitted input schema forbids edited input + confirmInterrupt.resolveInterrupt(true, { editedArgs: { cents: 1 } }) + // @ts-expect-error omitted approval branches forbid payloads + confirmInterrupt.resolveInterrupt(false, { + payload: { reason: 'no branch' }, + }) + + lookupInterrupt.resolveInterrupt({ accountId: 'account-1' }) + expectTypeOf(lookupInterrupt.resolveInterrupt) + .parameter(0) + .toEqualTypeOf<{ accountId: string }>() + expectTypeOf(genericInterrupt.resolveInterrupt) + .parameter(0) + .toEqualTypeOf() + } + void check + }) +}) diff --git a/packages/ai-vue/tests/use-chat.test.ts b/packages/ai-vue/tests/use-chat.test.ts index b513ef2c0..a334a392c 100644 --- a/packages/ai-vue/tests/use-chat.test.ts +++ b/packages/ai-vue/tests/use-chat.test.ts @@ -1,15 +1,97 @@ import type { ModelMessage } from '@tanstack/ai' +import { ChatClient } from '@tanstack/ai-client' import { flushPromises } from '@vue/test-utils' -import { describe, expect, it, vi } from 'vitest' +import { afterEach, describe, expect, it, vi } from 'vitest' import type { UIMessage } from '../src/types' import { createMockConnectionAdapter, + createInterruptResumeSnapshot, createTextChunks, createToolCallChunks, renderUseChat, } from './test-utils' describe('useChat', () => { + afterEach(() => { + vi.restoreAllMocks() + }) + + describe('interrupt state', () => { + it('projects one immutable reactive snapshot with the deprecated pending alias', async () => { + const onInterruptStateChange = vi.fn() + const { result } = renderUseChat({ + connection: createMockConnectionAdapter(), + initialResumeSnapshot: createInterruptResumeSnapshot(), + onInterruptStateChange, + }) + + expect(Object.isFrozen(result.current.interrupts)).toBe(true) + expect(result.current.pendingInterrupts).toBe(result.current.interrupts) + expect(result.current.interrupts[0]).toMatchObject({ + id: 'staged-interrupt', + status: 'staged', + }) + expect(result.current.interrupts[1]).toMatchObject({ + id: 'invalid-interrupt', + status: 'error', + error: { code: 'invalid-payload' }, + }) + expect(result.current.interruptErrors).toEqual([]) + expect(result.current.resuming).toBe(false) + expect(result.current.interrupts[0]).toEqual( + expect.objectContaining({ + resolveInterrupt: expect.any(Function), + cancel: expect.any(Function), + clearResolution: expect.any(Function), + }), + ) + + result.current.resolveInterrupts(false) + await flushPromises() + expect(result.current.interruptErrors[0]?.code).toBe( + 'unsupported-bulk-operation', + ) + expect(onInterruptStateChange).toHaveBeenLastCalledWith( + expect.objectContaining({ + interrupts: result.current.interrupts, + interruptErrors: result.current.interruptErrors, + }), + ) + }) + + it('delegates every root interrupt control to ChatClient', async () => { + const resolve = vi + .spyOn(ChatClient.prototype, 'resolveInterrupts') + .mockImplementation(() => {}) + const cancel = vi + .spyOn(ChatClient.prototype, 'cancelInterrupts') + .mockImplementation(() => {}) + const retry = vi + .spyOn(ChatClient.prototype, 'retryInterrupts') + .mockImplementation(() => {}) + const unsafe = vi + .spyOn(ChatClient.prototype, 'resumeInterruptsUnsafe') + .mockResolvedValue(true) + const { result } = renderUseChat({ + connection: createMockConnectionAdapter(), + }) + const resolver = () => undefined + const resume = [{ interruptId: 'one', status: 'cancelled' as const }] + + result.current.resolveInterrupts(resolver) + result.current.cancelInterrupts() + result.current.retryInterrupts() + await expect(result.current.resumeInterruptsUnsafe(resume)).resolves.toBe( + true, + ) + + expect(resolve).toHaveBeenCalledWith(resolver) + expect(cancel).toHaveBeenCalledOnce() + expect(retry).toHaveBeenCalledOnce() + expect(unsafe).toHaveBeenCalledWith(resume, undefined) + }) + }) + describe('initialization', () => { it('should initialize with default state', () => { const adapter = createMockConnectionAdapter() diff --git a/packages/ai/package.json b/packages/ai/package.json index f88be6ae6..bc8c79a1e 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -86,9 +86,12 @@ ], "dependencies": { "@ag-ui/core": "^0.0.57", + "@noble/hashes": "^2.0.1", "@standard-schema/spec": "^1.1.0", "@tanstack/ai-event-client": "workspace:*", "@tanstack/ai-utils": "workspace:*", + "ajv": "^8.20.0", + "ajv-formats": "^3.0.1", "partial-json": "^0.1.7" }, "peerDependencies": { diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index b8d9341ef..e39a8689b 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -11,6 +11,11 @@ import { stripToSpecMiddleware } from '../../strip-to-spec-middleware' import { streamToText } from '../../stream-to-response.js' import { resolveDebugOption } from '../../logger/resolve' import { EventType } from '../../types' +import { getInterruptPersistence } from '../../interrupts' +import { + canonicalInterruptJson, + digestInterruptJson, +} from '../../interrupt-serialization' import { normalizeToolResult } from '../../utilities/tool-result' import { isProviderExecutedToolCall } from '../../utilities/provider-executed' import { LazyToolManager } from './tools/lazy-tool-manager' @@ -25,6 +30,10 @@ import { isStandardSchema, parseWithStandardSchema, } from './tools/schema-converter' +import { + hashSchemaInput, + normalizeApprovalSchema, +} from './tools/approval-schema' import { maxIterations as maxIterationsStrategy } from './agent-loop-strategies' import { convertMessagesToModelMessages, generateMessageId } from './messages' import { MiddlewareRunner } from './middleware/compose' @@ -32,11 +41,18 @@ import { provideSandboxRuntime } from './middleware/sandbox-runtime' import { CapabilityRegistry } from './middleware/capabilities' import { validateCapabilities } from './middleware/validate' import { MCPManager } from './mcp/manager' +import type { + InterruptBinding, + InterruptRecoveryStateV1, + InterruptSubmissionError, + ToolApprovalResolution, +} from '../../interrupts' import type { ApprovalRequest, ClientToolRequest, ToolResult, } from './tools/tool-calls' +import type { ApprovalSchemaConfig } from './tools/tool-definition' import type { AnyTextAdapter, StructuredOutputOptions, @@ -52,6 +68,7 @@ import type { Interrupt, JSONSchema, LazyToolsConfig, + MessagesSnapshotEvent, ModelMessage, RunFinishedEvent, SchemaInput, @@ -96,6 +113,174 @@ import type { ChatMCPOptions } from './mcp/types' export const kind = 'text' as const type AnyRuntimeTool = AnyTool +type RuntimeToolWithApproval = AnyRuntimeTool & { + approvalSchema?: ApprovalSchemaConfig +} +const interruptBindingMetadataKey = 'tanstack:interruptBinding' + +interface StructuralInterruptFailure { + error: Error + errors: ReadonlyArray + recovery?: InterruptRecoveryStateV1 +} + +function isInterruptSubmissionError( + value: unknown, +): value is InterruptSubmissionError { + if (value === null || typeof value !== 'object' || Array.isArray(value)) { + return false + } + if ( + !('scope' in value) || + !('code' in value) || + !('message' in value) || + !('source' in value) || + !('retryable' in value) || + !('threadId' in value) || + !('interruptedRunId' in value) || + !('generation' in value) || + typeof value.code !== 'string' || + typeof value.message !== 'string' || + typeof value.retryable !== 'boolean' || + typeof value.threadId !== 'string' || + typeof value.interruptedRunId !== 'string' || + typeof value.generation !== 'number' + ) { + return false + } + if (value.scope === 'item') { + return ( + 'interruptId' in value && + typeof value.interruptId === 'string' && + (value.source === 'client' || value.source === 'server') + ) + } + return ( + value.scope === 'batch' && + 'interruptIds' in value && + Array.isArray(value.interruptIds) && + value.interruptIds.every((id) => typeof id === 'string') && + (value.source === 'client' || + value.source === 'server' || + value.source === 'transport') + ) +} + +function isInterruptRecoveryState( + value: unknown, +): value is InterruptRecoveryStateV1 { + return ( + value !== null && + typeof value === 'object' && + !Array.isArray(value) && + 'schemaVersion' in value && + value.schemaVersion === 1 && + 'state' in value && + typeof value.state === 'string' && + 'threadId' in value && + typeof value.threadId === 'string' && + 'interruptedRunId' in value && + typeof value.interruptedRunId === 'string' && + 'generation' in value && + typeof value.generation === 'number' && + 'pendingInterrupts' in value && + Array.isArray(value.pendingInterrupts) + ) +} + +function structuralInterruptFailure( + error: unknown, +): StructuralInterruptFailure | undefined { + if ( + !(error instanceof Error) || + error.name !== 'InterruptResumeValidationError' || + !('errors' in error) || + !Array.isArray(error.errors) || + error.errors.length === 0 || + !error.errors.every(isInterruptSubmissionError) + ) { + return undefined + } + const recovery = + 'recovery' in error && isInterruptRecoveryState(error.recovery) + ? error.recovery + : undefined + return { + error, + errors: error.errors, + ...(recovery !== undefined ? { recovery } : {}), + } +} + +function normalizePublicInterruptBinding( + value: unknown, + expectedInterruptId: string, +): InterruptBinding | undefined { + if (value === null || typeof value !== 'object' || Array.isArray(value)) { + return undefined + } + const binding: Record = Object.fromEntries( + Object.entries(value), + ) + if ( + binding.interruptId !== expectedInterruptId || + typeof binding.interruptedRunId !== 'string' || + typeof binding.generation !== 'number' || + !Number.isInteger(binding.generation) || + binding.generation < 0 || + typeof binding.responseSchemaHash !== 'string' || + (binding.expiresAt !== undefined && typeof binding.expiresAt !== 'string') + ) { + return undefined + } + const base = { + interruptId: binding.interruptId, + interruptedRunId: binding.interruptedRunId, + generation: binding.generation, + responseSchemaHash: binding.responseSchemaHash, + ...(typeof binding.expiresAt === 'string' + ? { expiresAt: binding.expiresAt } + : {}), + } + if (binding.kind === 'generic') { + return { kind: binding.kind, ...base } + } + if ( + typeof binding.toolName !== 'string' || + typeof binding.toolCallId !== 'string' + ) { + return undefined + } + if ( + binding.kind === 'client-tool-execution' && + typeof binding.outputSchemaHash === 'string' + ) { + return { + kind: binding.kind, + ...base, + toolName: binding.toolName, + toolCallId: binding.toolCallId, + outputSchemaHash: binding.outputSchemaHash, + } + } + if ( + binding.kind === 'tool-approval' && + Object.prototype.hasOwnProperty.call(binding, 'originalArgs') && + typeof binding.inputSchemaHash === 'string' && + typeof binding.approvalSchemaHash === 'string' + ) { + return { + kind: binding.kind, + ...base, + toolName: binding.toolName, + toolCallId: binding.toolCallId, + originalArgs: binding.originalArgs, + inputSchemaHash: binding.inputSchemaHash, + approvalSchemaHash: binding.approvalSchemaHash, + } + } + return undefined +} // The leaf context-inference primitives (KnownContext, MergeContext, // UnionToIntersection, DefinedContext, ContextFromTool, ContextFromMiddleware) @@ -255,6 +440,8 @@ export interface TextActivityOptions< runId?: TextOptions['runId'] /** Parent run ID for AG-UI protocol nested run correlation. */ parentRunId?: TextOptions['parentRunId'] + /** Application state mirrored in a STATE_SNAPSHOT before an interrupt terminal. */ + state?: TextOptions['state'] /** * AG-UI interrupt resume responses. Persistence middleware validates these * before accepting new input on a thread with pending interrupts. @@ -542,10 +729,12 @@ class TextEngine< private toolPhase: ToolPhaseResult = 'continue' private cyclePhase: CyclePhase = 'processText' // Client state extracted from initial messages (before conversion to ModelMessage) - private readonly initialApprovals: Map + private readonly initialApprovals: Map private readonly initialClientToolResults: Map - private readonly resumeApprovals = new Map() + private readonly resumeApprovals = new Map() private readonly resumeClientToolResults = new Map() + private readonly resumeDeniedToolResults = new Map() + private readonly resumeCancelledToolCallIds = new Set() // AG-UI protocol IDs private readonly threadId: string @@ -669,6 +858,7 @@ class TextEngine< requestId: this.requestId, streamId: this.streamId, runId: this.runIdOverride ?? this.requestId, + parentRunId: this.parentRunIdOverride, threadId: this.threadId, // Legacy alias kept on the ctx so middleware that reads // `ctx.conversationId` keeps working. Always equals `threadId`. @@ -926,6 +1116,41 @@ class TextEngine< } } } catch (error: unknown) { + if ( + error instanceof Error && + error.name === 'InterruptReplaySignal' && + 'continuationRunId' in error && + typeof error.continuationRunId === 'string' + ) { + this.terminalHookCalled = true + yield { + type: EventType.RUN_FINISHED, + timestamp: Date.now(), + threadId: this.threadId, + runId: this.runIdOverride ?? this.requestId, + finishReason: 'stop', + outcome: { type: 'success' }, + result: { + replayed: true, + continuationRunId: error.continuationRunId, + }, + } + return + } + const interruptFailure = structuralInterruptFailure(error) + if (interruptFailure) { + this.terminalHookCalled = true + this.logger.errors('chat interrupt resume failed', { + error, + threadId: this.middlewareCtx.threadId, + }) + await this.middlewareRunner.runOnError(this.middlewareCtx, { + error: interruptFailure.error, + duration: Date.now() - this.streamStartTime, + }) + yield this.buildInterruptRunErrorChunk(error) + return + } if (!this.terminalHookCalled) { this.terminalHookCalled = true if (error instanceof MiddlewareAbortError) { @@ -1061,6 +1286,13 @@ class TextEngine< : undefined const { approvals } = this.collectClientState() + const adapterApprovals = new Map() + for (const [approvalId, resolution] of approvals) { + adapterApprovals.set( + approvalId, + typeof resolution === 'boolean' ? resolution : resolution.approved, + ) + } for await (const chunk of this.adapter.chatStream({ model: this.params.model, @@ -1077,7 +1309,7 @@ class TextEngine< // Expose provided capabilities (e.g. sandbox) to harness adapters. capabilities: this.middlewareCtx, // Client approval decisions, for harness interactive-approval resolution. - approvals, + approvals: adapterApprovals, ...(combinedSchema ? { outputSchema: combinedSchema } : {}), })) { if (this.isCancelled()) { @@ -1391,6 +1623,10 @@ class TextEngine< }, this.middlewareCtx.context, this.toolAbortSignal, + { + deniedToolResults: this.resumeDeniedToolResults, + cancelledToolCallIds: this.resumeCancelledToolCallIds, + }, ) // Consume the async generator, yielding custom events and collecting the return value @@ -1410,13 +1646,6 @@ class TextEngine< needsClientExecution: executionResult.needsClientExecution, }) - // Build args lookup so buildToolResultChunks can emit TOOL_CALL_START + - // TOOL_CALL_ARGS before TOOL_CALL_END during continuation re-executions. - const argsMap = new Map() - for (const tc of pendingToolCalls) { - argsMap.set(tc.id, tc.function.arguments) - } - if ( executionResult.needsApproval.length > 0 || executionResult.needsClientExecution.length > 0 @@ -1427,28 +1656,23 @@ class TextEngine< for (const chunk of this.buildToolResultChunks( executionResult.results, finishEvent, - argsMap, )) { yield* this.pipeThroughMiddleware(chunk) } } - yield* this.pipeThroughMiddleware( - this.buildInterruptFinishedChunk( - finishEvent, - executionResult.needsApproval, - executionResult.needsClientExecution, - ), + const emitted = yield* this.emitActionableInterruptBoundary( + finishEvent, + executionResult.needsApproval, + executionResult.needsClientExecution, ) - - this.setToolPhase('wait') - return 'wait' + this.setToolPhase(emitted ? 'wait' : 'stop') + return emitted ? 'wait' : 'stop' } const toolResultChunks = this.buildToolResultChunks( executionResult.results, finishEvent, - argsMap, ) for (const chunk of toolResultChunks) { @@ -1550,6 +1774,10 @@ class TextEngine< }, this.middlewareCtx.context, this.toolAbortSignal, + { + deniedToolResults: this.resumeDeniedToolResults, + cancelledToolCallIds: this.resumeCancelledToolCallIds, + }, ) // Consume the async generator, yielding custom events and collecting the return value @@ -1584,15 +1812,12 @@ class TextEngine< } } - yield* this.pipeThroughMiddleware( - this.buildInterruptFinishedChunk( - finishEvent, - executionResult.needsApproval, - executionResult.needsClientExecution, - ), + const emitted = yield* this.emitActionableInterruptBoundary( + finishEvent, + executionResult.needsApproval, + executionResult.needsClientExecution, ) - - this.setToolPhase('wait') + this.setToolPhase(emitted ? 'wait' : 'stop') return } @@ -1677,10 +1902,10 @@ class TextEngine< private extractClientStateFromOriginalMessages( originalMessages: Array, ): { - approvals: Map + approvals: Map clientToolResults: Map } { - const approvals = new Map() + const approvals = new Map() const clientToolResults = new Map() for (const message of originalMessages) { @@ -1709,7 +1934,7 @@ class TextEngine< } private collectClientState(): { - approvals: Map + approvals: Map clientToolResults: Map } { // Start with the initial client state extracted from original messages @@ -1766,34 +1991,62 @@ class TextEngine< const interrupts: Array = [] for (const approval of approvals) { + const tool = this.tools.find( + (candidate) => candidate.name === approval.toolName, + ) as RuntimeToolWithApproval | undefined + const normalized = normalizeApprovalSchema( + tool?.approvalSchema, + tool?.inputSchema, + ) interrupts.push({ id: approval.approvalId, - reason: 'approval_required', + reason: 'tool_call', message: `Approval required to run ${approval.toolName}`, toolCallId: approval.toolCallId, - responseSchema: { - type: 'object', - properties: { approved: { type: 'boolean' } }, - required: ['approved'], - }, + responseSchema: normalized.responseSchema, metadata: { kind: 'approval', toolName: approval.toolName, input: approval.input, + [interruptBindingMetadataKey]: { + kind: 'tool-approval', + interruptId: approval.approvalId, + toolName: approval.toolName, + toolCallId: approval.toolCallId, + originalArgs: approval.input, + inputSchemaHash: hashSchemaInput(tool?.inputSchema), + approvalSchemaHash: normalized.approvalSchemaHash, + responseSchemaHash: normalized.responseSchemaHash, + }, }, }) } for (const clientTool of clientRequests) { + const tool = this.tools.find( + (candidate) => candidate.name === clientTool.toolName, + ) + const responseSchema = convertSchemaToJsonSchema(tool?.outputSchema) ?? {} interrupts.push({ id: `client_tool_${clientTool.toolCallId}`, - reason: 'client_tool_input', + reason: 'tanstack:client_tool_execution', message: `Client tool ${clientTool.toolName} is ready to run`, toolCallId: clientTool.toolCallId, + responseSchema, metadata: { kind: 'client_tool', toolName: clientTool.toolName, input: clientTool.input, + [interruptBindingMetadataKey]: { + kind: 'client-tool-execution', + interruptId: `client_tool_${clientTool.toolCallId}`, + toolName: clientTool.toolName, + toolCallId: clientTool.toolCallId, + outputSchemaHash: hashSchemaInput(tool?.outputSchema), + responseSchemaHash: digestInterruptJson( + canonicalInterruptJson(responseSchema), + ), + }, }, }) } @@ -1816,10 +2069,198 @@ class TextEngine< } } + private buildMessagesSnapshotChunk(): StreamChunk { + const messages: MessagesSnapshotEvent['messages'] = this.messages.map( + (message, index) => { + const content = + typeof message.content === 'string' + ? message.content + : message.content === null + ? undefined + : JSON.stringify(message.content) + return { + id: `snapshot_${this.runIdOverride ?? this.requestId}_${index}`, + role: message.role, + ...(content !== undefined ? { content } : {}), + ...('toolCalls' in message && message.toolCalls + ? { toolCalls: message.toolCalls } + : {}), + ...('toolCallId' in message && message.toolCallId + ? { toolCallId: message.toolCallId } + : {}), + } as MessagesSnapshotEvent['messages'][number] + }, + ) + return { + type: EventType.MESSAGES_SNAPSHOT, + timestamp: Date.now(), + model: this.params.model, + messages, + } + } + + private publicInterruptTerminal(chunk: StreamChunk): StreamChunk { + if ( + chunk.type !== EventType.RUN_FINISHED || + chunk.outcome?.type !== 'interrupt' + ) { + return chunk + } + return { + ...chunk, + outcome: { + ...chunk.outcome, + interrupts: chunk.outcome.interrupts.map((interrupt) => { + if ( + !interrupt.metadata || + typeof interrupt.metadata !== 'object' || + Array.isArray(interrupt.metadata) + ) { + return interrupt + } + const metadata = { ...interrupt.metadata } + const binding = normalizePublicInterruptBinding( + metadata[interruptBindingMetadataKey], + interrupt.id, + ) + if (binding) { + metadata[interruptBindingMetadataKey] = binding + } else { + delete metadata[interruptBindingMetadataKey] + } + return { ...interrupt, metadata } + }), + }, + } + } + + private interruptFailure(error: unknown): { + message: string + code: string + errors?: ReadonlyArray + recovery?: InterruptRecoveryStateV1 + } { + const structured = structuralInterruptFailure(error) + if (structured) { + return { + message: structured.error.message, + code: structured.errors[0]?.code ?? 'server', + errors: structured.errors, + ...(structured.recovery !== undefined + ? { recovery: structured.recovery } + : {}), + } + } + if (error && typeof error === 'object' && 'errors' in error) { + const errors = error.errors + if (Array.isArray(errors)) { + const first = errors[0] + if (first && typeof first === 'object') { + const message = + 'message' in first && typeof first.message === 'string' + ? first.message + : 'Interrupt persistence failed.' + const code = + 'code' in first && typeof first.code === 'string' + ? first.code + : 'server' + return { message, code } + } + } + } + return { + message: + error instanceof Error + ? error.message + : 'Interrupt persistence failed.', + code: 'server', + } + } + + private buildInterruptRunErrorChunk(error: unknown): StreamChunk { + const failure = this.interruptFailure(error) + return { + type: EventType.RUN_ERROR, + timestamp: Date.now(), + runId: this.runIdOverride ?? this.requestId, + threadId: this.threadId, + message: failure.message, + code: failure.code, + error: { message: failure.message, code: failure.code }, + ...(failure.errors !== undefined + ? { 'tanstack:interruptErrors': failure.errors } + : {}), + ...(failure.recovery !== undefined + ? { 'tanstack:interruptRecovery': failure.recovery } + : {}), + } + } + + private async *emitInterruptRunError( + error: unknown, + ): AsyncGenerator { + const failure = this.interruptFailure(error) + this.finalizationError = { + message: failure.message, + code: failure.code, + cause: error, + } + yield* this.pipeThroughMiddleware(this.buildInterruptRunErrorChunk(error)) + } + + private async *emitActionableInterruptBoundary( + finishEvent: RunFinishedEvent, + approvals: Array, + clientRequests: Array, + ): AsyncGenerator { + if (!getInterruptPersistence(this.middlewareCtx, { optional: true })) { + yield* this.emitInterruptRunError({ + errors: [ + { + code: 'persistence-required', + message: + 'Actionable interrupts require an atomic interrupt persistence capability.', + }, + ], + }) + return false + } + + const terminal = this.buildInterruptFinishedChunk( + finishEvent, + approvals, + clientRequests, + ) + let terminalOutputs: Array + try { + terminalOutputs = await this.middlewareRunner.runOnChunk( + this.middlewareCtx, + terminal, + ) + } catch (error) { + yield* this.emitInterruptRunError(error) + return false + } + + yield* this.pipeThroughMiddleware(this.buildMessagesSnapshotChunk()) + if (this.params.state !== undefined) { + yield* this.pipeThroughMiddleware({ + type: EventType.STATE_SNAPSHOT, + timestamp: Date.now(), + model: this.params.model, + snapshot: this.params.state, + }) + } + for (const output of terminalOutputs) { + yield this.publicInterruptTerminal(output) + this.middlewareCtx.chunkIndex++ + } + return true + } + private buildToolResultChunks( results: Array, finishEvent: RunFinishedEvent, - argsMap?: Map, ): Array { const chunks: Array = [] @@ -1834,41 +2275,6 @@ class TextEngine< const wireContent = typeof content === 'string' ? content : JSON.stringify(content) - // argsMap is set only on continuation re-executions, where the adapter - // never streamed these calls. Otherwise it already emitted END, so a - // second one here would be an orphan that fails verifyEvents (#519). - if (argsMap) { - chunks.push({ - type: 'TOOL_CALL_START', - timestamp: Date.now(), - model: finishEvent.model, - toolCallId: result.toolCallId, - toolCallName: result.toolName, - toolName: result.toolName, - } as StreamChunk) - - const args = argsMap.get(result.toolCallId) ?? '{}' - chunks.push({ - type: 'TOOL_CALL_ARGS', - timestamp: Date.now(), - model: finishEvent.model, - toolCallId: result.toolCallId, - delta: args, - args, - } as StreamChunk) - - chunks.push({ - type: 'TOOL_CALL_END', - timestamp: Date.now(), - model: finishEvent.model, - toolCallId: result.toolCallId, - toolCallName: result.toolName, - toolName: result.toolName, - result: wireContent, - ...(result.state !== undefined && { state: result.state }), - } as StreamChunk) - } - // AG-UI spec TOOL_CALL_RESULT event (content is string-only per spec) chunks.push({ type: 'TOOL_CALL_RESULT', @@ -2564,6 +2970,8 @@ class TextEngine< resumeToolState: { approvals: this.resumeApprovals, clientToolResults: this.resumeClientToolResults, + deniedToolResults: this.resumeDeniedToolResults, + cancelledToolCallIds: this.resumeCancelledToolCallIds, }, metadata: this.params.metadata, modelOptions: this.params.modelOptions, @@ -2572,8 +2980,8 @@ class TextEngine< private applyMiddlewareConfig(config: ChatMiddlewareConfig): void { if (config.resumeToolState?.approvals) { - for (const [approvalId, approved] of config.resumeToolState.approvals) { - this.resumeApprovals.set(approvalId, approved) + for (const [approvalId, resolution] of config.resumeToolState.approvals) { + this.resumeApprovals.set(approvalId, resolution) } } if (config.resumeToolState?.clientToolResults) { @@ -2582,6 +2990,17 @@ class TextEngine< this.resumeClientToolResults.set(toolCallId, result) } } + if (config.resumeToolState?.deniedToolResults) { + for (const [toolCallId, result] of config.resumeToolState + .deniedToolResults) { + this.resumeDeniedToolResults.set(toolCallId, result) + } + } + if (config.resumeToolState?.cancelledToolCallIds) { + for (const toolCallId of config.resumeToolState.cancelledToolCallIds) { + this.resumeCancelledToolCallIds.add(toolCallId) + } + } this.messages = config.messages this.systemPrompts = config.systemPrompts this.tools = config.tools diff --git a/packages/ai/src/activities/chat/mcp/manager.ts b/packages/ai/src/activities/chat/mcp/manager.ts index c52bff037..bac5e22ec 100644 --- a/packages/ai/src/activities/chat/mcp/manager.ts +++ b/packages/ai/src/activities/chat/mcp/manager.ts @@ -1,4 +1,4 @@ -import type { ServerTool } from '../tools/tool-definition' +import type { AnyServerTool } from '../tools/tool-definition' import type { ChatMCPOptions, MCPToolSource } from './types' /** @@ -18,7 +18,7 @@ import type { ChatMCPOptions, MCPToolSource } from './types' * drains. If discovery results were ever cached and reused across runs, this * would bind a closure over an already-closed source — bind onto a copy then. */ -function bindReadResource(tool: ServerTool, source: MCPToolSource): void { +function bindReadResource(tool: AnyServerTool, source: MCPToolSource): void { if (!source.readResource) return const meta = ( tool.metadata as { mcp?: { uiResourceUri?: string } } | undefined @@ -71,13 +71,13 @@ export class MCPManager { * (no `onDiscoveryError`, or it re-threw) or a duplicate tool name; in that * case it first closes any connected sources when the policy is 'close'. */ - async discover(): Promise> { + async discover(): Promise> { if (this.#sources.length === 0) return [] try { const settled = await Promise.allSettled( this.#sources.map((s) => s.tools({ lazy: this.#lazyTools })), ) - const tools: Array = [] + const tools: Array = [] const zipped = this.#sources.map( (source, i) => [source, settled[i]] as const, ) diff --git a/packages/ai/src/activities/chat/mcp/types.ts b/packages/ai/src/activities/chat/mcp/types.ts index 6fe73d5ef..96c7911ca 100644 --- a/packages/ai/src/activities/chat/mcp/types.ts +++ b/packages/ai/src/activities/chat/mcp/types.ts @@ -1,4 +1,4 @@ -import type { ServerTool } from '../tools/tool-definition' +import type { AnyServerTool } from '../tools/tool-definition' /** * The shape `readResource` resolves to — a structural subset of MCP's @@ -26,7 +26,7 @@ export interface MCPToolSource { // Keep the options shape in sync with ai-mcp's `ToolsOptions` — extra // optional fields added there still match structurally, but chat() only // forwards what is declared here. - tools: (options?: { lazy?: boolean }) => Promise> + tools: (options?: { lazy?: boolean }) => Promise> close: () => Promise /** * Reads an MCP resource by URI. Used by the chat manager to eagerly fetch diff --git a/packages/ai/src/activities/chat/middleware/index.ts b/packages/ai/src/activities/chat/middleware/index.ts index 6d231eedd..ad4fa9dad 100644 --- a/packages/ai/src/activities/chat/middleware/index.ts +++ b/packages/ai/src/activities/chat/middleware/index.ts @@ -4,6 +4,7 @@ export type { ChatMiddlewarePhase, ChatMiddlewareConfig, ChatResumeToolState, + ChatResumeGenericResolution, StructuredOutputMiddlewareConfig, ToolCallHookContext, BeforeToolCallDecision, diff --git a/packages/ai/src/activities/chat/middleware/types.ts b/packages/ai/src/activities/chat/middleware/types.ts index 741d2df45..06798374f 100644 --- a/packages/ai/src/activities/chat/middleware/types.ts +++ b/packages/ai/src/activities/chat/middleware/types.ts @@ -8,6 +8,7 @@ import type { ToolCall, } from '../../../types' import type { SystemPrompt } from '../../../system-prompts' +import type { ToolApprovalResolution } from '../../../interrupts' import type { Capability, CapabilityHandle, @@ -91,6 +92,8 @@ export interface ChatMiddlewareContext { streamId: string /** AG-UI run identifier for correlating client and server events */ runId: string + /** Interrupted or parent run correlated with this continuation. */ + parentRunId?: string /** * AG-UI thread identifier — a stable per-conversation ID used to * correlate client and server devtools events. Resolves to the @@ -221,10 +224,19 @@ export interface ChatMiddlewareConfig { * without relying on client message history. */ export interface ChatResumeToolState { - approvals?: ReadonlyMap | undefined + approvals?: ReadonlyMap | undefined clientToolResults?: ReadonlyMap | undefined + genericInterrupts?: + | ReadonlyMap + | undefined + deniedToolResults?: ReadonlyMap | undefined + cancelledToolCallIds?: ReadonlySet | undefined } +export type ChatResumeGenericResolution = + | { interruptId: string; status: 'resolved'; payload: unknown } + | { interruptId: string; status: 'cancelled'; payload?: never } + /** * Config passed to onStructuredOutputConfig. * diff --git a/packages/ai/src/activities/chat/tools/approval-schema.ts b/packages/ai/src/activities/chat/tools/approval-schema.ts new file mode 100644 index 000000000..7cfb47103 --- /dev/null +++ b/packages/ai/src/activities/chat/tools/approval-schema.ts @@ -0,0 +1,205 @@ +import { + canonicalInterruptJson, + digestInterruptJson, +} from '../../../interrupt-serialization' +import { compileJsonSchema202012 } from './json-schema-validator' +import { isStandardJSONSchema, isStandardSchema } from './schema-converter' +import type { StandardSchemaV1 } from '@standard-schema/spec' +import type { JSONSchema, SchemaInput } from '../../../types' +import type { ApprovalSchemaConfig } from './tool-definition' + +export interface NormalizedSchemaInput { + source: SchemaInput + validator?: StandardSchemaV1 + jsonSchema?: JSONSchema +} + +export interface NormalizedApprovalSchema { + branches: { + approve: NormalizedSchemaInput | null + reject: NormalizedSchemaInput | null + } + responseSchema: JSONSchema + responseSchemaHash: string + approvalSchemaHash: string +} + +const jsonSchemaKeywords = new Set([ + '$schema', + '$id', + '$ref', + '$defs', + 'type', + 'properties', + 'required', + 'additionalProperties', + 'items', + 'oneOf', + 'anyOf', + 'allOf', + 'enum', + 'const', + 'format', + 'minimum', + 'maximum', + 'minLength', + 'maxLength', + 'pattern', +]) + +function isPlainRecord(value: unknown): value is Record { + return ( + value !== null && + typeof value === 'object' && + !Array.isArray(value) && + (Object.getPrototypeOf(value) === Object.prototype || + Object.getPrototypeOf(value) === null) + ) +} + +function isRawJsonSchema(value: unknown): value is JSONSchema { + return ( + isPlainRecord(value) && + Object.keys(value).some((key) => jsonSchemaKeywords.has(key)) + ) +} + +function isSchemaInput(value: unknown): value is SchemaInput { + return ( + isStandardSchema(value) || + isStandardJSONSchema(value) || + isRawJsonSchema(value) + ) +} + +function isApprovalBranchMap( + value: unknown, +): value is { approve?: SchemaInput; reject?: SchemaInput } { + if (!isPlainRecord(value)) return false + const keys = Object.keys(value) + return ( + keys.length > 0 && + keys.every((key) => key === 'approve' || key === 'reject') && + keys.every((key) => isSchemaInput(value[key])) + ) +} + +function toJsonSchema(value: Record): JSONSchema { + const result: JSONSchema = {} + for (const [key, item] of Object.entries(value)) { + result[key] = item + } + return result +} + +function schemaToWire(schema: SchemaInput): NormalizedSchemaInput { + if (isStandardSchema(schema)) { + const jsonSchema = isStandardJSONSchema(schema) + ? toJsonSchema( + schema['~standard'].jsonSchema.input({ + target: 'draft-2020-12', + }), + ) + : undefined + return { + source: schema, + validator: schema, + ...(jsonSchema !== undefined && { jsonSchema }), + } + } + if (isRawJsonSchema(schema)) { + compileJsonSchema202012(schema) + return { source: schema, jsonSchema: schema } + } + throw new TypeError('Expected a supported SchemaInput.') +} + +function decisionEnvelope(input: { + approved: boolean + payload: NormalizedSchemaInput | null + inputSchema: NormalizedSchemaInput | null +}): JSONSchema { + const properties: Record = { + approved: { const: input.approved }, + } + const required = ['approved'] + if (input.approved && input.inputSchema) { + properties['editedArgs'] = input.inputSchema.jsonSchema ?? {} + } + if (input.payload) { + properties['payload'] = input.payload.jsonSchema ?? {} + required.push('payload') + } + return { + type: 'object', + properties, + required, + additionalProperties: false, + } +} + +export function normalizeApprovalSchema( + approvalSchema: ApprovalSchemaConfig | undefined, + inputSchema?: SchemaInput, +): NormalizedApprovalSchema { + const normalizedInput = + inputSchema === undefined ? null : schemaToWire(inputSchema) + let approve: NormalizedSchemaInput | null = null + let reject: NormalizedSchemaInput | null = null + + if (approvalSchema !== undefined) { + if (isStandardSchema(approvalSchema) || isRawJsonSchema(approvalSchema)) { + approve = schemaToWire(approvalSchema) + reject = approve + } else if (isApprovalBranchMap(approvalSchema)) { + approve = + approvalSchema.approve === undefined + ? null + : schemaToWire(approvalSchema.approve) + reject = + approvalSchema.reject === undefined + ? null + : schemaToWire(approvalSchema.reject) + } else { + throw new TypeError( + 'approvalSchema must be a SchemaInput or a nonempty map containing approve or reject.', + ) + } + } + + const responseSchema: JSONSchema = { + oneOf: [ + decisionEnvelope({ + approved: true, + payload: approve, + inputSchema: normalizedInput, + }), + decisionEnvelope({ + approved: false, + payload: reject, + inputSchema: null, + }), + ], + } + const responseCanonical = canonicalInterruptJson(responseSchema) + const approvalCanonical = canonicalInterruptJson({ + approve: approve?.jsonSchema ?? null, + reject: reject?.jsonSchema ?? null, + }) + + return { + branches: { approve, reject }, + responseSchema, + responseSchemaHash: digestInterruptJson(responseCanonical), + approvalSchemaHash: digestInterruptJson(approvalCanonical), + } +} + +export function hashSchemaInput(schema: SchemaInput | undefined): string { + if (schema === undefined) return digestInterruptJson('undefined') + const normalized = schemaToWire(schema) + const identity = normalized.jsonSchema ?? { + standardValidator: 'unserialized', + } + return digestInterruptJson(canonicalInterruptJson(identity)) +} diff --git a/packages/ai/src/activities/chat/tools/json-schema-validator.ts b/packages/ai/src/activities/chat/tools/json-schema-validator.ts new file mode 100644 index 000000000..753a89873 --- /dev/null +++ b/packages/ai/src/activities/chat/tools/json-schema-validator.ts @@ -0,0 +1,138 @@ +import Ajv2020 from 'ajv/dist/2020.js' +import addFormats from 'ajv-formats' +import type { ErrorObject } from 'ajv' +import type { JSONSchema } from '../../../types' + +export interface JsonSchemaValidationIssue { + keyword: string + message: string + path: ReadonlyArray +} + +export class JsonSchemaCompilationError extends Error { + constructor(message: string, options?: ErrorOptions) { + super(message, options) + this.name = 'JsonSchemaCompilationError' + } +} + +const draft202012 = 'https://json-schema.org/draft/2020-12/schema' + +function decodePointer(pointer: string): Array { + if (pointer === '') return [] + return pointer + .slice(1) + .split('/') + .map((part) => part.replaceAll('~1', '/').replaceAll('~0', '~')) +} + +function issuePath(error: ErrorObject): Array { + const path = decodePointer(error.instancePath) + if (error.keyword === 'required') { + path.push(String(error.params['missingProperty'])) + } + if (error.keyword === 'additionalProperties') { + path.push(String(error.params['additionalProperty'])) + } + return path +} + +function assertSupportedSchemaTree( + schema: unknown, +): asserts schema is JSONSchema { + const active = new WeakSet() + + const visit = (value: unknown): void => { + if ( + value === null || + typeof value === 'string' || + typeof value === 'boolean' + ) { + return + } + if (typeof value === 'number') { + if (Number.isFinite(value)) return + throw new JsonSchemaCompilationError( + 'Schema values must be JSON-compatible.', + ) + } + if (typeof value !== 'object') { + throw new JsonSchemaCompilationError( + 'Schema values must be JSON-compatible.', + ) + } + if ( + !Array.isArray(value) && + Object.getPrototypeOf(value) !== Object.prototype && + Object.getPrototypeOf(value) !== null + ) { + throw new JsonSchemaCompilationError( + 'Schema nodes must be plain JSON objects.', + ) + } + if (active.has(value)) { + throw new JsonSchemaCompilationError( + 'Schema values must not contain cycles.', + ) + } + + active.add(value) + if (Array.isArray(value)) { + value.forEach(visit) + active.delete(value) + return + } + + const record = value as Record + if ( + typeof record['$schema'] === 'string' && + record['$schema'] !== draft202012 + ) { + throw new JsonSchemaCompilationError('Only Draft 2020-12 is supported.') + } + if (typeof record['$ref'] === 'string' && !record['$ref'].startsWith('#')) { + throw new JsonSchemaCompilationError( + 'Only document-local $ref values are supported.', + ) + } + Object.values(record).forEach(visit) + active.delete(value) + } + + visit(schema) +} + +export function compileJsonSchema202012( + schema: unknown, +): (value: unknown) => ReadonlyArray { + assertSupportedSchemaTree(schema) + const ajv = new Ajv2020({ + allErrors: true, + strict: true, + strictRequired: false, + allowUnionTypes: true, + validateFormats: true, + coerceTypes: false, + useDefaults: false, + removeAdditional: false, + }) + addFormats(ajv, { mode: 'full' }) + + let validate + try { + validate = ajv.compile(schema) + } catch (cause) { + throw new JsonSchemaCompilationError('Invalid Draft 2020-12 schema.', { + cause, + }) + } + + return (value) => { + if (validate(value)) return [] + return (validate.errors ?? []).map((error) => ({ + keyword: error.keyword, + message: error.message ?? 'Schema validation failed.', + path: issuePath(error), + })) + } +} diff --git a/packages/ai/src/activities/chat/tools/tool-calls.ts b/packages/ai/src/activities/chat/tools/tool-calls.ts index 5049c1de0..fc5cae2f9 100644 --- a/packages/ai/src/activities/chat/tools/tool-calls.ts +++ b/packages/ai/src/activities/chat/tools/tool-calls.ts @@ -1,5 +1,6 @@ import { normalizeToolResult } from '../../../utilities/tool-result' import { isStandardSchema, parseWithStandardSchema } from './schema-converter' +import type { ToolApprovalResolution } from '../../../interrupts' import type { AnyTool, ContentPart, @@ -435,6 +436,36 @@ export interface ClientToolRequest { input: any } +export interface ToolResumeExecutionState { + deniedToolResults?: ReadonlyMap + cancelledToolCallIds?: ReadonlySet +} + +function approvalResolution( + approvals: ReadonlyMap, + toolCallId: string, +): ToolApprovalResolution | undefined { + return approvals.get(toolCallId) ?? approvals.get(`approval_${toolCallId}`) +} + +function isApproved(resolution: ToolApprovalResolution): boolean { + return typeof resolution === 'boolean' ? resolution : resolution.approved +} + +function editedApprovalArgs( + resolution: ToolApprovalResolution, +): unknown | undefined { + return typeof resolution === 'object' && resolution.approved + ? resolution.editedArgs + : undefined +} + +function deniedApprovalResult(resolution: ToolApprovalResolution): unknown { + return typeof resolution === 'object' && !resolution.approved + ? (resolution.payload ?? { error: 'User declined tool execution' }) + : { error: 'User declined tool execution' } +} + interface ExecuteToolCallsResult { /** Tool results ready to send to LLM */ results: Array @@ -685,7 +716,7 @@ function buildClientToolResult( export async function* executeToolCalls( toolCalls: Array, tools: ReadonlyArray, - approvals: Map = new Map(), + approvals: Map = new Map(), clientResults: Map = new Map(), createCustomEventChunk?: ( eventName: string, @@ -694,6 +725,7 @@ export async function* executeToolCalls( middlewareHooks?: ToolExecutionMiddlewareHooks, userContext?: TContext, abortSignal?: AbortSignal, + resumeState?: ToolResumeExecutionState, ): AsyncGenerator { const results: Array = [] const needsApproval: Array = [] @@ -709,7 +741,9 @@ export async function* executeToolCalls( // defer all execution so side effects don't happen before the user decides. const hasPendingApprovals = toolCalls.some((tc) => { const t = toolMap.get(tc.function.name) - return t?.needsApproval && !approvals.has(`approval_${tc.id}`) + return ( + t?.needsApproval && approvalResolution(approvals, tc.id) === undefined + ) }) for (const toolCall of toolCalls) { @@ -729,11 +763,24 @@ export async function* executeToolCalls( // Skip non-pending tools while approvals are outstanding if (hasPendingApprovals) { - if (!tool.needsApproval || approvals.has(`approval_${toolCall.id}`)) { + if ( + !tool.needsApproval || + approvalResolution(approvals, toolCall.id) !== undefined + ) { continue } } + if (resumeState?.cancelledToolCallIds?.has(toolCall.id)) { + results.push({ + toolCallId: toolCall.id, + toolName, + result: { error: 'Tool execution cancelled' }, + state: 'output-error', + }) + continue + } + // Parse arguments, throwing error if invalid JSON let input: unknown = {} const argsStr = toolCall.function.arguments.trim() || '{}' @@ -792,12 +839,14 @@ export async function* executeToolCalls( // Check if tool needs approval if (tool.needsApproval) { const approvalId = `approval_${toolCall.id}` + const resolution = approvalResolution(approvals, toolCall.id) // Check if approval decision exists - if (approvals.has(approvalId)) { - const approved = approvals.get(approvalId) + if (resolution !== undefined) { + const approved = isApproved(resolution) if (approved) { + input = editedApprovalArgs(resolution) ?? input // Approved - check if client has executed if (clientResults.has(toolCall.id)) { results.push( @@ -821,7 +870,9 @@ export async function* executeToolCalls( results.push({ toolCallId: toolCall.id, toolName, - result: { error: 'User declined tool execution' }, + result: + resumeState?.deniedToolResults?.get(toolCall.id) ?? + deniedApprovalResult(resolution), state: 'output-error', }) } @@ -860,12 +911,14 @@ export async function* executeToolCalls( // CASE 2: Server tool with approval required if (tool.needsApproval) { const approvalId = `approval_${toolCall.id}` + const resolution = approvalResolution(approvals, toolCall.id) // Check if approval decision exists - if (approvals.has(approvalId)) { - const approved = approvals.get(approvalId) + if (resolution !== undefined) { + const approved = isApproved(resolution) if (approved) { + input = editedApprovalArgs(resolution) ?? input // Apply middleware before-hook for approved tools if (middlewareHooks) { const decision = await applyBeforeToolCallDecision( @@ -895,7 +948,9 @@ export async function* executeToolCalls( results.push({ toolCallId: toolCall.id, toolName, - result: { error: 'User declined tool execution' }, + result: + resumeState?.deniedToolResults?.get(toolCall.id) ?? + deniedApprovalResult(resolution), state: 'output-error', }) } diff --git a/packages/ai/src/activities/chat/tools/tool-definition.ts b/packages/ai/src/activities/chat/tools/tool-definition.ts index fe62bd4ac..2849661eb 100644 --- a/packages/ai/src/activities/chat/tools/tool-definition.ts +++ b/packages/ai/src/activities/chat/tools/tool-definition.ts @@ -1,36 +1,113 @@ -import type { StandardJSONSchemaV1 } from '@standard-schema/spec' import type { + InferSchemaType, JSONSchema, SchemaInput, Tool, ToolExecuteFunction, } from '../../../types' +import type { + StandardJSONSchemaV1, + StandardSchemaV1, +} from '@standard-schema/spec' + +export declare const toolApprovalCapability: unique symbol + +export interface ToolApprovalCapabilityMarker< + TNeedsApproval extends boolean, + TApprovalSchema, +> { + readonly [toolApprovalCapability]?: { + needsApproval: TNeedsApproval + approvalSchema: TApprovalSchema + } +} + +export type ApprovalSchemaConfig = + | SchemaInput + | { approve: SchemaInput; reject?: SchemaInput } + | { approve?: SchemaInput; reject: SchemaInput } + +type ApprovalConfig< + TNeedsApproval extends boolean, + TApprovalSchema extends ApprovalSchemaConfig | undefined, +> = TNeedsApproval extends true + ? { needsApproval: TNeedsApproval; approvalSchema?: TApprovalSchema } + : { needsApproval?: TNeedsApproval; approvalSchema?: never } + +export type ApprovalCapabilityOf = + TTool extends ToolApprovalCapabilityMarker + ? TNeeds + : false + +export type ApprovalSchemaOf = + TTool extends ToolApprovalCapabilityMarker + ? TSchema + : undefined + +export declare const noSchema: unique symbol +export type NoSchema = typeof noSchema + +export type InputSchemaOf = TTool extends { + inputSchema: infer TInput +} + ? TInput extends undefined + ? NoSchema + : TInput + : NoSchema + +export type OutputSchemaOf = TTool extends { + outputSchema: infer TOutput +} + ? TOutput extends undefined + ? NoSchema + : TOutput + : NoSchema + +type BuiltToolSchemaFields< + TInput extends SchemaInput | undefined, + TOutput extends SchemaInput | undefined, + TApprovalSchema extends ApprovalSchemaConfig | undefined, +> = { + inputSchema: TInput + outputSchema: TOutput + approvalSchema: TApprovalSchema +} /** * Marker type for server-side tools */ export interface ServerTool< - TInput extends SchemaInput = SchemaInput, - TOutput extends SchemaInput = SchemaInput, + TInput extends SchemaInput | undefined = undefined, + TOutput extends SchemaInput | undefined = undefined, TName extends string = string, TContext = unknown, -> extends Tool { + TNeedsApproval extends boolean = false, + TApprovalSchema extends ApprovalSchemaConfig | undefined = undefined, +> + extends + Tool, + ToolApprovalCapabilityMarker { __toolSide: 'server' + inputSchema?: TInput + outputSchema?: TOutput + needsApproval?: TNeedsApproval + approvalSchema?: TApprovalSchema } /** * Marker type for client-side tools */ export interface ClientTool< - TInput extends SchemaInput = SchemaInput, - TOutput extends SchemaInput = SchemaInput, + TInput extends SchemaInput | undefined = undefined, + TOutput extends SchemaInput | undefined = undefined, TName extends string = string, TContext = unknown, // Captured as a literal (`true` / `false`) so downstream types — notably // the tool-call part's `approval` field — can be gated on it. Defaults to // `false` when the tool config omits `needsApproval`. TNeedsApproval extends boolean = false, -> { + TApprovalSchema extends ApprovalSchemaConfig | undefined = undefined, +> extends ToolApprovalCapabilityMarker { __toolSide: 'client' name: TName description: string @@ -42,35 +119,55 @@ export interface ClientTool< inputSchema?: TInput outputSchema?: TOutput needsApproval?: TNeedsApproval + approvalSchema?: TApprovalSchema lazy?: boolean metadata?: Record execute?: ToolExecuteFunction } +/** Broad server-tool shape for heterogeneous internal collections. */ +export type AnyServerTool = Omit< + ServerTool, + 'execute' +> & { + execute?: ((args: any, context?: any) => any) | undefined +} + /** * Tool definition that can be used directly or instantiated for server/client */ export interface ToolDefinitionInstance< - TInput extends SchemaInput = SchemaInput, - TOutput extends SchemaInput = SchemaInput, + TInput extends SchemaInput | undefined = undefined, + TOutput extends SchemaInput | undefined = undefined, TName extends string = string, TContext = unknown, TNeedsApproval extends boolean = false, + TApprovalSchema extends ApprovalSchemaConfig | undefined = undefined, > extends Tool { __toolSide: 'definition' // Narrow the base `needsApproval?: boolean` to the captured literal so it // survives into `ToolCallPartForTool`'s approval gate. + inputSchema: TInput + outputSchema: TOutput needsApproval?: TNeedsApproval + approvalSchema: TApprovalSchema + readonly [toolApprovalCapability]?: { + needsApproval: TNeedsApproval + approvalSchema: TApprovalSchema + } } /** * Union type for any kind of client-side tool (client tool or definition) */ export type AnyClientTool = - | (Omit, 'execute'> & { + | (Omit, 'execute'> & { execute?: ((args: any, context?: any) => any) | undefined }) - | (Omit, 'execute'> & { + | (Omit< + ToolDefinitionInstance, + 'execute' + > & { execute?: ((args: any, context?: any) => any) | undefined }) @@ -83,63 +180,73 @@ export type InferToolName = T extends { name: infer N } ? N : never * Extract the input type from a tool (inferred from Standard JSON Schema, or `unknown` for plain JSONSchema) */ export type InferToolInput = T extends { inputSchema?: infer TInput } - ? TInput extends StandardJSONSchemaV1 - ? TInferred - : TInput extends JSONSchema - ? unknown - : unknown + ? TInput extends JSONSchema + ? unknown + : InferSchemaType : unknown /** * Extract the output type from a tool (inferred from Standard JSON Schema, or `unknown` for plain JSONSchema) */ export type InferToolOutput = T extends { outputSchema?: infer TOutput } - ? TOutput extends StandardJSONSchemaV1 - ? TInferred - : TOutput extends JSONSchema - ? unknown - : unknown + ? TOutput extends StandardJSONSchemaV1 + ? InferSchemaType + : TOutput extends StandardSchemaV1 + ? InferSchemaType + : TOutput extends JSONSchema + ? unknown + : InferSchemaType : unknown /** * Tool definition configuration */ -export interface ToolDefinitionConfig< - TInput extends SchemaInput = SchemaInput, - TOutput extends SchemaInput = SchemaInput, +export type ToolDefinitionConfig< + TInput extends SchemaInput | undefined = undefined, + TOutput extends SchemaInput | undefined = undefined, TName extends string = string, TNeedsApproval extends boolean = false, -> { + TApprovalSchema extends ApprovalSchemaConfig | undefined = undefined, +> = { name: TName description: string inputSchema?: TInput outputSchema?: TOutput - needsApproval?: TNeedsApproval lazy?: boolean metadata?: Record -} +} & ApprovalConfig /** * Tool definition builder that allows creating server or client tools from a shared definition */ export interface ToolDefinition< - TInput extends SchemaInput = SchemaInput, - TOutput extends SchemaInput = SchemaInput, + TInput extends SchemaInput | undefined = undefined, + TOutput extends SchemaInput | undefined = undefined, TName extends string = string, TNeedsApproval extends boolean = false, + TApprovalSchema extends ApprovalSchemaConfig | undefined = undefined, > extends ToolDefinitionInstance< TInput, TOutput, TName, unknown, - TNeedsApproval + TNeedsApproval, + TApprovalSchema > { /** * Create a server-side tool with execute function */ server: ( execute: ToolExecuteFunction, - ) => ServerTool + ) => ServerTool< + TInput, + TOutput, + TName, + TContext, + TNeedsApproval, + TApprovalSchema + > & + BuiltToolSchemaFields /** * Create a client-side tool with optional execute function. @@ -148,7 +255,15 @@ export interface ToolDefinition< */ client: ( execute?: ToolExecuteFunction, - ) => ClientTool + ) => ClientTool< + TInput, + TOutput, + TName, + TContext, + TNeedsApproval, + TApprovalSchema + > & + BuiltToolSchemaFields } /** @@ -207,35 +322,84 @@ export interface ToolDefinition< * ``` */ export function toolDefinition< - TInput extends SchemaInput = SchemaInput, - TOutput extends SchemaInput = SchemaInput, + TInput extends SchemaInput | undefined = undefined, + TOutput extends SchemaInput | undefined = undefined, TName extends string = string, // `const` forces the literal (`true` / `false`) to be captured from the // config's optional `needsApproval` — without it TS widens to `boolean`, // which collapses the approval gate in `ToolCallPartForTool`. const TNeedsApproval extends boolean = false, + TApprovalSchema extends ApprovalSchemaConfig | undefined = undefined, >( - config: ToolDefinitionConfig, -): ToolDefinition { - const definition: ToolDefinition = { + config: ToolDefinitionConfig< + TInput, + TOutput, + TName, + TNeedsApproval, + TApprovalSchema + >, +): ToolDefinition { + if (config.approvalSchema !== undefined && config.needsApproval !== true) { + throw new TypeError('approvalSchema requires needsApproval: true.') + } + const inputSchema = config.inputSchema as TInput + const outputSchema = config.outputSchema as TOutput + const approvalSchema = config.approvalSchema as TApprovalSchema + const needsApproval = config.needsApproval as TNeedsApproval | undefined + + const definition: ToolDefinition< + TInput, + TOutput, + TName, + TNeedsApproval, + TApprovalSchema + > = { __toolSide: 'definition', ...config, + inputSchema, + outputSchema, + approvalSchema, + needsApproval, server( execute: ToolExecuteFunction, - ): ServerTool { + ): ServerTool< + TInput, + TOutput, + TName, + TContext, + TNeedsApproval, + TApprovalSchema + > & + BuiltToolSchemaFields { return { __toolSide: 'server', ...config, + inputSchema, + outputSchema, + approvalSchema, + needsApproval, execute, } }, client( execute?: ToolExecuteFunction, - ): ClientTool { + ): ClientTool< + TInput, + TOutput, + TName, + TContext, + TNeedsApproval, + TApprovalSchema + > & + BuiltToolSchemaFields { return { __toolSide: 'client', ...config, + inputSchema, + outputSchema, + approvalSchema, + needsApproval, ...(execute !== undefined && { execute }), } }, diff --git a/packages/ai/src/client.ts b/packages/ai/src/client.ts index 1a9d8e623..e9b84671f 100644 --- a/packages/ai/src/client.ts +++ b/packages/ai/src/client.ts @@ -232,17 +232,43 @@ export { type InferToolInput, type InferToolName, type InferToolOutput, + type ApprovalCapabilityOf, + type ApprovalSchemaConfig, + type ApprovalSchemaOf, + type InputSchemaOf, + type OutputSchemaOf, + type NoSchema, type ToolDefinition, type ToolDefinitionConfig, type ToolDefinitionInstance, } from './activities/chat/tools/tool-definition' +export { + hashSchemaInput, + normalizeApprovalSchema, + type NormalizedApprovalSchema, + type NormalizedSchemaInput, +} from './activities/chat/tools/approval-schema' + +export { + canonicalInterruptJson, + cloneAndDeepFreezeJson, + digestInterruptJson, +} from './interrupt-serialization' + export { convertSchemaToJsonSchema, isStandardSchema, parseWithStandardSchema, + validateWithStandardSchema, } from './activities/chat/tools/schema-converter' +export { + JsonSchemaCompilationError, + compileJsonSchema202012, + type JsonSchemaValidationIssue, +} from './activities/chat/tools/json-schema-validator' + export { convertMessagesToModelMessages, generateMessageId, @@ -316,6 +342,8 @@ export type { InferSchemaType, } from './types' +export * from './interrupts' + export type { AudioVisualization, RealtimeAdapter, diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 5558029a2..1bce2baad 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -45,12 +45,30 @@ export { type ToolDefinitionInstance, type ToolDefinitionConfig, type ServerTool, + type AnyServerTool, type ClientTool, type AnyClientTool, type InferToolName, type InferToolInput, type InferToolOutput, + type ApprovalCapabilityOf, + type ApprovalSchemaConfig, + type ApprovalSchemaOf, + type InputSchemaOf, + type OutputSchemaOf, + type NoSchema, } from './activities/chat/tools/tool-definition' +export { + hashSchemaInput, + normalizeApprovalSchema, + type NormalizedApprovalSchema, + type NormalizedSchemaInput, +} from './activities/chat/tools/approval-schema' +export { + canonicalInterruptJson, + cloneAndDeepFreezeJson, + digestInterruptJson, +} from './interrupt-serialization' // MCP chat option types export type { @@ -67,8 +85,14 @@ export { convertSchemaToJsonSchema, isStandardSchema, parseWithStandardSchema, + validateWithStandardSchema, StandardSchemaValidationError, } from './activities/chat/tools/schema-converter' +export { + JsonSchemaCompilationError, + compileJsonSchema202012, + type JsonSchemaValidationIssue, +} from './activities/chat/tools/json-schema-validator' // Stream utilities export { @@ -115,6 +139,7 @@ export type { ChatMiddlewarePhase, ChatMiddlewareConfig, ChatResumeToolState, + ChatResumeGenericResolution, StructuredOutputMiddlewareConfig, ToolCallHookContext, BeforeToolCallDecision, @@ -130,6 +155,8 @@ export type { ChatSandboxHooks, } from './activities/chat/middleware/index' +export * from './interrupts' + // Base, activity-agnostic middleware. The observe-only superset that media // activities accept via their `middleware` option; `ChatMiddleware` adds the // chat-only hooks on top. Pure types only — the `otelMiddleware` value lives at diff --git a/packages/ai/src/interrupt-serialization.ts b/packages/ai/src/interrupt-serialization.ts new file mode 100644 index 000000000..2441c2b6c --- /dev/null +++ b/packages/ai/src/interrupt-serialization.ts @@ -0,0 +1,68 @@ +import { sha256 } from '@noble/hashes/sha2.js' +import { bytesToHex, utf8ToBytes } from '@noble/hashes/utils.js' + +function canonical(value: unknown, active: WeakSet): string { + if (value === null) return 'null' + if (typeof value === 'string' || typeof value === 'boolean') { + return JSON.stringify(value) + } + if (typeof value === 'number' && Number.isFinite(value)) { + return JSON.stringify(value) + } + if (typeof value !== 'object') { + throw new TypeError('Interrupt values must be JSON-compatible.') + } + if (active.has(value)) { + throw new TypeError('Interrupt values must not cycle.') + } + if ( + !Array.isArray(value) && + Object.getPrototypeOf(value) !== Object.prototype && + Object.getPrototypeOf(value) !== null + ) { + throw new TypeError('Interrupt values must use plain JSON objects.') + } + + active.add(value) + let encoded: string + if (Array.isArray(value)) { + const items: Array = [] + for (let index = 0; index < value.length; index++) { + items.push(canonical(value[index], active)) + } + encoded = `[${items.join(',')}]` + } else { + const record = value as Record + encoded = `{${Object.keys(record) + .sort() + .map((key) => `${JSON.stringify(key)}:${canonical(record[key], active)}`) + .join(',')}}` + } + active.delete(value) + return encoded +} + +export function canonicalInterruptJson(value: unknown): string { + return canonical(value, new WeakSet()) +} + +export function digestInterruptJson(canonicalJson: string): string { + if (typeof canonicalJson !== 'string') { + throw new TypeError('Interrupt digests require canonical JSON text.') + } + return `sha256:${bytesToHex(sha256(utf8ToBytes(canonicalJson)))}` +} + +function freezeTree(value: unknown): void { + if (value === null || typeof value !== 'object' || Object.isFrozen(value)) { + return + } + Object.values(value).forEach(freezeTree) + Object.freeze(value) +} + +export function cloneAndDeepFreezeJson(value: T): T { + const clone: T = JSON.parse(canonicalInterruptJson(value)) + freezeTree(clone) + return clone +} diff --git a/packages/ai/src/interrupts.ts b/packages/ai/src/interrupts.ts new file mode 100644 index 000000000..8768c1d3a --- /dev/null +++ b/packages/ai/src/interrupts.ts @@ -0,0 +1,200 @@ +import { createCapability } from './activities/chat/middleware/capabilities' +import { + canonicalInterruptJson, + cloneAndDeepFreezeJson, + digestInterruptJson, +} from './interrupt-serialization' +import type { Interrupt, RunAgentResumeItem } from './types' + +export interface OpenInterruptBatchInput { + threadId: string + interruptedRunId: string + descriptors: ReadonlyArray + bindings: ReadonlyArray +} + +export interface CommitInterruptResolutionsInput { + threadId: string + interruptedRunId: string + continuationRunId: string + expectedGeneration: number + expectedInterruptIds: ReadonlyArray + resolutions: ReadonlyArray + fingerprint: string + canonicalResolutions: string +} + +export interface InterruptCorrelation { + threadId: string + interruptedRunId: string + generation: number + submissionId?: string + continuationRunId?: string +} + +export type ItemInterruptErrorCode = + | 'invalid-payload' + | 'invalid-edited-args' + | 'invalid-tool-output' + | 'invalid-response-schema' + | 'unknown-interrupt' + | 'expired' + | 'stale' + | 'conflict' + | 'legacy-unsupported' + +export type BatchInterruptErrorCode = + | 'incomplete-batch' + | 'item-validation-failed' + | 'unsupported-bulk-operation' + | 'async-resolver' + | 'inactive-transaction' + | 'mixed-provenance' + | 'transport' + | 'server' + | 'protocol' + | 'invalid-response-schema' + | 'expired' + | 'stale' + | 'conflict' + | 'persistence-required' + | 'atomic-commit-unsupported' + | 'recovery-unavailable' + | 'legacy-submit-failed' + +export interface ItemInterruptError extends InterruptCorrelation { + scope: 'item' + interruptId: string + code: ItemInterruptErrorCode + message: string + path?: ReadonlyArray + source: 'client' | 'server' + retryable: boolean +} + +export interface BatchInterruptError extends InterruptCorrelation { + scope: 'batch' + code: BatchInterruptErrorCode + message: string + source: 'client' | 'server' | 'transport' + retryable: boolean + interruptIds: ReadonlyArray +} + +export type InterruptSubmissionError = ItemInterruptError | BatchInterruptError + +export type InterruptBinding = + | { + kind: 'tool-approval' + interruptId: string + interruptedRunId: string + generation: number + toolName: string + toolCallId: string + originalArgs: unknown + inputSchemaHash: string + approvalSchemaHash: string + responseSchemaHash: string + expiresAt?: string + } + | { + kind: 'client-tool-execution' + interruptId: string + interruptedRunId: string + generation: number + toolName: string + toolCallId: string + outputSchemaHash: string + responseSchemaHash: string + expiresAt?: string + } + | { + kind: 'generic' + interruptId: string + interruptedRunId: string + generation: number + responseSchemaHash: string + expiresAt?: string + } + +export type UnopenedInterruptBinding = InterruptBinding extends infer TBinding + ? TBinding extends InterruptBinding + ? Omit + : never + : never + +export type InterruptCommitResult = + | { status: 'committed'; continuationRunId: string } + | { status: 'replayed'; continuationRunId: string } + | { + status: 'conflict' + authoritativeState: InterruptRecoveryStateV1 + } + +export interface InterruptRecoveryQuery { + threadId: string + interruptedRunId: string + knownGeneration: number +} + +export interface InterruptRecoveryStateV1 extends InterruptCorrelation { + schemaVersion: 1 + state: 'pending' | 'committed' | 'expired' | 'missing' | 'legacy-committed' + pendingInterrupts: ReadonlyArray + committed?: { + fingerprint: string + resolutions?: ReadonlyArray + continuationRunId?: string + committedAt: string + } +} + +export interface InterruptPersistenceGateway { + openInterruptBatch: ( + input: OpenInterruptBatchInput, + ) => Promise<{ generation: number; descriptors: ReadonlyArray }> + commitInterruptResolutions: ( + input: CommitInterruptResolutionsInput, + ) => Promise + getInterruptRecoveryState: ( + input: InterruptRecoveryQuery, + ) => Promise +} + +export type ToolApprovalResolution = + | boolean + | { + approved: true + editedArgs?: unknown + payload?: unknown + } + | { + approved: false + payload?: unknown + editedArgs?: never + } + +export function canonicalizeInterruptResolutions( + resolutions: ReadonlyArray, +): { + resolutions: ReadonlyArray + canonicalResolutions: string + fingerprint: string +} { + const sorted = [...resolutions].sort((left, right) => + left.interruptId.localeCompare(right.interruptId), + ) + const frozen = cloneAndDeepFreezeJson(sorted) + const canonicalResolutions = canonicalInterruptJson(frozen) + return Object.freeze({ + resolutions: frozen, + canonicalResolutions, + fingerprint: digestInterruptJson(canonicalResolutions), + }) +} + +export const InterruptPersistenceCapability = + createCapability()('interrupt-persistence') + +export const [getInterruptPersistence, provideInterruptPersistence] = + InterruptPersistenceCapability diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index a585cac23..7e3c0abd5 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -5,6 +5,10 @@ import type { import type { InternalLogger } from './logger/internal-logger' import type { SystemPrompt } from './system-prompts' import type { CapabilityContext } from './activities/chat/middleware/capabilities' +import type { + InterruptRecoveryStateV1, + InterruptSubmissionError, +} from './interrupts' // The canonical usage types live in the leaf `@tanstack/ai-event-client` // package (which `@tanstack/ai` already depends on) so there is a single source // of truth without a dependency cycle. They are re-exported below. @@ -572,8 +576,8 @@ export type ToolExecutionContext = } export type ToolExecuteFunction< - TInput extends SchemaInput = SchemaInput, - TOutput extends SchemaInput = SchemaInput, + TInput extends SchemaInput | undefined = SchemaInput, + TOutput extends SchemaInput | undefined = SchemaInput, TContext = unknown, > = undefined extends TContext ? ( @@ -599,8 +603,8 @@ export type ToolExecuteFunction< * @see https://standardschema.dev/json-schema */ export interface Tool< - TInput extends SchemaInput = SchemaInput, - TOutput extends SchemaInput = SchemaInput, + TInput extends SchemaInput | undefined = SchemaInput, + TOutput extends SchemaInput | undefined = SchemaInput, TName extends string = string, TContext = unknown, > { @@ -982,6 +986,9 @@ export interface TextOptions< */ parentRunId?: string + /** Application state mirrored in a STATE_SNAPSHOT before an interrupt terminal. */ + state?: unknown + /** * AG-UI interrupt resume responses supplied by the client on a follow-up run. * Threaded through request parsing now so later runtime behavior can resolve @@ -1108,6 +1115,10 @@ export interface RunFinishedEvent extends AGUIRunFinishedEvent { export interface RunErrorEvent extends AGUIRunErrorEvent { /** Model identifier for multi-model support */ model?: string + /** Exhaustive TanStack interrupt submission failures for this run. */ + 'tanstack:interruptErrors'?: ReadonlyArray + /** Authoritative interrupt state returned with stale/conflicting submissions. */ + 'tanstack:interruptRecovery'?: InterruptRecoveryStateV1 /** * @deprecated Use top-level `message` and `code` fields instead. * Kept for backward compatibility. @@ -1368,6 +1379,10 @@ export interface StructuredOutputStartEvent extends CustomEvent { * (the agent-loop branch of `runStreamingStructuredOutputImpl` in * `activities/chat/index.ts` forwards CUSTOM events from `TextEngine.run()`). */ +/** + * @deprecated Native interrupts use RUN_FINISHED interrupt outcomes. This + * compatibility event remains readable until 1.0. + */ export interface ApprovalRequestedEvent extends CustomEvent { name: 'approval-requested' value: { @@ -1384,6 +1399,10 @@ export interface ApprovalRequestedEvent extends CustomEvent { * will not fire for that run. Shape fixed by the agent-loop forwarding in * `runStreamingStructuredOutputImpl` in `activities/chat/index.ts`. */ +/** + * @deprecated Native interrupts use RUN_FINISHED interrupt outcomes. This + * compatibility event remains readable until 1.0. + */ export interface ToolInputAvailableEvent extends CustomEvent { name: 'tool-input-available' value: { diff --git a/packages/ai/tests/chat-mcp-manager.test.ts b/packages/ai/tests/chat-mcp-manager.test.ts index de8e982de..73f3c78a2 100644 --- a/packages/ai/tests/chat-mcp-manager.test.ts +++ b/packages/ai/tests/chat-mcp-manager.test.ts @@ -8,7 +8,7 @@ import { type ToolResult, } from '../src/activities/chat/tools/tool-calls' import { EventType } from '../src/types' -import type { ServerTool } from '../src' +import type { JSONSchema, ServerTool } from '../src' import type { CustomEvent, ToolCall, @@ -41,11 +41,13 @@ interface ReadResourceResult { } /** A ui-linked ServerTool whose `metadata.mcp` is statically typed. */ -interface UiServerTool extends ServerTool { +type TestServerTool = ServerTool + +interface UiServerTool extends TestServerTool { metadata: McpAppMetadata } -function tool(name: string): ServerTool { +function tool(name: string): TestServerTool { return { __toolSide: 'server', name, @@ -55,7 +57,7 @@ function tool(name: string): ServerTool { } } -function source(tools: Array, opts: { fail?: boolean } = {}) { +function source(tools: Array, opts: { fail?: boolean } = {}) { const s = { closed: false, tools: async (_o?: { lazy?: boolean }) => { @@ -193,7 +195,7 @@ type EmittedEvent = { name: string; value: UIResourceEvent['value'] } * CUSTOM events emitted via the same `emitCustomEvent` closure chat() wires in. */ async function runToolResult( - tool: ServerTool, + tool: TestServerTool, toolCallId: string, onEvent: (event: EmittedEvent) => void, ): Promise> { diff --git a/packages/ai/tests/chat.test.ts b/packages/ai/tests/chat.test.ts index c92169229..3786b5091 100644 --- a/packages/ai/tests/chat.test.ts +++ b/packages/ai/tests/chat.test.ts @@ -1,8 +1,15 @@ import { describe, expect, it, vi } from 'vitest' import { chat, createChatOptions } from '../src/activities/chat/index' +import { defineChatMiddleware } from '../src/activities/chat/middleware/define' import { DISCOVERY_TOOL_NAME } from '../src/activities/chat/tools/lazy-tool-manager' +import { + InterruptPersistenceCapability, + provideInterruptPersistence, +} from '../src/interrupts' import { EventType } from '../src/types' import type { StreamChunk, Tool } from '../src/types' +import type { InterruptPersistenceGateway } from '../src/interrupts' +import type { ChatResumeToolState } from '../src/activities/chat/middleware/types' import { chunk, ev, @@ -33,6 +40,60 @@ function expectSingleRunFinished( return terminals[0]! } +function interruptPersistenceMiddleware(sequence?: string[]) { + const gateway: InterruptPersistenceGateway = { + openInterruptBatch: async (input) => { + sequence?.push('persist') + return { generation: 1, descriptors: input.descriptors } + }, + commitInterruptResolutions: async (input) => ({ + status: 'committed', + continuationRunId: input.continuationRunId, + }), + getInterruptRecoveryState: async (input) => ({ + schemaVersion: 1, + state: 'missing', + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + generation: input.knownGeneration, + pendingInterrupts: [], + }), + } + return defineChatMiddleware({ + name: 'test-interrupt-persistence', + provides: [InterruptPersistenceCapability], + setup(ctx) { + provideInterruptPersistence(ctx, gateway) + }, + async onChunk(_ctx, value) { + if ( + value.type === EventType.RUN_FINISHED && + value.outcome?.type === 'interrupt' + ) { + await gateway.openInterruptBatch({ + threadId: value.threadId, + interruptedRunId: value.runId, + descriptors: value.outcome.interrupts, + bindings: [], + }) + } else if (value.type === EventType.MESSAGES_SNAPSHOT) { + sequence?.push('messages') + } else if (value.type === EventType.STATE_SNAPSHOT) { + sequence?.push('state') + } + }, + }) +} + +function resumeStateMiddleware(resumeToolState: ChatResumeToolState) { + return defineChatMiddleware({ + name: 'test-interrupt-resume-state', + onConfig() { + return { resumeToolState } + }, + }) +} + // ============================================================================ // Tests // ============================================================================ @@ -42,6 +103,115 @@ describe('chat()', () => { // Streaming text (no tools) // ========================================================================== describe('streaming text (no tools)', () => { + it('turns an exact interrupt replay signal into a successful terminal', async () => { + const { adapter, calls } = createMockAdapter({ + iterations: [[ev.runStarted(), ev.runFinished('stop')]], + }) + const replay = defineChatMiddleware({ + name: 'test-interrupt-replay', + onConfig(ctx) { + if (ctx.phase !== 'init') return + const error = new Error('already committed') + error.name = 'InterruptReplaySignal' + Object.defineProperty(error, 'continuationRunId', { + value: 'committed-run', + enumerable: true, + }) + throw error + }, + }) + + const chunks = await collectChunks( + chat({ + adapter, + messages: [], + threadId: 'thread-1', + runId: 'replay-run', + parentRunId: 'interrupted-run', + middleware: [replay], + }) as AsyncIterable, + ) + + expect(calls).toHaveLength(0) + expect(chunks).toEqual([ + expect.objectContaining({ + type: EventType.RUN_FINISHED, + threadId: 'thread-1', + runId: 'replay-run', + outcome: { type: 'success' }, + }), + ]) + }) + + it('turns interrupt validation failures into one structured error terminal', async () => { + const { adapter, calls } = createMockAdapter({ + iterations: [[ev.runStarted(), ev.runFinished('stop')]], + }) + const onError = vi.fn() + const errors = [ + { + scope: 'item', + interruptId: 'interrupt-1', + code: 'invalid-payload', + message: 'The interrupt payload is invalid.', + source: 'server', + retryable: false, + threadId: 'thread-1', + interruptedRunId: 'interrupted-run', + generation: 1, + }, + ] as const + const recovery = { + schemaVersion: 1, + state: 'pending', + threadId: 'thread-1', + interruptedRunId: 'interrupted-run', + generation: 1, + pendingInterrupts: [], + } as const + const validation = defineChatMiddleware({ + name: 'test-interrupt-validation-failure', + onConfig(ctx) { + if (ctx.phase !== 'init') return + const error = new Error(errors[0].message) + error.name = 'InterruptResumeValidationError' + Object.defineProperties(error, { + errors: { value: errors, enumerable: true }, + recovery: { value: recovery, enumerable: true }, + }) + throw error + }, + onError(_ctx, info) { + onError(info.error) + }, + }) + + const chunks = await collectChunks( + chat({ + adapter, + messages: [], + stream: true, + threadId: 'thread-1', + runId: 'continuation-run', + parentRunId: 'interrupted-run', + middleware: [validation], + }), + ) + + expect(calls).toHaveLength(0) + expect(onError).toHaveBeenCalledTimes(1) + expect(chunks).toEqual([ + expect.objectContaining({ + type: EventType.RUN_ERROR, + threadId: 'thread-1', + runId: 'continuation-run', + message: errors[0].message, + 'tanstack:interruptErrors': errors, + 'tanstack:interruptRecovery': recovery, + }), + ]) + }) + it('should return an async iterable that yields all adapter chunks', async () => { const { adapter } = createMockAdapter({ iterations: [ @@ -420,6 +590,83 @@ describe('chat()', () => { // Client tools (no execute) // ========================================================================== describe('client tools (no execute)', () => { + it('emits one structured error instead of an unpersisted interrupt terminal', async () => { + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('call_1', 'clientSearch'), + ev.toolArgs('call_1', '{"query":"test"}'), + ev.runFinished('tool_calls'), + ], + ], + }) + + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Search' }], + tools: [clientTool('clientSearch')], + }) as AsyncIterable, + ) + + expect( + chunks.filter((value) => value.type === EventType.RUN_ERROR), + ).toHaveLength(1) + expect( + chunks.filter((value) => value.type === EventType.RUN_FINISHED), + ).toHaveLength(0) + expect(chunks.at(-1)).toMatchObject({ + type: EventType.RUN_ERROR, + code: 'persistence-required', + }) + }) + + it('persists before ordered snapshots and emits canonical bound interrupts', async () => { + const sequence: string[] = [] + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('call_1', 'clientSearch'), + ev.toolArgs('call_1', '{"query":"test"}'), + ev.runFinished('tool_calls'), + ], + ], + }) + + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Search' }], + tools: [clientTool('clientSearch')], + state: { screen: 'search' }, + middleware: [interruptPersistenceMiddleware(sequence)], + }) as AsyncIterable, + ) + + expect(sequence).toEqual(['persist', 'messages', 'state']) + expect(chunks.slice(-3).map((value) => value.type)).toEqual([ + EventType.MESSAGES_SNAPSHOT, + EventType.STATE_SNAPSHOT, + EventType.RUN_FINISHED, + ]) + const terminal = expectSingleRunFinished(chunks) + expect(terminal).toMatchObject({ + outcome: { + type: 'interrupt', + interrupts: [ + { + id: 'client_tool_call_1', + reason: 'tanstack:client_tool_execution', + toolCallId: 'call_1', + responseSchema: {}, + }, + ], + }, + }) + }) + it('should yield an interrupt outcome for client tools', async () => { const { adapter } = createMockAdapter({ iterations: [ @@ -436,6 +683,7 @@ describe('chat()', () => { adapter, messages: [{ role: 'user', content: 'Search for test' }], tools: [clientTool('clientSearch')], + middleware: [interruptPersistenceMiddleware()], }) const chunks = await collectChunks(stream as AsyncIterable) @@ -449,7 +697,7 @@ describe('chat()', () => { interrupts: [ { id: 'client_tool_call_1', - reason: 'client_tool_input', + reason: 'tanstack:client_tool_execution', toolCallId: 'call_1', metadata: { kind: 'client_tool', @@ -489,6 +737,7 @@ describe('chat()', () => { required: ['status'], }, stream: true, + middleware: [interruptPersistenceMiddleware()], }) const chunks = await collectChunks( @@ -553,6 +802,7 @@ describe('chat()', () => { serverTool('searchTools', searchExecute), clientTool('showNotification'), ], + middleware: [interruptPersistenceMiddleware()], }) const chunks = await collectChunks(stream as AsyncIterable) @@ -581,7 +831,7 @@ describe('chat()', () => { type: 'interrupt', interrupts: [ { - reason: 'client_tool_input', + reason: 'tanstack:client_tool_execution', toolCallId: 'call_client', metadata: { kind: 'client_tool', @@ -641,6 +891,7 @@ describe('chat()', () => { serverTool('getWeather', weatherExecute), clientTool('showNotification'), ], + middleware: [interruptPersistenceMiddleware()], }) const chunks = await collectChunks(stream as AsyncIterable) @@ -669,7 +920,7 @@ describe('chat()', () => { type: 'interrupt', interrupts: [ { - reason: 'client_tool_input', + reason: 'tanstack:client_tool_execution', toolCallId: 'call_client', metadata: { kind: 'client_tool', @@ -696,6 +947,65 @@ describe('chat()', () => { // Approval flow // ========================================================================== describe('approval flow', () => { + it('applies approved argument edits and emits only a result for a resumed tool call', async () => { + const execute = vi.fn().mockImplementation((input) => input) + const { adapter } = createMockAdapter({ + iterations: [[ev.runStarted(), ev.runFinished('stop')]], + }) + const stream = chat({ + adapter, + messages: [ + { role: 'user', content: 'Change it' }, + { + role: 'assistant', + content: '', + toolCalls: [ + { + id: 'call_1', + type: 'function', + function: { + name: 'dangerousTool', + arguments: '{"action":"original"}', + }, + }, + ], + }, + ], + tools: [ + { + ...serverTool('dangerousTool', execute), + needsApproval: true, + inputSchema: { + type: 'object', + properties: { action: { type: 'string' } }, + required: ['action'], + }, + }, + ], + middleware: [ + interruptPersistenceMiddleware(), + resumeStateMiddleware({ + approvals: new Map([ + ['call_1', { approved: true, editedArgs: { action: 'edited' } }], + ]), + }), + ], + }) + + const chunks = await collectChunks(stream as AsyncIterable) + expect(execute).toHaveBeenCalledWith( + { action: 'edited' }, + expect.any(Object), + ) + expect( + chunks + .filter( + (value) => 'toolCallId' in value && value.toolCallId === 'call_1', + ) + .map((value) => value.type), + ).toEqual([EventType.TOOL_CALL_RESULT]) + }) + it('should end with an interrupt outcome for tools with needsApproval', async () => { const { adapter } = createMockAdapter({ iterations: [ @@ -715,6 +1025,7 @@ describe('chat()', () => { ...t, needsApproval: true, })), + middleware: [interruptPersistenceMiddleware()], }) const chunks = await collectChunks(stream as AsyncIterable) @@ -729,14 +1040,10 @@ describe('chat()', () => { interrupts: [ { id: 'approval_call_1', - reason: 'approval_required', + reason: 'tool_call', message: 'Approval required to run dangerousTool', toolCallId: 'call_1', - responseSchema: { - type: 'object', - properties: { approved: { type: 'boolean' } }, - required: ['approved'], - }, + responseSchema: { oneOf: expect.any(Array) }, metadata: { kind: 'approval', toolName: 'dangerousTool', @@ -771,6 +1078,7 @@ describe('chat()', () => { adapter, messages: [{ role: 'user', content: 'Do something' }], tools: [clientTool('clientDanger', { needsApproval: true })], + middleware: [interruptPersistenceMiddleware()], }) const chunks = await collectChunks(stream as AsyncIterable) @@ -784,7 +1092,7 @@ describe('chat()', () => { interrupts: [ { id: 'approval_call_1', - reason: 'approval_required', + reason: 'tool_call', toolCallId: 'call_1', metadata: { kind: 'approval', @@ -813,6 +1121,7 @@ describe('chat()', () => { adapter, messages: [{ role: 'user', content: 'Search for test' }], tools: [clientTool('clientSearch')], + middleware: [interruptPersistenceMiddleware()], }) const chunks = await collectChunks(stream as AsyncIterable) @@ -826,7 +1135,7 @@ describe('chat()', () => { interrupts: [ { id: 'client_tool_call_1', - reason: 'client_tool_input', + reason: 'tanstack:client_tool_execution', message: 'Client tool clientSearch is ready to run', toolCallId: 'call_1', metadata: { @@ -973,7 +1282,7 @@ describe('chat()', () => { expect(calls).toHaveLength(1) }) - it('should emit TOOL_CALL_START and TOOL_CALL_ARGS before TOOL_CALL_END for pending tool calls', async () => { + it('should emit only TOOL_CALL_RESULT for pending tool calls', async () => { const executeSpy = vi.fn().mockReturnValue({ temp: 72 }) const { adapter } = createMockAdapter({ @@ -1014,42 +1323,14 @@ describe('chat()', () => { // Tool should have been executed expect(executeSpy).toHaveBeenCalledTimes(1) - // The continuation re-execution should emit the full chunk sequence: - // TOOL_CALL_START -> TOOL_CALL_ARGS -> TOOL_CALL_END - // Without the fix, only TOOL_CALL_END is emitted, causing the client - // to store the tool call with empty arguments {}. - const toolStartChunks = chunks.filter( - (c) => - c.type === 'TOOL_CALL_START' && (c as any).toolCallId === 'call_1', - ) - expect(toolStartChunks).toHaveLength(1) - // Both AG-UI spec field `toolCallName` and deprecated alias `toolName` - // must be set so consumers reading either get a valid name (issue #532). - expect((toolStartChunks[0] as any).toolCallName).toBe('getWeather') - expect((toolStartChunks[0] as any).toolName).toBe('getWeather') - - const toolArgsChunks = chunks.filter( - (c) => - c.type === 'TOOL_CALL_ARGS' && (c as any).toolCallId === 'call_1', - ) - expect(toolArgsChunks).toHaveLength(1) - expect((toolArgsChunks[0] as any).delta).toBe('{"city":"NYC"}') - expect((toolArgsChunks[0] as any).args).toBe('{"city":"NYC"}') - - const toolEndChunks = chunks.filter( - (c) => c.type === 'TOOL_CALL_END' && (c as any).toolCallId === 'call_1', - ) - expect(toolEndChunks).toHaveLength(1) - - // Verify ordering: START before ARGS before END - const startIdx = chunks.indexOf(toolStartChunks[0]!) - const argsIdx = chunks.indexOf(toolArgsChunks[0]!) - const endIdx = chunks.indexOf(toolEndChunks[0]!) - expect(startIdx).toBeLessThan(argsIdx) - expect(argsIdx).toBeLessThan(endIdx) + expect( + chunks + .filter((c) => 'toolCallId' in c && c.toolCallId === 'call_1') + .map((c) => c.type), + ).toEqual([EventType.TOOL_CALL_RESULT]) }) - it('should emit TOOL_CALL_START and TOOL_CALL_ARGS for each pending tool call in a batch', async () => { + it('should emit only TOOL_CALL_RESULT for each pending tool call in a batch', async () => { const weatherSpy = vi.fn().mockReturnValue({ temp: 72 }) const timeSpy = vi.fn().mockReturnValue({ time: '3pm' }) @@ -1099,39 +1380,16 @@ describe('chat()', () => { expect(weatherSpy).toHaveBeenCalledTimes(1) expect(timeSpy).toHaveBeenCalledTimes(1) - // Each pending tool should get the full START -> ARGS -> END sequence - for (const { id, name, args } of [ - { id: 'call_weather', name: 'getWeather', args: '{"city":"NYC"}' }, - { id: 'call_time', name: 'getTime', args: '{"tz":"EST"}' }, - ]) { - const starts = chunks.filter( - (c) => c.type === 'TOOL_CALL_START' && (c as any).toolCallId === id, - ) - expect(starts).toHaveLength(1) - expect((starts[0] as any).toolCallName).toBe(name) - expect((starts[0] as any).toolName).toBe(name) - - const argChunks = chunks.filter( - (c) => c.type === 'TOOL_CALL_ARGS' && (c as any).toolCallId === id, - ) - expect(argChunks).toHaveLength(1) - expect((argChunks[0] as any).delta).toBe(args) - - const ends = chunks.filter( - (c) => c.type === 'TOOL_CALL_END' && (c as any).toolCallId === id, - ) - expect(ends).toHaveLength(1) - - // Verify ordering - const startIdx = chunks.indexOf(starts[0]!) - const argsIdx = chunks.indexOf(argChunks[0]!) - const endIdx = chunks.indexOf(ends[0]!) - expect(startIdx).toBeLessThan(argsIdx) - expect(argsIdx).toBeLessThan(endIdx) + for (const id of ['call_weather', 'call_time']) { + expect( + chunks + .filter((c) => 'toolCallId' in c && c.toolCallId === id) + .map((c) => c.type), + ).toEqual([EventType.TOOL_CALL_RESULT]) } }) - it('should emit TOOL_CALL_START and TOOL_CALL_ARGS for the server tool in a mixed pending batch', async () => { + it('should emit only TOOL_CALL_RESULT for the server tool in a mixed pending batch', async () => { const weatherSpy = vi.fn().mockReturnValue({ temp: 72 }) const { adapter } = createMockAdapter({ iterations: [] }) @@ -1165,6 +1423,7 @@ describe('chat()', () => { serverTool('getWeather', weatherSpy), clientTool('showNotification'), ], + middleware: [interruptPersistenceMiddleware()], }) const chunks = await collectChunks(stream as AsyncIterable) @@ -1172,36 +1431,11 @@ describe('chat()', () => { // Server tool should have executed expect(weatherSpy).toHaveBeenCalledTimes(1) - // The executed server tool should get the full START -> ARGS -> END - const starts = chunks.filter( - (c) => - c.type === 'TOOL_CALL_START' && - (c as any).toolCallId === 'call_server', - ) - expect(starts).toHaveLength(1) - expect((starts[0] as any).toolCallName).toBe('getWeather') - expect((starts[0] as any).toolName).toBe('getWeather') - - const argChunks = chunks.filter( - (c) => - c.type === 'TOOL_CALL_ARGS' && - (c as any).toolCallId === 'call_server', - ) - expect(argChunks).toHaveLength(1) - expect((argChunks[0] as any).delta).toBe('{"city":"NYC"}') - - const ends = chunks.filter( - (c) => - c.type === 'TOOL_CALL_END' && (c as any).toolCallId === 'call_server', - ) - expect(ends).toHaveLength(1) - - // Verify ordering - const startIdx = chunks.indexOf(starts[0]!) - const argsIdx = chunks.indexOf(argChunks[0]!) - const endIdx = chunks.indexOf(ends[0]!) - expect(startIdx).toBeLessThan(argsIdx) - expect(argsIdx).toBeLessThan(endIdx) + expect( + chunks + .filter((c) => 'toolCallId' in c && c.toolCallId === 'call_server') + .map((c) => c.type), + ).toEqual([EventType.TOOL_CALL_RESULT]) }) it('should replace pendingExecution placeholder with the real tool result and supply both toolCallName/toolName (issue #532)', async () => { @@ -1261,19 +1495,11 @@ describe('chat()', () => { // it as pendingExecution. expect(executeSpy).toHaveBeenCalledTimes(1) - // Synthesized TOOL_CALL_START must include both `toolCallName` (AG-UI - // spec) and `toolName` (deprecated alias). Without `toolCallName` the - // chat-client's StreamProcessor would create a tool-call part with - // name=undefined and the next outbound request would fail at Anthropic - // with `tool_use.name: String should have at least 1 character`. - const toolStart = chunks.find( - (c) => - c.type === 'TOOL_CALL_START' && - (c as any).toolCallId === 'call_approval', - ) - expect(toolStart).toBeDefined() - expect((toolStart as any).toolCallName).toBe('approvedTool') - expect((toolStart as any).toolName).toBe('approvedTool') + expect( + chunks + .filter((c) => 'toolCallId' in c && c.toolCallId === 'call_approval') + .map((c) => c.type), + ).toEqual([EventType.TOOL_CALL_RESULT]) // The follow-up adapter call (after the tool ran) must see the real // tool result, not the placeholder. With the placeholder still in the diff --git a/packages/ai/tests/interrupts-types.test-d.ts b/packages/ai/tests/interrupts-types.test-d.ts new file mode 100644 index 000000000..99b501e2f --- /dev/null +++ b/packages/ai/tests/interrupts-types.test-d.ts @@ -0,0 +1,74 @@ +import { expectTypeOf } from 'vitest' +import { z } from 'zod' +import { + toolDefinition, + type ApprovalCapabilityOf, + type ApprovalSchemaOf, + type InferToolInput, + type InferToolOutput, + type InputSchemaOf, + type NoSchema, + type RunErrorEvent, + type ChatMiddlewareContext, + type ChatResumeToolState, +} from '../src' +import type { InterruptSubmissionError } from '../src/interrupts' + +const transfer = toolDefinition({ + name: 'transfer', + description: 'Transfer funds', + needsApproval: true, + inputSchema: z.object({ cents: z.number() }), + outputSchema: z.object({ receipt: z.string() }), + approvalSchema: { + approve: z.object({ note: z.string() }), + reject: z.object({ reason: z.string() }), + }, +}) + +expectTypeOf>().toEqualTypeOf<{ + cents: number +}>() +expectTypeOf< + InferToolOutput> +>().toEqualTypeOf<{ + receipt: string +}>() +expectTypeOf< + ApprovalCapabilityOf> +>().toEqualTypeOf() +expectTypeOf< + ApprovalSchemaOf> +>().toEqualTypeOf() + +toolDefinition({ + name: 'invalid', + description: 'Cannot declare an approval payload', + // @ts-expect-error approvalSchema requires needsApproval: true + approvalSchema: z.object({ note: z.string() }), +}) + +const noInputDefinition = toolDefinition({ + name: 'noInput', + description: 'Approval without editable input', + needsApproval: true, +}) +const noInputClient = noInputDefinition.client() +const noInputServer = noInputDefinition.server(async () => ({ ok: true })) +expectTypeOf(noInputClient.inputSchema).toEqualTypeOf() +expectTypeOf(noInputServer.inputSchema).toEqualTypeOf() +expectTypeOf< + InputSchemaOf +>().toEqualTypeOf() +expectTypeOf>().toEqualTypeOf() +expectTypeOf>().toEqualTypeOf() + +expectTypeOf().toEqualTypeOf< + readonly InterruptSubmissionError[] | undefined +>() +expectTypeOf().toEqualTypeOf< + string | undefined +>() +expectTypeOf().toEqualTypeOf< + ReadonlySet | undefined +>() diff --git a/packages/ai/tests/interrupts.test.ts b/packages/ai/tests/interrupts.test.ts index 1c99e0ab5..a809c587b 100644 --- a/packages/ai/tests/interrupts.test.ts +++ b/packages/ai/tests/interrupts.test.ts @@ -1,5 +1,24 @@ import { describe, expect, expectTypeOf, it } from 'vitest' import { EventType } from '@ag-ui/core' +import { + JsonSchemaCompilationError, + compileJsonSchema202012, +} from '../src/activities/chat/tools/json-schema-validator' +import { + hashSchemaInput, + normalizeApprovalSchema, +} from '../src/activities/chat/tools/approval-schema' +import { + canonicalInterruptJson, + cloneAndDeepFreezeJson, + digestInterruptJson, +} from '../src/interrupt-serialization' +import { + InterruptPersistenceCapability, + canonicalizeInterruptResolutions, + type InterruptRecoveryStateV1, + type InterruptSubmissionError, +} from '../src/interrupts' import type { Interrupt, RunAgentResumeItem, @@ -81,3 +100,218 @@ describe('AG-UI interrupt protocol types', () => { >() }) }) + +describe('Draft 2020-12 interrupt schema validation', () => { + it('normalizes every Draft 2020-12 issue without mutating input', () => { + const validate = compileJsonSchema202012({ + $schema: 'https://json-schema.org/draft/2020-12/schema', + type: 'object', + properties: { + 'a/b': { type: 'string', minLength: 3 }, + email: { type: 'string', format: 'email' }, + }, + required: ['a/b', 'email'], + additionalProperties: false, + }) + const input = { 'a/b': 'x', email: 'bad', extra: true } + + expect(validate(input)).toEqual([ + expect.objectContaining({ + keyword: 'additionalProperties', + path: ['extra'], + }), + expect.objectContaining({ path: ['a/b'] }), + expect.objectContaining({ path: ['email'] }), + ]) + expect(input).toEqual({ 'a/b': 'x', email: 'bad', extra: true }) + }) + + it('accepts #, local $defs, and repeated acyclic object identities', () => { + const sharedProperty = { type: 'string', minLength: 2 } as const + const sharedObject = { + type: 'object', + properties: { label: sharedProperty }, + required: ['label'], + additionalProperties: false, + } as const + const root = { + type: 'object', + $defs: { entry: sharedObject }, + properties: { + selfCopy: { $ref: '#' }, + first: { $ref: '#/$defs/entry' }, + second: sharedObject, + third: sharedObject, + }, + additionalProperties: false, + } as const + const validate = compileJsonSchema202012(root) + + expect( + validate({ + first: { label: 'ok' }, + second: { label: 'yes' }, + third: { label: 'no' }, + }), + ).toEqual([]) + }) + + it('rejects unresolved local references, true cycles, Date, and Map', () => { + expect(() => compileJsonSchema202012({ $ref: '#/$defs/missing' })).toThrow( + JsonSchemaCompilationError, + ) + + const cyclic: Record = { type: 'object' } + cyclic['properties'] = { self: cyclic } + expect(() => compileJsonSchema202012(cyclic)).toThrow(/cycles/) + expect(() => compileJsonSchema202012(new Date())).toThrow(/plain JSON/) + expect(() => compileJsonSchema202012(new Map())).toThrow(/plain JSON/) + }) + + it('accepts canonical JSON primitives and rejects non-JSON values', () => { + const validate = compileJsonSchema202012({ + type: 'array', + items: { type: ['string', 'number', 'boolean', 'null'] }, + }) + expect(validate(['text', 1, true, null])).toEqual([]) + expect(() => compileJsonSchema202012({ const: undefined })).toThrow( + /JSON-compatible/, + ) + }) + + it('rejects other dialects, remote references, and unknown formats', () => { + expect(() => + compileJsonSchema202012({ + $schema: 'http://json-schema.org/draft-07/schema#', + type: 'string', + }), + ).toThrow(JsonSchemaCompilationError) + expect(() => + compileJsonSchema202012({ + $ref: 'https://example.com/schema.json', + }), + ).toThrow(/document-local/) + expect(() => + compileJsonSchema202012({ type: 'string', format: 'private-id' }), + ).toThrow(JsonSchemaCompilationError) + }) +}) + +describe('approval schema normalization and interrupt serialization', () => { + it('builds branch-aware approval envelopes and rejects malformed maps', () => { + const normalized = normalizeApprovalSchema({ + approve: { type: 'object', required: ['note'] }, + }) + expect(normalized.responseSchema).toMatchObject({ + oneOf: expect.any(Array), + }) + expect(normalized.branches.reject).toBeNull() + expect(() => normalizeApprovalSchema({})).toThrow(/approve or reject/) + expect(() => + normalizeApprovalSchema({ + approve: { unsupported: Symbol('bad') }, + }), + ).toThrow(/SchemaInput/) + }) + + it('normalizes shared and omitted payload branches with stable hashes', () => { + const shared = { type: 'string', minLength: 1 } + const left = normalizeApprovalSchema(shared, { + type: 'object', + properties: { value: { type: 'number' } }, + }) + const right = normalizeApprovalSchema( + { minLength: 1, type: 'string' }, + { + properties: { value: { type: 'number' } }, + type: 'object', + }, + ) + const omitted = normalizeApprovalSchema(undefined) + + expect(left.branches.approve?.jsonSchema).toBe(shared) + expect(left.branches.reject?.jsonSchema).toBe(shared) + expect(left.responseSchemaHash).toBe(right.responseSchemaHash) + expect(left.approvalSchemaHash).toBe(right.approvalSchemaHash) + expect(omitted.branches).toEqual({ approve: null, reject: null }) + expect(hashSchemaInput(undefined)).toMatch(/^sha256:[0-9a-f]{64}$/) + }) + + it('canonicalizes JSON without hashing functions and deeply freezes clones', () => { + const canonical = canonicalInterruptJson({ z: 1, a: { value: true } }) + expect(canonical).toBe('{"a":{"value":true},"z":1}') + expect(digestInterruptJson(canonical)).toMatch(/^sha256:[0-9a-f]{64}$/) + + const frozen = cloneAndDeepFreezeJson({ nested: { ok: true } }) + expect(Object.isFrozen(frozen)).toBe(true) + expect(Object.isFrozen(frozen.nested)).toBe(true) + expect(() => canonicalInterruptJson({ handler: () => undefined })).toThrow( + /JSON-compatible/, + ) + + const cyclic: Record = {} + cyclic['self'] = cyclic + expect(() => canonicalInterruptJson(cyclic)).toThrow(/cycle/) + }) +}) + +describe('core interrupt correlation and persistence seam', () => { + it('correlates interrupt errors and recovery to the interrupted run', () => { + const recovery = { + schemaVersion: 1, + state: 'committed', + threadId: 'thread-1', + interruptedRunId: 'run-old', + generation: 2, + pendingInterrupts: [], + committed: { + fingerprint: 'sha256:abc', + resolutions: [], + continuationRunId: 'run-new', + committedAt: '2026-07-13T10:00:00.000Z', + }, + } satisfies InterruptRecoveryStateV1 + const error: InterruptSubmissionError = { + scope: 'batch', + code: 'conflict', + message: 'Another response batch won.', + source: 'server', + retryable: false, + interruptIds: ['interrupt-1'], + threadId: 'thread-1', + interruptedRunId: 'run-old', + generation: 2, + } + + expect(recovery.committed.continuationRunId).toBe('run-new') + expect(error.interruptedRunId).toBe('run-old') + expect(InterruptPersistenceCapability.capabilityName).toBe( + 'interrupt-persistence', + ) + }) + + it('canonicalizes and deeply freezes response batches in core', () => { + const left = canonicalizeInterruptResolutions([ + { interruptId: 'b', status: 'cancelled' }, + { + interruptId: 'a', + status: 'resolved', + payload: { z: 1, a: 2 }, + }, + ]) + const right = canonicalizeInterruptResolutions([ + { + interruptId: 'a', + status: 'resolved', + payload: { a: 2, z: 1 }, + }, + { interruptId: 'b', status: 'cancelled' }, + ]) + + expect(left.canonicalResolutions).toBe(right.canonicalResolutions) + expect(left.fingerprint).toMatch(/^sha256:[0-9a-f]{64}$/) + expect(left.fingerprint).toBe(right.fingerprint) + expect(Object.isFrozen(left.resolutions)).toBe(true) + expect(Object.isFrozen(left.resolutions[0])).toBe(true) + }) +}) diff --git a/packages/ai/tests/stream-processor.test.ts b/packages/ai/tests/stream-processor.test.ts index d9bb3827b..af2500026 100644 --- a/packages/ai/tests/stream-processor.test.ts +++ b/packages/ai/tests/stream-processor.test.ts @@ -4,6 +4,11 @@ import { createReplayStream, } from '../src/activities/chat/stream/processor' import { chat } from '../src/activities/chat/index' +import { defineChatMiddleware } from '../src/activities/chat/middleware/define' +import { + InterruptPersistenceCapability, + provideInterruptPersistence, +} from '../src/interrupts' import { EventType } from '../src/types' import { createMockAdapter, @@ -23,11 +28,39 @@ import type { UIMessage, } from '../src/types' import type { Message as AGUIMessage } from '@ag-ui/core' +import type { InterruptPersistenceGateway } from '../src/interrupts' // ============================================================================ // Helpers // ============================================================================ +const testInterruptGateway: InterruptPersistenceGateway = { + openInterruptBatch: async (input) => ({ + generation: 1, + descriptors: input.descriptors, + }), + commitInterruptResolutions: async (input) => ({ + status: 'committed', + continuationRunId: input.continuationRunId, + }), + getInterruptRecoveryState: async (input) => ({ + schemaVersion: 1, + state: 'missing', + threadId: input.threadId, + interruptedRunId: input.interruptedRunId, + generation: input.knownGeneration, + pendingInterrupts: [], + }), +} + +const testInterruptPersistence = defineChatMiddleware({ + name: 'test-interrupt-persistence', + provides: [InterruptPersistenceCapability], + setup(ctx) { + provideInterruptPersistence(ctx, testInterruptGateway) + }, +}) + /** Create a typed StreamChunk by event type. Narrows the return to the * matching variant via `Extract`. */ function chunk( @@ -1209,7 +1242,7 @@ describe('StreamProcessor', () => { interrupts: [ { id: 'approval-1', - reason: 'approval_required', + reason: 'tool_call', toolCallId: 'tc-1', metadata: { kind: 'approval', @@ -1252,7 +1285,7 @@ describe('StreamProcessor', () => { interrupts: [ { id: 'client_tool_tc-1', - reason: 'client_tool_input', + reason: 'tanstack:client_tool_execution', toolCallId: 'tc-1', metadata: { kind: 'client_tool', @@ -1302,6 +1335,7 @@ describe('StreamProcessor', () => { serverTool('searchTools', () => ({ results: ['a', 'b'] })), clientTool('showNotification'), ], + middleware: [testInterruptPersistence], }) const chunks = await collectChunks(stream as AsyncIterable) @@ -1314,7 +1348,7 @@ describe('StreamProcessor', () => { toolName: 'showNotification', input: { message: 'done' }, }) - expect(order).toEqual(['tool-call', 'stream-end']) + expect(order).toEqual(['tool-call']) }) it('addToolApprovalResponse should approve a tool call', () => { diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 9e775c9f1..0eeed9b50 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1320,6 +1320,9 @@ importers: '@ag-ui/core': specifier: ^0.0.57 version: 0.0.57 + '@noble/hashes': + specifier: ^2.0.1 + version: 2.2.0 '@standard-schema/spec': specifier: ^1.1.0 version: 1.1.0 @@ -1329,6 +1332,12 @@ importers: '@tanstack/ai-utils': specifier: workspace:* version: link:../ai-utils + ajv: + specifier: ^8.20.0 + version: 8.20.0 + ajv-formats: + specifier: ^3.0.1 + version: 3.0.1(ajv@8.20.0) partial-json: specifier: ^0.1.7 version: 0.1.7 @@ -5845,6 +5854,10 @@ packages: resolution: {integrity: sha512-P06o9TpxrJbiRbHQkiwy/rUrlXRupc+Z8KT4MiJfmcdWxvIdzjCaJOdnNkcOTs6DMyzIOefG5tvk/HLdtjqr0g==} engines: {node: '>= 10'} + '@noble/hashes@2.2.0': + resolution: {integrity: sha512-IYqDGiTXab6FniAgnSdZwgWbomxpy9FtYvLKs7wCUs2a8RkITG+DFGO1DM9cr+E3/RgADRpFjrKVaJ1z6sjtEg==} + engines: {node: '>= 20.19.0'} + '@nodelib/fs.scandir@2.1.5': resolution: {integrity: sha512-vq24Bq3ym5HEQm2NKCr3yXDwjc7vTsEThRDnkp2DK9p1uqLR+DHurm/NOTo0KG7HYHU7eppKZj3MyqYuMBf62g==} engines: {node: '>= 8'} @@ -18851,7 +18864,7 @@ snapshots: '@eslint/eslintrc@3.3.3': dependencies: - ajv: 6.12.6 + ajv: 6.15.0 debug: 4.4.3 espree: 10.4.0 globals: 14.0.0 @@ -19932,6 +19945,8 @@ snapshots: '@ngrok/ngrok-win32-ia32-msvc': 1.7.0 '@ngrok/ngrok-win32-x64-msvc': 1.7.0 + '@noble/hashes@2.2.0': {} + '@nodelib/fs.scandir@2.1.5': dependencies: '@nodelib/fs.stat': 2.0.5 diff --git a/testing/e2e/src/lib/interrupts-v2-fixture.ts b/testing/e2e/src/lib/interrupts-v2-fixture.ts new file mode 100644 index 000000000..71328739e --- /dev/null +++ b/testing/e2e/src/lib/interrupts-v2-fixture.ts @@ -0,0 +1,393 @@ +import { + EventType, + canonicalInterruptJson, + convertSchemaToJsonSchema, + digestInterruptJson, + hashSchemaInput, + normalizeApprovalSchema, + toolDefinition, +} from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { z } from 'zod' +import type { + AnyTextAdapter, + Interrupt, + InterruptBinding, + ModelMessage, + RunAgentResumeItem, + StreamChunk, +} from '@tanstack/ai' + +const editableInputSchema = z.object({ action: z.string().min(1) }) +const sharedApprovalSchema = z.object({ note: z.string().optional() }) +const branchApprovalSchema = { + approve: z.object({ note: z.string() }), + reject: z.object({ reason: z.string() }), +} +const toolOutputSchema = z.object({ ok: z.boolean() }) +const clientToolOutputSchema = z.object({ browserValue: z.string() }) + +const editableActionDefinition = toolDefinition({ + name: 'editable_action', + description: 'Run an editable action after approval.', + inputSchema: editableInputSchema, + outputSchema: toolOutputSchema, + needsApproval: true, + approvalSchema: sharedApprovalSchema, +}) + +const branchActionDefinition = toolDefinition({ + name: 'branch_action', + description: 'Run an action with branch-specific approval payloads.', + inputSchema: z.object({ action: z.string() }), + outputSchema: toolOutputSchema, + needsApproval: true, + approvalSchema: branchApprovalSchema, +}) + +const payloadlessActionDefinition = toolDefinition({ + name: 'payloadless_action', + description: 'Run an action with a payloadless approval decision.', + outputSchema: toolOutputSchema, + needsApproval: true, +}) + +const browserActionDefinition = toolDefinition({ + name: 'browser_action', + description: 'Resolve an action in the browser.', + inputSchema: z.object({ value: z.string() }), + outputSchema: clientToolOutputSchema, +}) + +export const interruptFixtureServerTools = [ + editableActionDefinition.server(() => Promise.resolve({ ok: true })), + branchActionDefinition.server(() => Promise.resolve({ ok: true })), + payloadlessActionDefinition.server(() => Promise.resolve({ ok: true })), + browserActionDefinition, +] as const + +export const interruptFixtureTools = [ + editableActionDefinition.client(), + branchActionDefinition.client(), + payloadlessActionDefinition.client(), + browserActionDefinition.client(() => + Promise.resolve({ browserValue: 'done' }), + ), +] as const + +const genericResponseSchema = { + $schema: 'https://json-schema.org/draft/2020-12/schema', + type: 'object', + properties: { + answer: { type: 'string', minLength: 2 }, + }, + required: ['answer'], + additionalProperties: false, +} as const + +const bindingKey = 'tanstack:interruptBinding' + +function approvalInterrupt( + interruptId: string, + tool: + | typeof editableActionDefinition + | typeof branchActionDefinition + | typeof payloadlessActionDefinition, + originalArgs: unknown, + interruptedRunId: string, + generation: number, +): Interrupt { + const normalized = normalizeApprovalSchema( + tool.approvalSchema, + tool.inputSchema, + ) + const binding = { + kind: 'tool-approval', + interruptId, + interruptedRunId, + generation, + toolName: tool.name, + toolCallId: `tool-call-${interruptId}`, + originalArgs, + inputSchemaHash: hashSchemaInput(tool.inputSchema), + approvalSchemaHash: normalized.approvalSchemaHash, + responseSchemaHash: normalized.responseSchemaHash, + } satisfies InterruptBinding + + return { + id: interruptId, + reason: 'tool_call', + message: `Approval required for ${tool.name}`, + toolCallId: binding.toolCallId, + responseSchema: normalized.responseSchema, + metadata: { + kind: 'approval', + toolName: tool.name, + input: originalArgs, + [bindingKey]: binding, + }, + } +} + +function clientToolInterrupt( + interruptedRunId: string, + generation: number, +): Interrupt { + const responseSchema = + convertSchemaToJsonSchema(browserActionDefinition.outputSchema) ?? {} + const binding = { + kind: 'client-tool-execution', + interruptId: 'client-tool-1', + interruptedRunId, + generation, + toolName: browserActionDefinition.name, + toolCallId: 'tool-call-client-tool-1', + outputSchemaHash: hashSchemaInput(browserActionDefinition.outputSchema), + responseSchemaHash: digestInterruptJson( + canonicalInterruptJson(responseSchema), + ), + } satisfies InterruptBinding + + return { + id: binding.interruptId, + reason: 'tanstack:client_tool_execution', + message: 'Browser action is ready to run.', + toolCallId: binding.toolCallId, + responseSchema, + metadata: { + kind: 'client_tool', + toolName: binding.toolName, + input: { value: 'fixture' }, + [bindingKey]: binding, + }, + } +} + +function genericInterrupt( + interruptedRunId: string, + generation: number, +): Interrupt { + const binding = { + kind: 'generic', + interruptId: 'question-1', + interruptedRunId, + generation, + responseSchemaHash: digestInterruptJson( + canonicalInterruptJson(genericResponseSchema), + ), + } satisfies InterruptBinding + return { + id: binding.interruptId, + reason: 'fixture_question', + message: 'Answer the fixture question.', + responseSchema: genericResponseSchema, + metadata: { + kind: 'generic', + [bindingKey]: binding, + }, + } +} + +function scenarioInterrupts( + scenario: string, + interruptedRunId: string, + generation: number, +): ReadonlyArray { + const editable = () => + approvalInterrupt( + 'approval-1', + editableActionDefinition, + { + action: 'original-action', + }, + interruptedRunId, + generation, + ) + switch (scenario) { + case 'singleton-approval': + case 'commit-then-truncate': + case 'two-tab-conflict': + return [editable()] + case 'three-tool-approvals': + return [ + editable(), + approvalInterrupt( + 'approval-2', + branchActionDefinition, + { + action: 'branch-action', + }, + interruptedRunId, + generation, + ), + approvalInterrupt( + 'approval-3', + payloadlessActionDefinition, + {}, + interruptedRunId, + generation, + ), + ] + case 'heterogeneous-callback': + return [ + editable(), + approvalInterrupt( + 'approval-2', + branchActionDefinition, + { + action: 'branch-action', + }, + interruptedRunId, + generation, + ), + genericInterrupt(interruptedRunId, generation), + ] + case 'generic-validation': + return [genericInterrupt(interruptedRunId, generation)] + case 'two-invalid': + case 'partial-draft-reload': + return [editable(), genericInterrupt(interruptedRunId, generation)] + case 'client-tool-and-approval': + return [clientToolInterrupt(interruptedRunId, generation), editable()] + default: + return [editable()] + } +} + +export interface InterruptFixture { + persistence: ReturnType + continuationCount: number + continuationRunIds: Array + continuationChunks: Map> + decisions: Array + edits: Record + auditHistory: Array + resultEventNames: Array + storedHistory: Array + truncateCommittedResponseOnce: boolean + truncatedResponses: number + replayCount: number + joinedContinuationRunId?: string +} + +const fixtures = new Map() + +export function getInterruptFixture(testId: string): InterruptFixture { + const existing = fixtures.get(testId) + if (existing) return existing + const created: InterruptFixture = { + persistence: memoryPersistence(), + continuationCount: 0, + continuationRunIds: [], + continuationChunks: new Map(), + decisions: [], + edits: {}, + auditHistory: [], + resultEventNames: [], + storedHistory: [], + truncateCommittedResponseOnce: true, + truncatedResponses: 0, + replayCount: 0, + } + fixtures.set(testId, created) + return created +} + +export function createInterruptFixtureAdapter( + scenario: string, + resuming: boolean, +): AnyTextAdapter { + return { + kind: 'text', + name: 'interrupts-v2-fixture', + model: 'interrupts-v2-fixture', + '~types': { + providerOptions: {}, + inputModalities: ['text'], + messageMetadataByModality: {}, + toolCapabilities: [], + toolCallMetadata: undefined, + systemPromptMetadata: undefined, + }, + async *chatStream(options): AsyncGenerator { + await Promise.resolve() + const runId = options.runId ?? crypto.randomUUID() + const threadId = options.threadId ?? `thread-${runId}` + const model = 'interrupts-v2-fixture' + yield { + type: EventType.RUN_STARTED, + runId, + threadId, + model, + timestamp: Date.now(), + } + + if (!resuming) { + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + model, + finishReason: 'stop', + outcome: { + type: 'interrupt', + interrupts: [...scenarioInterrupts(scenario, runId, 1)], + }, + timestamp: Date.now(), + } + return + } + + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + model, + finishReason: 'stop', + outcome: { type: 'success' }, + timestamp: Date.now(), + } + }, + structuredOutput: () => Promise.resolve({ data: {}, rawText: '{}' }), + } +} + +export function fixtureResumeMessages( + scenario: string, +): Array | undefined { + if (scenario !== 'three-tool-approvals') return undefined + return [ + { + role: 'assistant', + content: '', + toolCalls: [ + { + id: 'tool-call-approval-1', + type: 'function', + function: { + name: editableActionDefinition.name, + arguments: JSON.stringify({ action: 'original-action' }), + }, + }, + { + id: 'tool-call-approval-2', + type: 'function', + function: { + name: branchActionDefinition.name, + arguments: JSON.stringify({ action: 'branch-action' }), + }, + }, + ], + }, + ] +} + +export function responseDecision(item: RunAgentResumeItem): string { + if (item.status === 'cancelled') return 'cancel' + const payload = item.payload + if (payload && typeof payload === 'object' && 'approved' in payload) { + return payload.approved === true ? 'approve' : 'deny' + } + if (item.interruptId.startsWith('client-tool')) return 'client-tool' + return 'generic' +} diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 8abc6f917..03d49d81d 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -12,6 +12,7 @@ import { Route as rootRouteImport } from './routes/__root' import { Route as ToolsTestRouteImport } from './routes/tools-test' import { Route as MiddlewareTestRouteImport } from './routes/middleware-test' import { Route as MarkdownCjkRouteImport } from './routes/markdown-cjk' +import { Route as InterruptsV2RouteImport } from './routes/interrupts-v2' import { Route as DevtoolsToolsRouteImport } from './routes/devtools-tools' import { Route as DevtoolsStructuredRouteImport } from './routes/devtools-structured' import { Route as DevtoolsRouteBRouteImport } from './routes/devtools-route-b' @@ -44,6 +45,7 @@ import { Route as ApiMcpAppsServerRouteImport } from './routes/api.mcp-apps-serv import { Route as ApiMcpAppsChatRouteImport } from './routes/api.mcp-apps-chat' import { Route as ApiMcpAppsCallRouteImport } from './routes/api.mcp-apps-call' import { Route as ApiLazyToolsWireRouteImport } from './routes/api.lazy-tools-wire' +import { Route as ApiInterruptsV2RouteImport } from './routes/api.interrupts-v2' import { Route as ApiImageRouteImport } from './routes/api.image' import { Route as ApiDurableDeliveryRouteImport } from './routes/api.durable-delivery' import { Route as ApiChatRouteImport } from './routes/api.chat' @@ -56,6 +58,7 @@ import { Route as ProviderFeatureRouteImport } from './routes/$provider/$feature import { Route as ApiVideoStreamRouteImport } from './routes/api.video.stream' import { Route as ApiTtsStreamRouteImport } from './routes/api.tts.stream' import { Route as ApiTranscriptionStreamRouteImport } from './routes/api.transcription.stream' +import { Route as ApiInterruptsV2RecoveryRouteImport } from './routes/api.interrupts-v2.recovery' import { Route as ApiImageStreamRouteImport } from './routes/api.image.stream' import { Route as ApiAudioStreamRouteImport } from './routes/api.audio.stream' @@ -74,6 +77,11 @@ const MarkdownCjkRoute = MarkdownCjkRouteImport.update({ path: '/markdown-cjk', getParentRoute: () => rootRouteImport, } as any) +const InterruptsV2Route = InterruptsV2RouteImport.update({ + id: '/interrupts-v2', + path: '/interrupts-v2', + getParentRoute: () => rootRouteImport, +} as any) const DevtoolsToolsRoute = DevtoolsToolsRouteImport.update({ id: '/devtools-tools', path: '/devtools-tools', @@ -238,6 +246,11 @@ const ApiLazyToolsWireRoute = ApiLazyToolsWireRouteImport.update({ path: '/api/lazy-tools-wire', getParentRoute: () => rootRouteImport, } as any) +const ApiInterruptsV2Route = ApiInterruptsV2RouteImport.update({ + id: '/api/interrupts-v2', + path: '/api/interrupts-v2', + getParentRoute: () => rootRouteImport, +} as any) const ApiImageRoute = ApiImageRouteImport.update({ id: '/api/image', path: '/api/image', @@ -299,6 +312,11 @@ const ApiTranscriptionStreamRoute = ApiTranscriptionStreamRouteImport.update({ path: '/stream', getParentRoute: () => ApiTranscriptionRoute, } as any) +const ApiInterruptsV2RecoveryRoute = ApiInterruptsV2RecoveryRouteImport.update({ + id: '/recovery', + path: '/recovery', + getParentRoute: () => ApiInterruptsV2Route, +} as any) const ApiImageStreamRoute = ApiImageStreamRouteImport.update({ id: '/stream', path: '/stream', @@ -319,6 +337,7 @@ export interface FileRoutesByFullPath { '/devtools-route-b': typeof DevtoolsRouteBRoute '/devtools-structured': typeof DevtoolsStructuredRoute '/devtools-tools': typeof DevtoolsToolsRoute + '/interrupts-v2': typeof InterruptsV2Route '/markdown-cjk': typeof MarkdownCjkRoute '/middleware-test': typeof MiddlewareTestRoute '/tools-test': typeof ToolsTestRoute @@ -331,6 +350,7 @@ export interface FileRoutesByFullPath { '/api/chat': typeof ApiChatRoute '/api/durable-delivery': typeof ApiDurableDeliveryRoute '/api/image': typeof ApiImageRouteWithChildren + '/api/interrupts-v2': typeof ApiInterruptsV2RouteWithChildren '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute '/api/mcp-apps-call': typeof ApiMcpAppsCallRoute '/api/mcp-apps-chat': typeof ApiMcpAppsChatRoute @@ -357,6 +377,7 @@ export interface FileRoutesByFullPath { '/$provider/': typeof ProviderIndexRoute '/api/audio/stream': typeof ApiAudioStreamRoute '/api/image/stream': typeof ApiImageStreamRoute + '/api/interrupts-v2/recovery': typeof ApiInterruptsV2RecoveryRoute '/api/transcription/stream': typeof ApiTranscriptionStreamRoute '/api/tts/stream': typeof ApiTtsStreamRoute '/api/video/stream': typeof ApiVideoStreamRoute @@ -370,6 +391,7 @@ export interface FileRoutesByTo { '/devtools-route-b': typeof DevtoolsRouteBRoute '/devtools-structured': typeof DevtoolsStructuredRoute '/devtools-tools': typeof DevtoolsToolsRoute + '/interrupts-v2': typeof InterruptsV2Route '/markdown-cjk': typeof MarkdownCjkRoute '/middleware-test': typeof MiddlewareTestRoute '/tools-test': typeof ToolsTestRoute @@ -382,6 +404,7 @@ export interface FileRoutesByTo { '/api/chat': typeof ApiChatRoute '/api/durable-delivery': typeof ApiDurableDeliveryRoute '/api/image': typeof ApiImageRouteWithChildren + '/api/interrupts-v2': typeof ApiInterruptsV2RouteWithChildren '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute '/api/mcp-apps-call': typeof ApiMcpAppsCallRoute '/api/mcp-apps-chat': typeof ApiMcpAppsChatRoute @@ -408,6 +431,7 @@ export interface FileRoutesByTo { '/$provider': typeof ProviderIndexRoute '/api/audio/stream': typeof ApiAudioStreamRoute '/api/image/stream': typeof ApiImageStreamRoute + '/api/interrupts-v2/recovery': typeof ApiInterruptsV2RecoveryRoute '/api/transcription/stream': typeof ApiTranscriptionStreamRoute '/api/tts/stream': typeof ApiTtsStreamRoute '/api/video/stream': typeof ApiVideoStreamRoute @@ -422,6 +446,7 @@ export interface FileRoutesById { '/devtools-route-b': typeof DevtoolsRouteBRoute '/devtools-structured': typeof DevtoolsStructuredRoute '/devtools-tools': typeof DevtoolsToolsRoute + '/interrupts-v2': typeof InterruptsV2Route '/markdown-cjk': typeof MarkdownCjkRoute '/middleware-test': typeof MiddlewareTestRoute '/tools-test': typeof ToolsTestRoute @@ -434,6 +459,7 @@ export interface FileRoutesById { '/api/chat': typeof ApiChatRoute '/api/durable-delivery': typeof ApiDurableDeliveryRoute '/api/image': typeof ApiImageRouteWithChildren + '/api/interrupts-v2': typeof ApiInterruptsV2RouteWithChildren '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute '/api/mcp-apps-call': typeof ApiMcpAppsCallRoute '/api/mcp-apps-chat': typeof ApiMcpAppsChatRoute @@ -460,6 +486,7 @@ export interface FileRoutesById { '/$provider/': typeof ProviderIndexRoute '/api/audio/stream': typeof ApiAudioStreamRoute '/api/image/stream': typeof ApiImageStreamRoute + '/api/interrupts-v2/recovery': typeof ApiInterruptsV2RecoveryRoute '/api/transcription/stream': typeof ApiTranscriptionStreamRoute '/api/tts/stream': typeof ApiTtsStreamRoute '/api/video/stream': typeof ApiVideoStreamRoute @@ -475,6 +502,7 @@ export interface FileRouteTypes { | '/devtools-route-b' | '/devtools-structured' | '/devtools-tools' + | '/interrupts-v2' | '/markdown-cjk' | '/middleware-test' | '/tools-test' @@ -487,6 +515,7 @@ export interface FileRouteTypes { | '/api/chat' | '/api/durable-delivery' | '/api/image' + | '/api/interrupts-v2' | '/api/lazy-tools-wire' | '/api/mcp-apps-call' | '/api/mcp-apps-chat' @@ -513,6 +542,7 @@ export interface FileRouteTypes { | '/$provider/' | '/api/audio/stream' | '/api/image/stream' + | '/api/interrupts-v2/recovery' | '/api/transcription/stream' | '/api/tts/stream' | '/api/video/stream' @@ -526,6 +556,7 @@ export interface FileRouteTypes { | '/devtools-route-b' | '/devtools-structured' | '/devtools-tools' + | '/interrupts-v2' | '/markdown-cjk' | '/middleware-test' | '/tools-test' @@ -538,6 +569,7 @@ export interface FileRouteTypes { | '/api/chat' | '/api/durable-delivery' | '/api/image' + | '/api/interrupts-v2' | '/api/lazy-tools-wire' | '/api/mcp-apps-call' | '/api/mcp-apps-chat' @@ -564,6 +596,7 @@ export interface FileRouteTypes { | '/$provider' | '/api/audio/stream' | '/api/image/stream' + | '/api/interrupts-v2/recovery' | '/api/transcription/stream' | '/api/tts/stream' | '/api/video/stream' @@ -577,6 +610,7 @@ export interface FileRouteTypes { | '/devtools-route-b' | '/devtools-structured' | '/devtools-tools' + | '/interrupts-v2' | '/markdown-cjk' | '/middleware-test' | '/tools-test' @@ -589,6 +623,7 @@ export interface FileRouteTypes { | '/api/chat' | '/api/durable-delivery' | '/api/image' + | '/api/interrupts-v2' | '/api/lazy-tools-wire' | '/api/mcp-apps-call' | '/api/mcp-apps-chat' @@ -615,6 +650,7 @@ export interface FileRouteTypes { | '/$provider/' | '/api/audio/stream' | '/api/image/stream' + | '/api/interrupts-v2/recovery' | '/api/transcription/stream' | '/api/tts/stream' | '/api/video/stream' @@ -629,6 +665,7 @@ export interface RootRouteChildren { DevtoolsRouteBRoute: typeof DevtoolsRouteBRoute DevtoolsStructuredRoute: typeof DevtoolsStructuredRoute DevtoolsToolsRoute: typeof DevtoolsToolsRoute + InterruptsV2Route: typeof InterruptsV2Route MarkdownCjkRoute: typeof MarkdownCjkRoute MiddlewareTestRoute: typeof MiddlewareTestRoute ToolsTestRoute: typeof ToolsTestRoute @@ -641,6 +678,7 @@ export interface RootRouteChildren { ApiChatRoute: typeof ApiChatRoute ApiDurableDeliveryRoute: typeof ApiDurableDeliveryRoute ApiImageRoute: typeof ApiImageRouteWithChildren + ApiInterruptsV2Route: typeof ApiInterruptsV2RouteWithChildren ApiLazyToolsWireRoute: typeof ApiLazyToolsWireRoute ApiMcpAppsCallRoute: typeof ApiMcpAppsCallRoute ApiMcpAppsChatRoute: typeof ApiMcpAppsChatRoute @@ -690,6 +728,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof MarkdownCjkRouteImport parentRoute: typeof rootRouteImport } + '/interrupts-v2': { + id: '/interrupts-v2' + path: '/interrupts-v2' + fullPath: '/interrupts-v2' + preLoaderRoute: typeof InterruptsV2RouteImport + parentRoute: typeof rootRouteImport + } '/devtools-tools': { id: '/devtools-tools' path: '/devtools-tools' @@ -914,6 +959,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiLazyToolsWireRouteImport parentRoute: typeof rootRouteImport } + '/api/interrupts-v2': { + id: '/api/interrupts-v2' + path: '/api/interrupts-v2' + fullPath: '/api/interrupts-v2' + preLoaderRoute: typeof ApiInterruptsV2RouteImport + parentRoute: typeof rootRouteImport + } '/api/image': { id: '/api/image' path: '/api/image' @@ -998,6 +1050,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiTranscriptionStreamRouteImport parentRoute: typeof ApiTranscriptionRoute } + '/api/interrupts-v2/recovery': { + id: '/api/interrupts-v2/recovery' + path: '/recovery' + fullPath: '/api/interrupts-v2/recovery' + preLoaderRoute: typeof ApiInterruptsV2RecoveryRouteImport + parentRoute: typeof ApiInterruptsV2Route + } '/api/image/stream': { id: '/api/image/stream' path: '/stream' @@ -1039,6 +1098,18 @@ const ApiImageRouteWithChildren = ApiImageRoute._addFileChildren( ApiImageRouteChildren, ) +interface ApiInterruptsV2RouteChildren { + ApiInterruptsV2RecoveryRoute: typeof ApiInterruptsV2RecoveryRoute +} + +const ApiInterruptsV2RouteChildren: ApiInterruptsV2RouteChildren = { + ApiInterruptsV2RecoveryRoute: ApiInterruptsV2RecoveryRoute, +} + +const ApiInterruptsV2RouteWithChildren = ApiInterruptsV2Route._addFileChildren( + ApiInterruptsV2RouteChildren, +) + interface ApiTranscriptionRouteChildren { ApiTranscriptionStreamRoute: typeof ApiTranscriptionStreamRoute } @@ -1082,6 +1153,7 @@ const rootRouteChildren: RootRouteChildren = { DevtoolsRouteBRoute: DevtoolsRouteBRoute, DevtoolsStructuredRoute: DevtoolsStructuredRoute, DevtoolsToolsRoute: DevtoolsToolsRoute, + InterruptsV2Route: InterruptsV2Route, MarkdownCjkRoute: MarkdownCjkRoute, MiddlewareTestRoute: MiddlewareTestRoute, ToolsTestRoute: ToolsTestRoute, @@ -1094,6 +1166,7 @@ const rootRouteChildren: RootRouteChildren = { ApiChatRoute: ApiChatRoute, ApiDurableDeliveryRoute: ApiDurableDeliveryRoute, ApiImageRoute: ApiImageRouteWithChildren, + ApiInterruptsV2Route: ApiInterruptsV2RouteWithChildren, ApiLazyToolsWireRoute: ApiLazyToolsWireRoute, ApiMcpAppsCallRoute: ApiMcpAppsCallRoute, ApiMcpAppsChatRoute: ApiMcpAppsChatRoute, diff --git a/testing/e2e/src/routes/api.interrupts-v2.recovery.ts b/testing/e2e/src/routes/api.interrupts-v2.recovery.ts new file mode 100644 index 000000000..f8e94d80f --- /dev/null +++ b/testing/e2e/src/routes/api.interrupts-v2.recovery.ts @@ -0,0 +1,24 @@ +import { createFileRoute } from '@tanstack/react-router' +import { createInterruptRecoveryHandler } from '@tanstack/ai-persistence' +import { getInterruptFixture } from '../lib/interrupts-v2-fixture' + +export const Route = createFileRoute('/api/interrupts-v2/recovery')({ + server: { + handlers: { + POST: async ({ request }) => { + const testId = new URL(request.url).searchParams.get('testId') + if (!testId) return new Response('Missing testId', { status: 400 }) + const fixture = getInterruptFixture(testId) + const handleRecovery = createInterruptRecoveryHandler({ + gateway: fixture.persistence.stores.interrupts, + authorize: (recoveryRequest) => ({ + authorized: + recoveryRequest.headers.get('x-interrupt-fixture') === testId, + includeResolutions: true, + }), + }) + return handleRecovery(request) + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/api.interrupts-v2.ts b/testing/e2e/src/routes/api.interrupts-v2.ts new file mode 100644 index 000000000..f2355c7b5 --- /dev/null +++ b/testing/e2e/src/routes/api.interrupts-v2.ts @@ -0,0 +1,235 @@ +import { + EventType, + chat, + chatParamsFromRequestBody, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { + InterruptResumeValidationError, + withChatPersistence, +} from '@tanstack/ai-persistence' +import { createFileRoute } from '@tanstack/react-router' +import { + createInterruptFixtureAdapter, + fixtureResumeMessages, + getInterruptFixture, + interruptFixtureServerTools, + responseDecision, +} from '../lib/interrupts-v2-fixture' +import type { RunAgentResumeItem, StreamChunk } from '@tanstack/ai' +import type { InterruptFixture } from '../lib/interrupts-v2-fixture' + +async function collectChunks( + stream: AsyncIterable, +): Promise> { + const chunks: Array = [] + for await (const chunk of stream) chunks.push(chunk) + return chunks +} + +function replayChunks( + chunks: ReadonlyArray, +): AsyncIterable { + return { + async *[Symbol.asyncIterator]() { + await Promise.resolve() + for (const chunk of chunks) yield chunk + }, + } +} + +function truncatedSseResponse(): Response { + const encoder = new TextEncoder() + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode('event: RUN_STARTED\ndata: {')) + controller.close() + }, + }), + { headers: { 'Content-Type': 'text/event-stream' } }, + ) +} + +function errorStream( + error: InterruptResumeValidationError, + runId: string, +): AsyncIterable { + return { + async *[Symbol.asyncIterator]() { + await Promise.resolve() + const base: StreamChunk = { + type: EventType.RUN_ERROR, + runId, + message: error.message, + code: 'INTERRUPT_RESUME_VALIDATION', + timestamp: Date.now(), + } + yield Object.assign(base, { + 'tanstack:interruptErrors': error.errors, + ...(error.recovery === undefined + ? {} + : { 'tanstack:interruptRecovery': error.recovery }), + }) + }, + } +} + +function objectValue(value: unknown): Record | undefined { + return value !== null && typeof value === 'object' && !Array.isArray(value) + ? Object.fromEntries(Object.entries(value)) + : undefined +} + +function recordSubmission( + fixture: InterruptFixture, + resume: ReadonlyArray, + chunks: ReadonlyArray, + runId: string, +): void { + fixture.continuationCount += 1 + fixture.continuationRunIds.push(runId) + fixture.continuationChunks.set(runId, chunks) + fixture.decisions = resume.map(responseDecision) + fixture.auditHistory = resume.map( + (item) => `${item.interruptId}:${responseDecision(item)}`, + ) + fixture.resultEventNames = chunks + .map((chunk) => chunk.type) + .filter( + (type) => + type === EventType.RUN_STARTED || + type === EventType.TOOL_CALL_RESULT || + type === EventType.RUN_FINISHED, + ) + fixture.storedHistory.push(fixture.decisions.join(',')) + + for (const item of resume) { + const payload = objectValue(item.payload) + const editedArgs = payload?.editedArgs + const edited = objectValue(editedArgs) + if (edited) fixture.edits = edited + } +} + +function replayedContinuationId( + chunks: ReadonlyArray, +): string | undefined { + for (const chunk of chunks) { + if (chunk.type !== EventType.RUN_FINISHED || !('result' in chunk)) continue + const result = objectValue(chunk.result) + if (result?.replayed !== true) continue + const continuationRunId = result.continuationRunId + if (typeof continuationRunId === 'string') return continuationRunId + } + return undefined +} + +function hasInterruptFailure(chunks: ReadonlyArray): boolean { + return chunks.some( + (chunk) => + chunk.type === EventType.RUN_ERROR && + (chunk['tanstack:interruptErrors']?.length ?? 0) > 0, + ) +} + +export const Route = createFileRoute('/api/interrupts-v2')({ + server: { + handlers: { + GET: ({ request }) => { + const url = new URL(request.url) + const testId = url.searchParams.get('testId') + if (!testId) return new Response('Missing testId', { status: 400 }) + const fixture = getInterruptFixture(testId) + + if (url.searchParams.get('stats') === '1') { + return Response.json({ + continuationCount: fixture.continuationCount, + continuationRunIds: fixture.continuationRunIds, + decisions: fixture.decisions, + edits: fixture.edits, + auditHistory: fixture.auditHistory, + resultEventNames: fixture.resultEventNames, + storedHistory: fixture.storedHistory, + truncatedResponses: fixture.truncatedResponses, + replayCount: fixture.replayCount, + joinedContinuationRunId: fixture.joinedContinuationRunId, + }) + } + + const runId = url.searchParams.get('runId') + if (!runId) return new Response('Missing runId', { status: 400 }) + const chunks = fixture.continuationChunks.get(runId) + if (!chunks) return new Response('Unknown run', { status: 404 }) + fixture.replayCount += 1 + fixture.joinedContinuationRunId = runId + return toServerSentEventsResponse(replayChunks(chunks)) + }, + POST: async ({ request }) => { + const url = new URL(request.url) + const testId = url.searchParams.get('testId') + const scenario = url.searchParams.get('scenario') + if (!testId || !scenario) { + return new Response('Missing fixture correlation', { status: 400 }) + } + + let params + try { + params = await chatParamsFromRequestBody(await request.json()) + } catch (error) { + return new Response( + error instanceof Error ? error.message : 'Bad request', + { status: 400 }, + ) + } + + const fixture = getInterruptFixture(testId) + const resumeMessages = + params.resume === undefined + ? undefined + : fixtureResumeMessages(scenario) + const stream = chat({ + ...params, + ...(resumeMessages === undefined ? {} : { messages: resumeMessages }), + adapter: createInterruptFixtureAdapter( + scenario, + params.resume !== undefined, + ), + tools: [...interruptFixtureServerTools], + middleware: [withChatPersistence(fixture.persistence)], + }) + + let chunks: Array + try { + chunks = await collectChunks(stream) + } catch (error) { + if (error instanceof InterruptResumeValidationError) { + return toServerSentEventsResponse(errorStream(error, params.runId)) + } + throw error + } + + if (params.resume !== undefined) { + const replayedId = replayedContinuationId(chunks) + if (replayedId === undefined && !hasInterruptFailure(chunks)) { + recordSubmission(fixture, params.resume, chunks, params.runId) + } else { + fixture.joinedContinuationRunId = replayedId + } + } + + if ( + scenario === 'commit-then-truncate' && + params.resume !== undefined && + fixture.truncateCommittedResponseOnce + ) { + fixture.truncateCommittedResponseOnce = false + fixture.truncatedResponses += 1 + return truncatedSseResponse() + } + + return toServerSentEventsResponse(replayChunks(chunks)) + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/interrupts-v2.tsx b/testing/e2e/src/routes/interrupts-v2.tsx new file mode 100644 index 000000000..1ad14e833 --- /dev/null +++ b/testing/e2e/src/routes/interrupts-v2.tsx @@ -0,0 +1,552 @@ +import { createFileRoute } from '@tanstack/react-router' +import { + fetchServerSentEvents, + localStoragePersistence, +} from '@tanstack/ai-client' +import { useChat } from '@tanstack/ai-react' +import { useCallback, useEffect, useMemo, useState } from 'react' +import { interruptFixtureTools } from '../lib/interrupts-v2-fixture' +import type { + InterruptRecoveryQuery, + InterruptRecoveryStateV1, +} from '@tanstack/ai' +import type { + ChatResumeSnapshot, + ConnectConnectionAdapter, +} from '@tanstack/ai-client' + +interface FixtureStats { + continuationCount: number + decisions: Array + edits: Record + auditHistory: Array + resultEventNames: Array + storedHistory: Array +} + +interface LocalDrafts { + action: string + generic: string + stagedFirst: string + clientOutput: string + addToolResultCount: number +} + +const emptyStats: FixtureStats = { + continuationCount: 0, + decisions: [], + edits: {}, + auditHistory: [], + resultEventNames: [], + storedHistory: [], +} + +const emptyDrafts: LocalDrafts = { + action: '', + generic: '', + stagedFirst: '', + clientOutput: '', + addToolResultCount: 0, +} + +function createInterruptStateFetcher( + url: () => string, + authorization: string, +): NonNullable { + return async ( + query: InterruptRecoveryQuery, + abortSignal?: AbortSignal, + ): Promise => { + const recoveryUrl = new URL(url(), window.location.origin) + recoveryUrl.searchParams.set('threadId', query.threadId) + recoveryUrl.searchParams.set('interruptedRunId', query.interruptedRunId) + recoveryUrl.searchParams.set( + 'knownGeneration', + String(query.knownGeneration), + ) + const response = await fetch(recoveryUrl, { + method: 'POST', + headers: { 'x-interrupt-fixture': authorization }, + ...(abortSignal === undefined ? {} : { signal: abortSignal }), + }) + if (!response.ok) { + throw new Error(`Interrupt recovery failed with ${response.status}.`) + } + return response.json() + } +} + +function readDrafts(key: string): LocalDrafts { + if (typeof window === 'undefined') return emptyDrafts + const value = window.localStorage.getItem(key) + if (!value) return emptyDrafts + try { + const parsed: unknown = JSON.parse(value) + if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) { + return emptyDrafts + } + const record = Object.fromEntries(Object.entries(parsed)) + return { + action: typeof record.action === 'string' ? record.action : '', + generic: typeof record.generic === 'string' ? record.generic : '', + stagedFirst: + typeof record.stagedFirst === 'string' ? record.stagedFirst : '', + clientOutput: + typeof record.clientOutput === 'string' ? record.clientOutput : '', + addToolResultCount: + typeof record.addToolResultCount === 'number' + ? record.addToolResultCount + : 0, + } + } catch { + return emptyDrafts + } +} + +function InterruptsV2Page() { + const { testId, scenario } = Route.useSearch() + const draftKey = `interrupts-v2:drafts:${testId}` + const [drafts, setDrafts] = useState(() => readDrafts(draftKey)) + const [stats, setStats] = useState(emptyStats) + const [callbackReturns, setCallbackReturns] = useState('') + const [retryVisible, setRetryVisible] = useState(false) + const [hydrated, setHydrated] = useState(false) + + const updateDrafts = useCallback( + (update: Partial) => { + setDrafts((current) => { + const next = { ...current, ...update } + if (typeof window !== 'undefined') { + window.localStorage.setItem(draftKey, JSON.stringify(next)) + } + return next + }) + }, + [draftKey], + ) + + const refreshStats = useCallback(async () => { + if (typeof window === 'undefined') return + const response = await fetch( + new URL( + `/api/interrupts-v2?testId=${encodeURIComponent(testId)}&stats=1`, + window.location.origin, + ), + ) + if (response.ok) setStats(await response.json()) + }, [testId]) + + const connection = useMemo(() => { + const chatUrl = () => + `/api/interrupts-v2?testId=${encodeURIComponent(testId)}&scenario=${encodeURIComponent(scenario)}` + const recoveryUrl = () => + `/api/interrupts-v2/recovery?testId=${encodeURIComponent(testId)}` + return { + ...fetchServerSentEvents(chatUrl), + loadInterruptState: createInterruptStateFetcher(recoveryUrl, testId), + } + }, [scenario, testId]) + + const resumePersistence = useMemo( + () => + typeof window === 'undefined' + ? undefined + : localStoragePersistence({ + keyPrefix: 'tanstack-ai:interrupts-v2:', + serialize: (value) => JSON.stringify(value), + deserialize: (value) => JSON.parse(value), + }), + [], + ) + + const { + sendMessage, + interrupts, + interruptErrors, + error, + isLoading, + resuming, + resolveInterrupts, + retryInterrupts, + addToolResult, + } = useChat({ + id: `interrupts-v2-${testId}`, + threadId: `interrupts-v2-${testId}`, + connection, + tools: interruptFixtureTools, + ...(resumePersistence === undefined + ? {} + : { persistence: { server: resumePersistence } }), + onChunk: (chunk) => { + if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') { + void refreshStats() + } + }, + onError: () => { + setRetryVisible(true) + void refreshStats() + }, + }) + + useEffect(() => { + void refreshStats() + }, [isLoading, refreshStats, resuming, interrupts.length]) + + useEffect(() => setHydrated(true), []) + + const first = interrupts.find( + (interrupt) => interrupt.interruptId === 'approval-1', + ) + const second = interrupts.find( + (interrupt) => interrupt.interruptId === 'question-1', + ) + const clientTool = interrupts.find( + (interrupt) => interrupt.kind === 'client-tool-execution', + ) + const invalidItemIds = interrupts + .filter((interrupt) => interrupt.error !== undefined) + .map((interrupt) => interrupt.interruptId) + const validationAggregate = + invalidItemIds.length === 0 + ? '' + : `item-validation-failed:${invalidItemIds.join(',')}` + + const approveFirst = useCallback(() => { + if (first?.kind !== 'tool-approval') return + const editedArgs = drafts.action + ? { editedArgs: { action: drafts.action } } + : {} + if (first.toolName === 'editable_action') { + first.resolveInterrupt(true, { payload: {}, ...editedArgs }) + updateDrafts({ + stagedFirst: JSON.stringify({ approved: true, ...editedArgs }), + }) + } + }, [drafts.action, first, updateDrafts]) + + const parseGenericDraft = useCallback((): unknown => { + try { + return JSON.parse(drafts.generic) + } catch { + return drafts.generic + } + }, [drafts.generic]) + + return ( +
+

Interrupts v2 fixture

+ + + + {stats.continuationCount} + + + {stats.decisions.join(',')} + + + {JSON.stringify(stats.edits)} + + + {stats.auditHistory.join('|')} + + + {stats.resultEventNames.join(',')} + + + {stats.storedHistory.join('|')} + + {callbackReturns} + {drafts.stagedFirst} + {drafts.clientOutput} + + {drafts.addToolResultCount} + + + {clientTool?.status ?? ''} + + + + {first?.error ? `${first.error.code}:${first.error.message}` : ''} + + + {second?.error ? `${second.error.code}:${second.error.message}` : ''} + + + {second?.error ? `${second.error.code}:${second.error.message}` : ''} + + + {interruptErrors + .map( + (submissionError) => + `${submissionError.code}:${submissionError.interruptIds.join(',')}:${submissionError.message}`, + ) + .join('|')} + {interruptErrors.length > 0 && validationAggregate ? '|' : ''} + {validationAggregate} + + {error?.message ?? ''} + + {retryVisible ? ( +
Retry available
+ ) : null} + + + {interrupts.map((interrupt) => ( +
+ {interrupt.interruptId} + + {interrupt.kind} + + {interrupt.kind === 'tool-approval' ? ( + <> + {interrupt.toolName === 'editable_action' ? ( + <> + + updateDrafts({ action: event.currentTarget.value }) + } + /> + + + + ) : interrupt.toolName === 'branch_action' ? ( + <> + + + + ) : ( + <> + + + + )} + + + ) : null} + {interrupt.kind === 'generic' ? ( + <> +