diff --git a/.changeset/generic-interrupts.md b/.changeset/generic-interrupts.md new file mode 100644 index 0000000000..ddcd3d62bd --- /dev/null +++ b/.changeset/generic-interrupts.md @@ -0,0 +1,17 @@ +--- +'@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': minor +--- + +Add first-party generic interrupts. + +Use `defineInterrupt()` to describe a pause, register it on `chat()` and the client hooks, and return requests from `onInterruptBoundary`. The client gets typed payloads and `resolveInterrupt`. Resume validates the answer and runs `onInterruptResolution`. + +`GenericInterrupt` types one bound card. `INTERRUPT_BOUNDARY_PHASES` and `INTERRUPT_TOOL_RESUMES` are the shared phase and resume lists. diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml index c65aac8bc7..46195d5832 100644 --- a/.github/workflows/e2e.yml +++ b/.github/workflows/e2e.yml @@ -19,7 +19,7 @@ jobs: e2e: name: E2E Tests runs-on: ubuntu-latest - timeout-minutes: 15 + timeout-minutes: 30 steps: - name: Checkout uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 diff --git a/docs/advanced/middleware.md b/docs/advanced/middleware.md index 80fe5fed46..07d381f929 100644 --- a/docs/advanced/middleware.md +++ b/docs/advanced/middleware.md @@ -332,6 +332,180 @@ const budget: ChatMiddleware = { For a full per-turn + cumulative tool budget recipe, see [Tool-call budgets](../chat/agentic-cycle#tool-call-budgets-middleware-recipe). +### onInterruptBoundary and onInterruptResolution + +Use these hooks when middleware needs data from the client. Define the request +with `defineInterrupt()` and register it with `chat({ interrupts })` and +`useChat({ interrupts })`. Do not emit raw AG-UI events from middleware. + +`onInterruptBoundary` runs at four points in an agent iteration: + +- `beforeModel`, before the adapter starts. +- `afterModel`, after the model response is complete. +- `beforeTools`, before tool execution starts. +- `afterTools`, after the tool phase is complete. + +Each middleware can return requests from one boundary. The engine combines all +requests from that boundary into one AG-UI interrupt batch. The batch ends the +run with one interrupt outcome. + +This hook cannot change config. Its only legal return is `{ interrupts }` or +nothing. The continuation is a new `chat()` call, so the hook runs again. Skip +the emit when `ctx.parentRunId` is set if this pause belongs to the original +request only. + +What is in `ctx` at each phase, and when to use each phase, is in +[Lifecycle Boundaries](../interrupts/boundaries). + +Create one shared definition. Both the server and the client import this value, +so the definition ID and response shape stay the same on both sides. + +```typescript title="review-plan.ts" +import { defineInterrupt, type ChatMiddleware } from '@tanstack/ai' +import { z } from 'zod' + +export const reviewPlan = defineInterrupt({ + id: 'review-plan', + payloadSchema: z.object({ title: z.string() }), + responseSchema: z.object({ approved: z.boolean() }), +}) + +export const reviewMiddleware: ChatMiddleware = { + name: 'review-plan', + onInterruptBoundary(ctx) { + if (ctx.phase !== 'beforeTools') return + if (ctx.parentRunId) return + return { + interrupts: [ + reviewPlan.interrupt({ + key: 'release-plan', + reason: 'review-required', + message: 'Approve this plan?', + payload: { title: 'Release plan' }, + }), + ], + } + }, + onInterruptResolution(_ctx, resumedInterrupts) { + for (const result of resumedInterrupts.for(reviewPlan)) { + if (result.status === 'resolved' && !result.response.approved) { + return { toolResume: 'stop' } + } + } + }, +} +``` + +Register the definition on the server. Forward `parentRunId` and `resume` so +a client resolution starts the continuation with its full context. + +```typescript title="route.ts" +import { + chat, + chatParamsFromRequestBody, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { openaiText } from '@tanstack/ai-openai' +import { reviewMiddleware, reviewPlan } from './review-plan' + +export async function POST(request: Request) { + const params = await chatParamsFromRequestBody(await request.json()) + const stream = chat({ + adapter: openaiText('gpt-5.5'), + messages: params.messages, + threadId: params.threadId, + runId: params.runId, + ...(params.parentRunId ? { parentRunId: params.parentRunId } : {}), + ...(params.resume ? { resume: params.resume } : {}), + interrupts: [reviewPlan], + middleware: [reviewMiddleware], + }) + + return toServerSentEventsResponse(stream) +} +``` + +Register the same definition on the client. Check `kind` and `definitionId`. +TypeScript then treats the item as `GenericInterrupt`. +`resolveInterrupt` uses the response shape from `reviewPlan.responseSchema`. + +```tsx title="review-plan-panel.tsx" +import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' +import type { GenericInterrupt } from '@tanstack/ai-react' +import { reviewPlan } from './review-plan' + +function ReviewCard({ + interrupt, +}: { + interrupt: GenericInterrupt +}) { + return ( + + ) +} + +export function ReviewPlanPanel() { + const { interrupts, sendMessage } = useChat({ + connection: fetchServerSentEvents('/api/chat'), + interrupts: [reviewPlan], + }) + + return ( + <> + + {interrupts.map((interrupt) => { + if (interrupt.kind !== 'generic') return null + if (!('definitionId' in interrupt)) return null + if (interrupt.definitionId !== reviewPlan.id) return null + return + })} + + ) +} +``` + +`onInterruptResolution` does not run in the `chat()` call that paused. It +runs once at the start of the next `chat()` call, after the client answers. + +``` +setup +onConfig (phase is init) +onInterruptResolution (phase is still init) +onStart +then stop, or continue the agent loop +``` + +`useChat` sends `parentRunId` and `resume` on that second request. Each +generic resume item includes the original request in `metadata`. If `resume` +is present and `parentRunId` is missing, the server throws. + +Use `resumedInterrupts.for(definition)` for one typed definition. Use +`resumedInterrupts.all()` for every registered definition. Use +`resumedInterrupts.all(definitionA, definitionB)` to read a typed subset. + +The hook can return `toolResume: 'continue'`, `'cancel'`, or `'stop'`. Results +from all middleware combine by the most restrictive rule: `stop` wins over +`cancel`, and `cancel` wins over `continue`. + +This hook cannot change prompts, tools, or messages. Store the answer on a +capability, then return those fields from `onConfig` when +`ctx.phase === 'beforeModel'`. + +| Hook | Can change | +| --- | --- | +| `onInterruptBoundary` | Nothing. It can only pause. | +| `onInterruptResolution` | Pending-tool policy (`toolResume`) | +| `onConfig` | `messages`, `systemPrompts`, `tools`, `modelOptions`, `metadata` | + +The full resume order, plus an example that writes a user note into the +system prompt, is in [Apply Answers](../interrupts/apply-answers). + ### onBeforeToolCall Called before each tool executes. The first middleware that returns a non-void decision short-circuits — remaining middleware are skipped for that tool call. @@ -721,9 +895,32 @@ If you drop `withCounter` from the array, `chat()` reports a compile-time error `createChatMiddleware()` builds the array through chained `.use()` calls and enforces **provider-before-consumer ordering at compile time**: each `.use()` requires that the middleware's `requires` are already covered by capabilities provided by earlier `.use()` calls. ```typescript -import { chat, createChatMiddleware } from "@tanstack/ai"; +import { + chat, + createCapability, + createChatMiddleware, + defineChatMiddleware, +} from "@tanstack/ai"; import { openaiText } from "@tanstack/ai-openai"; -import { withCounter, countsChunks } from "./counter-middleware"; + +const counterCapability = createCapability<{ value: number }>()("counter"); +const [getCounter, provideCounter] = counterCapability; + +const withCounter = defineChatMiddleware({ + name: "with-counter", + provides: [counterCapability], + setup(ctx) { + provideCounter(ctx, { value: 0 }); + }, +}); + +const countsChunks = defineChatMiddleware({ + name: "counts-chunks", + requires: [counterCapability], + onChunk(ctx) { + getCounter(ctx).value++; + }, +}); const middleware = createChatMiddleware() .use(withCounter) // provides "counter" diff --git a/docs/config.json b/docs/config.json index e1a5aa44b6..7d13be1762 100644 --- a/docs/config.json +++ b/docs/config.json @@ -188,27 +188,43 @@ { "label": "Overview", "to": "interrupts/overview", - "addedAt": "2026-08-04" + "addedAt": "2026-08-04", + "updatedAt": "2026-08-14" }, { "label": "Tool Approval", "to": "interrupts/tool-approval", - "addedAt": "2026-08-04" + "addedAt": "2026-08-04", + "updatedAt": "2026-08-14" }, { "label": "Multiple Interrupts", "to": "interrupts/multiple", - "addedAt": "2026-08-04" + "addedAt": "2026-08-04", + "updatedAt": "2026-08-14" }, { "label": "Generic Interrupts", "to": "interrupts/generic", - "addedAt": "2026-08-04" + "addedAt": "2026-08-04", + "updatedAt": "2026-08-14" + }, + { + "label": "Lifecycle Boundaries", + "to": "interrupts/boundaries", + "addedAt": "2026-08-13" + }, + { + "label": "Apply Answers", + "to": "interrupts/apply-answers", + "addedAt": "2026-08-13", + "updatedAt": "2026-08-14" }, { "label": "Migration", "to": "interrupts/migration", - "addedAt": "2026-08-04" + "addedAt": "2026-08-04", + "updatedAt": "2026-08-14" } ] }, @@ -245,7 +261,8 @@ { "label": "Chat Persistence", "to": "persistence/chat-persistence", - "addedAt": "2026-08-04" + "addedAt": "2026-08-04", + "updatedAt": "2026-08-13" }, { "label": "Client Persistence", @@ -308,7 +325,8 @@ { "label": "Store Reference", "to": "persistence/store-reference", - "addedAt": "2026-08-04" + "addedAt": "2026-08-04", + "updatedAt": "2026-08-14" }, { "label": "How Persistence Works", @@ -480,7 +498,7 @@ "label": "Middleware", "to": "advanced/middleware", "addedAt": "2026-04-15", - "updatedAt": "2026-07-21" + "updatedAt": "2026-08-14" }, { "label": "Built-in Middleware", diff --git a/docs/interrupts/apply-answers.md b/docs/interrupts/apply-answers.md new file mode 100644 index 0000000000..fafc52852e --- /dev/null +++ b/docs/interrupts/apply-answers.md @@ -0,0 +1,311 @@ +--- +title: Apply Answers +id: interrupts-apply-answers +order: 6 +description: "Read a generic interrupt answer on the continuation run and apply it to the next model turn." +keywords: + - tanstack ai + - onInterruptResolution + - toolResume + - generic interrupt + - parentRunId + - createCapability +--- + +# Apply Answers + +The user approved a plan or typed a note. You need that value on the server +before the next model call. `onInterruptResolution` is where you read it. It +does not change config. `onConfig` is where you apply it. + +By the end of this page you know when the resolution hook runs, what it can +return, and how to turn the answer into a prompt, a tool list, or a stop. + +Define and emit the interrupt first. See [Generic Interrupts](./generic). + +## Two `chat()` calls + +The pause spans two runs. One user-visible turn. Two `chat()` calls. + +**Call 1 (pause).** `onInterruptBoundary` returns `{ interrupts }`. The run +ends with `RUN_FINISHED` and `outcome: interrupt`. The resolution hook does +not run. + +**Call 2 (resume).** The client starts a new request after +`resolveInterrupt()` or `cancel()`. The body includes: + +- a new `runId` +- `parentRunId` set to the paused run +- `resume` with the answers. Each generic item also has `metadata` with the + original request (`tanstack:interruptContinuation`) + +`useChat` sends those fields for you. If you POST by hand, include all three. +If `resume` is present and `parentRunId` is missing, the server throws. + +A hand-built generic resume item looks like this: + +```ts +import { wrapGenericInterruptContinuation } from '@tanstack/ai' + +const resumeItem = { + interruptId: 'generic-1', + status: 'resolved' as const, + payload: { approved: true }, + metadata: wrapGenericInterruptContinuation({ + v: 1, + definitionId: 'review-plan', + key: 'turn-1', + batchIndex: 0, + reason: 'review', + message: 'Review the plan', + }), +} +``` + +```mermaid +sequenceDiagram + participant User + participant Client + participant Server + + Client->>Server: first chat() request + Server-->>Client: RUN_FINISHED outcome interrupt + Client->>User: interrupts array + User->>Client: resolveInterrupt or cancel + Client->>Server: second chat() with parentRunId and resume + Note over Server: onInterruptResolution runs here + Server-->>Client: continue, cancel tools, or stop +``` + +## Exact place in the second call + +``` +setup +onConfig (phase is init) +onInterruptResolution (phase is still init) +onStart +then stop, or continue the agent loop +``` + +The hook runs **once**, at the start of the continuation. It does not run at +`beforeModel`, `afterModel`, `beforeTools`, or `afterTools`. + +`ctx.phase` is `'init'`. `ctx.iteration` is `0`. No model call has started. + +The hook runs once per batch. Two cards in one pause still produce one hook +call with both answers. + +The hook does not run on the first user message. It does not run for a +tool-approval or client-tool batch that has no generic interrupt. + +If that continuation pauses again, that is a third `chat()` call. The hook +runs again at the start of that third call. + +## Read the typed answers + +```ts +import type { ChatMiddleware } from '@tanstack/ai' +import { reviewPlan } from './interrupts' + +export const applyReview: ChatMiddleware = { + name: 'apply-review', + onInterruptResolution(_ctx, resumedInterrupts) { + for (const result of resumedInterrupts.for(reviewPlan)) { + if (result.status === 'resolved' && !result.response.approved) { + return { toolResume: 'stop' } + } + } + }, +} +``` + +- `resumedInterrupts.for(reviewPlan)` keeps the response type for that + definition +- `resumedInterrupts.all()` returns every registered answer +- `resumedInterrupts.all(reviewPlan, otherDefinition)` narrows to those + definitions + +Each item is `resolved` (with `response`) or `cancelled`. + +## What the hook can return + +Return `toolResume` to decide what happens to **pending tools** from the +paused turn: + +| Value | Effect | +| --- | --- | +| `continue` | Run the pending tools | +| `cancel` | Mark the pending tools as cancelled. Do not run them | +| `stop` | End the run after `onStart`. No tools. No model call | + +When more than one middleware returns a value, the engine keeps the stricter +one. `stop` wins over `cancel`. `cancel` wins over `continue`. + +If a generic interrupt shares a batch with client tools, the client does not +run those client tools until `toolResume` is `continue`. `cancel` and `stop` +skip them. After `continue`, the engine emits the client-tool wait again. + +After a `continue`, the engine picks up from the phase that paused: + +| First run paused at | Next step | +| --- | --- | +| `beforeModel` | The model call | +| `afterModel` | Tools, if the model asked for them | +| `beforeTools` | Tool execution | +| `afterTools` | The next model turn (the tools already ran) | + +## What the hook cannot change + +`onInterruptResolution` cannot change config. It cannot change the model or +the adapter. + +These fields change only from `onConfig` (or `onStructuredOutputConfig`): + +- `messages` +- `systemPrompts` +- `tools` +- `modelOptions` +- `metadata` + +The user answer is **not** added to `messages` by itself. If the model must +see the note, you add it. + +## Apply the answer in `onConfig` + +Order on the continuation: + +1. `onConfig` with `phase: 'init'`. The answers are not applied yet. +2. `onInterruptResolution`. Read the answers and store them on the run. +3. `onConfig` with `phase: 'beforeModel'`. Return the new prompts, tools, or + messages. + +Store the answer in a [middleware capability](../advanced/middleware#capabilities). +The value lives on `ctx` for this `chat()` call. Another middleware can +declare `requires` and read the same note. Two overlapping `chat()` calls do +not share the value. + +If you list the capability in `provides`, you must provide it in `setup`. +You do not have the user answer yet. Provide an empty box first. Write the +answer in `onInterruptResolution`. + +```ts +import { createCapability, type ChatMiddleware } from '@tanstack/ai' +import { reviewPlan } from './interrupts' + +export const reviewNote = createCapability<{ note?: string }>()('review-note') +export const [getReviewNote, provideReviewNote] = reviewNote + +export const reviewMiddleware: ChatMiddleware = { + name: 'review-plan', + provides: [reviewNote], + setup(ctx) { + provideReviewNote(ctx, {}) + }, + onInterruptBoundary(ctx) { + if (ctx.phase !== 'beforeModel') return + if (ctx.parentRunId) return + return { + interrupts: [ + reviewPlan.interrupt({ + key: 'initial-plan', + reason: 'review-required', + message: 'Review the proposed plan.', + payload: { + title: 'Release plan', + changes: ['Add search'], + }, + }), + ], + } + }, + onInterruptResolution(ctx, resumed) { + const [result] = resumed.for(reviewPlan) + if (result?.status !== 'resolved') return + provideReviewNote(ctx, { note: result.response.note }) + if (!result.response.approved) { + return { toolResume: 'stop' } + } + }, + onConfig(ctx, config) { + if (ctx.phase !== 'beforeModel') return + const note = getReviewNote(ctx).note + if (!note) return + return { + systemPrompts: [ + ...config.systemPrompts, + `User review note: ${note}`, + ], + } + }, +} +``` + +Register the object on `chat({ middleware: [reviewMiddleware] })`. + +A later middleware can read the same note: + +```ts +import { type ChatMiddleware } from '@tanstack/ai' +import { getReviewNote, reviewNote } from './review-plan' + +export const applyVoice: ChatMiddleware = { + name: 'apply-voice', + requires: [reviewNote], + onConfig(ctx, config) { + const note = getReviewNote(ctx).note + if (!note) return + return { + systemPrompts: [...config.systemPrompts, `Voice note: ${note}`], + } + }, +} +``` + +Other hooks that can change behavior, but not this resume payload: + +- `onBeforeToolCall` can rewrite args, skip a tool, or abort +- `onChunk` can rewrite or drop stream events +- `onShouldContinue` can return `false` to stop the loop with a normal finish + +See [Middleware](../advanced/middleware) for those hooks. + +## Persistence + +With [chat persistence](../persistence/chat-persistence), the hook still runs +at the same moment: after init `onConfig`, before `onStart`. + +Persistence rebuilds the pending requests from the store and clears +`config.resume` so the engine does not rebuild them from client history. You +still read answers from `resumedInterrupts.for(definition)`. + +## Register both sides + +The continuation needs the same definitions and the same middleware as the +paused run. Forward `parentRunId` and `resume` from the request. + +```ts +// app/api/chat/route.ts +import { + chat, + chatParamsFromRequest, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { openaiText } from '@tanstack/ai-openai' +import { reviewMiddleware } from '../../chat-middleware' +import { reviewPlan } from '../../interrupts' + +export async function POST(request: Request) { + const params = await chatParamsFromRequest(request) + const stream = chat({ + adapter: openaiText('gpt-5.5'), + messages: params.messages, + threadId: params.threadId, + runId: params.runId, + ...(params.parentRunId ? { parentRunId: params.parentRunId } : {}), + ...(params.resume ? { resume: params.resume } : {}), + interrupts: [reviewPlan], + middleware: [reviewMiddleware], + }) + return toServerSentEventsResponse(stream) +} +``` diff --git a/docs/interrupts/boundaries.md b/docs/interrupts/boundaries.md new file mode 100644 index 0000000000..e4efb499c5 --- /dev/null +++ b/docs/interrupts/boundaries.md @@ -0,0 +1,212 @@ +--- +title: Lifecycle Boundaries +id: interrupts-boundaries +order: 5 +description: "Pick the chat lifecycle phase where a generic interrupt pauses the run." +keywords: + - tanstack ai + - generic interrupt + - onInterruptBoundary + - INTERRUPT_BOUNDARY_PHASES + - beforeModel + - afterModel + - beforeTools + - afterTools +--- + +# Lifecycle Boundaries + +You know you need a generic interrupt. You do not know which phase to pause +in. If you pause too early, the model has no draft to review. If you pause too +late, a tool has already run. + +By the end of this page you can pick one phase from +`INTERRUPT_BOUNDARY_PHASES` and write the `ctx.phase` guard for it. + +For the define, register, and resolve steps, see +[Generic Interrupts](./generic). + +## The four phases + +`onInterruptBoundary` runs at each of these points in one agent iteration. +Return `{ interrupts }` to pause. Return nothing to let the run continue. + +| Phase | When it runs | Typical question | +| --- | --- | --- | +| `beforeModel` | After `onConfig` for this iteration, before the adapter call | Do we have enough from the user to spend tokens? | +| `afterModel` | After the model stream ends, before tools run | Is this draft or these tool calls acceptable? | +| `beforeTools` | After the assistant tool-call message is in `messages`, before execution | May these tools run? | +| `afterTools` | After tools finish and their result messages are in `messages` | May these results go back to the model? | + +The engine combines every request from every middleware at the same phase into +one interrupt batch. That batch ends the current run with one `interrupt` +outcome. + +```ts +import type { ChatMiddleware } from '@tanstack/ai' +import { reviewPlan } from './interrupts' + +export const requestReview: ChatMiddleware = { + name: 'request-review', + onInterruptBoundary(ctx) { + if (ctx.phase !== 'beforeModel') return + if (ctx.parentRunId) return + if (ctx.iteration !== 0) return + return { + interrupts: [ + reviewPlan.interrupt({ + key: 'initial-plan', + reason: 'review-required', + message: 'Review the proposed plan.', + payload: { + title: 'Release plan', + changes: ['Add search', 'Add tests'], + }, + }), + ], + } + }, +} +``` + +`onInterruptBoundary` cannot change config. It can only pause. To change +prompts, tools, or messages from the user answer, see +[Apply Answers](./apply-answers). + +## Skip the pause on the continuation + +The continuation is a new `chat()` call. Every boundary hook runs again. + +If you return the same request, the run pauses again. If `ctx.parentRunId` is +set, skip the emit. Use that skip when the pause belongs to the original +request only. + +```ts ignore +onInterruptBoundary(ctx) { + if (ctx.phase !== 'beforeModel') return + if (ctx.parentRunId) return + return { + interrupts: [ + reviewPlan.interrupt({ + key: 'initial-plan', + reason: 'review-required', + message: 'Review the proposed plan.', + payload: { + title: 'Release plan', + changes: ['Add search'], + }, + }), + ], + } +} +``` + +## What you can read at each phase + +Every hook receives the same `ChatMiddlewareContext`. The **contents** change. + +Useful fields on `ctx`: + +- `phase` and `iteration` +- `parentRunId` (set on a continuation) +- `messages` (read-only view) +- `systemPrompts`, `toolNames`, `hasTools`, `modelOptions` +- `accumulatedContent` (assistant text for this model turn) +- `model` and `provider` (fixed for this request) +- `context` (your `chat({ context })` value) +- `abort(reason)` and `defer(promise)` + +Mutating `ctx.messages` does not change the engine config. + +### `beforeModel` + +`onConfig` for this iteration has already run. Prompts, tools, messages, and +`modelOptions` are the values after that merge. + +- `accumulatedContent` is empty +- `messages` is the conversation so far + +Use this phase when you need data **before** you pay for a model call. + +### `afterModel` + +The model stream is complete. + +- `accumulatedContent` has this turn's assistant text +- `messages` does **not** include this turn yet +- Proposed tool calls are not on `ctx`. They are not in `messages` yet + +If you need tool names or args, wait for `beforeTools`, or watch `onChunk` +during `modelStream`. + +### `beforeTools` + +The engine has added the assistant message with `toolCalls` to `messages`. +Tools have not run. + +Use this phase to inspect the proposed calls before any side effect. + +### `afterTools` + +Tools have run. Result messages with `role: 'tool'` are already in +`messages`. `onToolPhaseComplete` has already run. + +Use this phase to inspect results before the next model turn. + +## Real uses + +### `beforeModel`: collect a choice first + +Ask for a plan, an audience, or a locale before the model writes. + +Examples: + +- "Which brand voice should this reply use?" +- "Which ticket should I work on?" +- "Is this request in scope for this agent?" + +### `afterModel`: review the draft + +The model has written text. You want a human to accept it before tools run. + +Examples: + +- Review an email draft before `sendEmail` +- Review a SQL query before `runQuery` +- Review a support reply before it reaches the user + +### `beforeTools`: gate the side effects + +The model asked for tools. Nothing has executed. + +Examples: + +- Confirm a bulk delete +- Confirm a payment +- Confirm a deploy + +This is close to [tool approval](./tool-approval). Use a generic interrupt +when the question is not a yes or no on one tool. Also use it when several +tools must be judged as one batch. + +### `afterTools`: audit the results + +The tools have already run. You want a human to see the output before the +model uses it. + +Examples: + +- A search returned customer PII. Ask if it may stay in context +- A code run produced a diff. Ask if the next turn may apply it +- A lookup returned a low-confidence match. Ask which record to keep + +## Try it + +The React chat example has a playground for all four phases. + +1. Start `examples/ts-react-chat`. +2. Open `/generic-interrupts`. +3. Pick a phase and resolve both cards. + +When you are ready to apply the answer to prompts or to stop the run, go to +[Apply Answers](./apply-answers). diff --git a/docs/interrupts/generic.md b/docs/interrupts/generic.md index b9d57daf2d..1c75db77ec 100644 --- a/docs/interrupts/generic.md +++ b/docs/interrupts/generic.md @@ -2,124 +2,240 @@ title: Generic Interrupts id: interrupts-generic order: 4 -description: "Pause a run to ask the user something that isn't a tool call, validate their answer, and continue." +description: "Ask for typed data from the client during a chat lifecycle boundary." keywords: - tanstack ai - generic interrupt - - responseSchema - - fromJSONSchema - - resolveInterrupt + - defineInterrupt + - onInterruptBoundary + - onInterruptResolution + - INTERRUPT_BOUNDARY_PHASES --- # Generic Interrupts -Sometimes the agent needs an answer that isn't a tool call at all. Mid-run it has -to ask the user to pick a shipping speed, confirm an address, or choose which of -two drafts to keep. There's no tool to approve here, just a question your app -asks and the user answers. +Use a generic interrupt when the server needs data from the client but no tool +call caused the request. For example, ask the user to select a plan before the +model runs. -A generic interrupt is that question. You end the run with a pause that carries a -`responseSchema` describing the answer you expect, render a form for it, and -continue the run once the user submits a valid value. +`defineInterrupt` defines the data that the client sees and the data it must +return. Register the same definition with `chat()` and `useChat()`. This gives +both sides the same types. -Because the pause is defined by your app, you own both ends: the server emits it, -and the client resolves it. +## Define and emit an interrupt -## Resolve it on the client +Put the definition in a module that both the route and the client can import. +The schemas must export JSON Schema. -The schema arrives over the wire, so its value is `unknown` at compile time. -Validate the user's answer against it before resolving. Build the value from your -form fields and pass it straight to the schema: - -```tsx -// app/refund-reason.tsx -import { useState } from 'react' -import type { GenericAGUIInterrupt } from '@tanstack/ai-client' -import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' +```ts +// app/interrupts.ts +import { defineInterrupt } from '@tanstack/ai' import { z } from 'zod' -// You emitted this pause, so you know the shape of the answer. Here it is a -// single reason string chosen from a dropdown. -function RefundReasonForm({ interrupt }: { interrupt: GenericAGUIInterrupt }) { - const [reason, setReason] = useState('damaged') - const [errors, setErrors] = useState>([]) +export const reviewPlan = defineInterrupt({ + id: 'review-plan', + payloadSchema: z.object({ title: z.string(), changes: z.array(z.string()) }), + responseSchema: z.object({ approved: z.boolean(), note: z.string().optional() }), +}) +``` - const submit = () => { - if (!interrupt.responseSchema) { - setErrors(['This interrupt has no response schema.']) - return - } - const result = z - .fromJSONSchema(interrupt.responseSchema) - .safeParse({ reason }) - if (!result.success) { - setErrors(result.error.issues.map((issue) => issue.message)) - return +`payloadSchema` describes display data that travels from the server to the +client. `responseSchema` describes data that travels from the client back to +the server. The display payload is optional. The response schema is required. + +Return requests from `onInterruptBoundary`. The hook can run at each value in +`INTERRUPT_BOUNDARY_PHASES`: + +- `beforeModel`: before the adapter call +- `afterModel`: after the model stream ends +- `beforeTools`: before tool execution +- `afterTools`: after tool results are in `messages` + +Requests from every middleware in the same boundary form one interrupt batch. + +Pick the phase with a `ctx.phase` guard. If `ctx.parentRunId` is set, skip +the emit. If you do not, the same pause happens again. + +When to use each phase, and what `ctx` contains there, is in +[Lifecycle Boundaries](./boundaries). + +```ts +// app/chat-middleware.ts +import type { ChatMiddleware } from '@tanstack/ai' +import { reviewPlan } from './interrupts' + +export const requestReview: ChatMiddleware = { + name: 'request-review', + onInterruptBoundary(ctx) { + if (ctx.phase !== 'beforeModel') return + if (ctx.parentRunId) return + if (ctx.iteration !== 0) return + return { + interrupts: [ + reviewPlan.interrupt({ + key: 'initial-plan', + reason: 'review-required', + message: 'Review the proposed plan.', + payload: { title: 'Release plan', changes: ['Add search', 'Add tests'] }, + }), + ], } - interrupt.resolveInterrupt(result.data) - setErrors([]) - } + }, +} +``` + +`interrupt()` accepts only `key`, `reason`, `message`, `expiresAt`, and, when +declared, `payload`. It returns an immutable request. Do not put secrets in the +payload because the client receives it and persistence can store it. + +## Register it on the server + +Register all definitions on `chat({ interrupts })`. A duplicate definition id +fails before the adapter starts. + +```ts +// app/api/chat/route.ts +import { chat, chatParamsFromRequest, toServerSentEventsResponse } from '@tanstack/ai' +import { openaiText } from '@tanstack/ai-openai' +import { requestReview } from '../../chat-middleware' +import { reviewPlan } from '../../interrupts' + +export async function POST(request: Request) { + const params = await chatParamsFromRequest(request) + const stream = chat({ + adapter: openaiText('gpt-5.5'), + messages: params.messages, + threadId: params.threadId, + runId: params.runId, + ...(params.parentRunId ? { parentRunId: params.parentRunId } : {}), + ...(params.resume ? { resume: params.resume } : {}), + interrupts: [reviewPlan], + middleware: [requestReview], + }) + return toServerSentEventsResponse(stream) +} +``` + +An interrupt ends the current AG-UI run with one `RUN_FINISHED` event whose +outcome is `interrupt`. Resolving it starts a new run with `parentRunId` set to +the interrupted run. The continuation carries the response to the registered +middleware. + +## Resolve it in React + +Register the same definitions with `useChat`. A bound generic interrupt has its +definition id, typed display payload, and typed `resolveInterrupt` method. +Check `kind` and `definitionId`. Then pass the item to a card that takes +`GenericInterrupt`. + +```tsx +// app/plan-review.tsx +import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' +import type { GenericInterrupt } from '@tanstack/ai-react' +import { reviewPlan } from './interrupts' + +function ReviewCard({ + interrupt, +}: { + interrupt: GenericInterrupt +}) { return ( -
-

{interrupt.message ?? interrupt.reason}

- - - {errors.map((message) => ( -

{message}

- ))} -
+ + ) } -export function RefundReasons() { +export function PlanReview() { const { interrupts } = useChat({ - threadId: 'order-7', + threadId: 'release-42', connection: fetchServerSentEvents('/api/chat'), + interrupts: [reviewPlan], }) return ( <> - {interrupts.map((interrupt) => - interrupt.kind === 'generic' ? ( - - ) : null, - )} + {interrupts.map((interrupt) => { + if (interrupt.kind !== 'generic') return null + if (!('definitionId' in interrupt)) return null + if (interrupt.definitionId !== reviewPlan.id) return null + return + })} ) } ``` -`z.fromJSONSchema` gives you a runtime validator, not a trustworthy static type. -The library does not validate the wire schema for you. Whatever you pass to -`resolveInterrupt` is sent as-is, so validate the value here on the client, and -again on the server if you need to trust it, the same way you would treat any -other user input. - -## Emit it on the server - -Tool approvals are rebuilt by `chat()` from message history for free. Generic -pauses are not, because only your app knows when to ask and what to ask. You emit -the descriptor and validate the answer yourself: - -1. End a run with `RUN_FINISHED` and `outcome.type === 'interrupt'`, carrying a - `generic` descriptor with your `responseSchema`. A small middleware is the - usual place to do this. -2. On the continuation request, correlate the incoming `resume` against that - same pending descriptor with `validateInterruptResumeBatch`. It checks the - batch is complete and matches the pending item; it does not validate your - generic value, that is yours to do. Then append the answer and continue. - -The interrupt lab in `examples/ts-react-chat` has a complete middleware that -emits a generic pause and correlates its answer. Without the server half, a -generic answer fails resume validation with `unknown-interrupt` or -`incomplete-batch`. - -> Gating a tool instead of asking a free-form question? A tool -> [approval](./tool-approval) gives you typed branches on top of validation. +`resolveInterrupt` stages the answer. The client sends one continuation only +after every bound interrupt in the batch is resolved or cancelled. Use +`cancel()` when the user declines to provide data. + +## Read resumed values in middleware + +`onInterruptResolution` does not run in the `chat()` call that paused. It runs +once at the start of the **next** `chat()` call, after the client answers. + +That second call is a new run. `useChat` sends `parentRunId` and `resume`. +Each generic resume item carries the original request in `metadata`. The hook +runs after init `onConfig` and before `onStart`. `ctx.phase` is still `'init'`. + +`for(definition)` keeps the response type for that definition. `all()` reads +all registered definitions. `all(definitionA, definitionB)` narrows the +result to those definitions. + +```ts +import type { ChatMiddleware } from '@tanstack/ai' +import { reviewPlan } from './interrupts' + +export const applyReview: ChatMiddleware = { + name: 'apply-review', + onInterruptResolution(_ctx, resumedInterrupts) { + for (const result of resumedInterrupts.for(reviewPlan)) { + if (result.status === 'resolved' && !result.response.approved) { + return { toolResume: 'stop' } + } + } + }, +} +``` + +Middleware can return `toolResume: 'continue'`, `'cancel'`, or `'stop'`. +When more than one middleware returns a value, `stop` wins over `cancel`, and +`cancel` wins over `continue`. + +This hook cannot change prompts, tools, or messages. Store the answer on a +capability, then return those fields from `onConfig` when +`ctx.phase === 'beforeModel'`. The full order, plus a working example, is in +[Apply Answers](./apply-answers). + +## Try the four lifecycle phases + +The React chat example has a playground for `beforeModel`, `afterModel`, +`beforeTools`, and `afterTools`. Each pause shows two typed cards: +`reviewPlan` and `chooseAudience`. + +1. Start `examples/ts-react-chat`. +2. Open `/generic-interrupts`. +3. Pick a phase, resolve both cards, and watch the selected policy. + +## External generic interrupts + +An external system can emit a standard AG-UI generic interrupt. TanStack AI +shows it as `kind: 'unbound'` unless it has a valid TanStack binding. It stays +visible, but it has no resolve or cancel method. This keeps another system from +receiving a continuation that TanStack AI owns. + +See [Multiple Interrupts](./multiple) for mixed tool approvals and generic +interrupts. See [Chat persistence](../persistence/chat-persistence) when an +interrupt must survive a restart. + +| You want to | Page | +| --- | --- | +| Pick `beforeModel`, `afterModel`, `beforeTools`, or `afterTools` | [Lifecycle Boundaries](./boundaries) | +| Apply the user answer to prompts, tools, or `toolResume` | [Apply Answers](./apply-answers) | +| Mix generic interrupts with tool approvals | [Multiple Interrupts](./multiple) | diff --git a/docs/interrupts/migration.md b/docs/interrupts/migration.md index bdd4d867f9..b82a2586b9 100644 --- a/docs/interrupts/migration.md +++ b/docs/interrupts/migration.md @@ -18,8 +18,8 @@ AG-UI interrupt descriptors. Native runs end with `RUN_FINISHED.outcome.type === 'interrupt'`, and the continuation is a new run whose `parentRunId` is the interrupted run. -There's no codemod. Migrate the server lifecycle and client rendering together. -Legacy readers stay temporarily for old streams but can't provide the full +There is no codemod. Migrate the server lifecycle and client rendering together. +Legacy readers stay temporarily for old streams but cannot provide the full native contract. Start from [Overview](./overview). ## API mapping @@ -130,7 +130,7 @@ Deprecated readers recognize well-formed historical `approval-requested` and cloned-history follow-up. They do **not** support edited arguments, custom approval payloads, generic responses, payloadless cancellation, or expiry/ schema-hash reconciliation; those fail with `legacy-unsupported`. Native and -legacy items can't mix in one batch; a failed legacy transport keeps staged +legacy items cannot mix in one batch; a failed legacy transport keeps staged decisions and reports `legacy-submit-failed`. ## Checklist diff --git a/docs/interrupts/multiple.md b/docs/interrupts/multiple.md index 0b0ebd78b4..c069c95be3 100644 --- a/docs/interrupts/multiple.md +++ b/docs/interrupts/multiple.md @@ -2,7 +2,7 @@ title: Multiple Interrupts id: interrupts-multiple order: 3 -description: "Render a queue of pending decisions and resolve them item by item or all at once as one atomic batch." +description: "Render a queue of pending decisions and resolve them item by item or all at once as one validated batch." keywords: - tanstack ai - ag-ui interrupts @@ -13,9 +13,9 @@ keywords: # Multiple Interrupts -One run can pause on several decisions at once. The model lines up three -transfers, or an approval and a question land together. You want to show the -whole queue and send the answers back together, not one round trip each. +One run can pause on several decisions at once. A tool approval and a generic +middleware request can land in the same batch. You want to show the whole queue +and send the answers back together, not one round trip each. ## Two ways to resolve @@ -27,8 +27,8 @@ call a method on the item itself. interrupt.resolveInterrupt(true) ``` -When several are pending, it is often easier to answer them all from one place. -The `useChat` hook gives you root helpers that act on the whole queue: +When several client-owned items are pending, it is often easier to answer them +all from one place. The `useChat` hook gives you root helpers for those items: ```ts ignore // All at once: one callback decides every pending item. @@ -42,8 +42,19 @@ resolveInterrupts((interrupt) => { ``` Both stage local drafts. Nothing goes to the server until every pending item has -an answer, then the whole set submits at once. The server accepts all of them or -none, so you never end up with half a batch applied. +an answer. Then the client submits the whole set at once. The server validates +the whole set before it accepts any item. + +When persistence is enabled, `InterruptStore.commitBatch()` can commit the +validated batch in one database transaction. Without `commitBatch()`, the +compatibility fallback writes items in sequence. That fallback is not atomic. + +For a mixed batch, switch on `kind`. Registered generic items use their +definition id to narrow their payload and response. Tool approvals keep their +own controls. A generic item with a valid raw binding has untyped controls and +joins the resume batch. An `unbound` item has a missing, malformed, or +unsupported binding. It remains visible, has no controls, and does not block +the resumable items. ## Render the queue @@ -110,9 +121,10 @@ export function DecisionQueue() { ## Resolve every item from one callback -`resolveInterrupts(callback)` runs your callback once per item inside a single -synchronous transaction. It must resolve or cancel every item. If it throws or -leaves one item unanswered, nothing submits: +`resolveInterrupts(callback)` runs your callback once per resumable item in a +single synchronous transaction. This includes valid raw external generic +bindings. It must resolve or cancel every such item. If it throws or leaves one +unanswered, nothing submits: ```ts ignore resolveInterrupts((interrupt) => { @@ -130,7 +142,7 @@ Two shortcuts cover the common cases: whole queue. It works only when every item is a tool approval that needs no payload or edits. Generic items, mixed queues, or required payloads are rejected. -- `cancelInterrupts()` cancels every item with no payload. +- `cancelInterrupts()` cancels every resumable item with no payload. ## When an answer is wrong @@ -170,8 +182,8 @@ export function RobustQueue() { tools: [transferTool] as const, }) - // Retry only helps a transport failure. Expired or stale batches can't be - // retried, so don't offer it for those. + // Retry only helps a transport failure. Expired or stale batches cannot + // be retried, so do not offer it for those. const canRetry = interruptErrors.some((error) => error.code === 'transport') return ( @@ -233,6 +245,8 @@ The two recovery paths, side by side: - `interrupt.clearResolution()` drops one item's draft so the user can answer it again from scratch. Fixing a form and calling `resolveInterrupt` again works too, the draft is replaced, not stacked. -- `retryInterrupts()` re-sends the whole staged batch after a transport failure. - It does nothing for expired or stale batches, start a fresh run to get a new - set of interrupts for those. +- `retryInterrupts()` retries a staged batch after a transport failure. With + sequential durable storage, a failed request may have committed earlier + entries already. Reload the current pending records, then retry only the + remaining batch. It does nothing for expired or stale batches, start a fresh + run to get a new set of interrupts for those. diff --git a/docs/interrupts/overview.md b/docs/interrupts/overview.md index 9a5233b106..09cb256c45 100644 --- a/docs/interrupts/overview.md +++ b/docs/interrupts/overview.md @@ -14,7 +14,7 @@ keywords: # Interrupts Most agent runs are fire and forget. The model calls tools, they run, you get an -answer back. But some steps shouldn't happen on their own: moving money, +answer back. But some steps must not happen on their own: moving money, deleting a project, sending an email. And sometimes the agent needs an answer only the user can give before it can go on. @@ -36,16 +36,16 @@ sequenceDiagram participant Client participant Server - Client->>Server: send message — run starts - Server-->>Client: interrupt outcome — run ends without a final answer + Client->>Server: send message, run starts + Server-->>Client: interrupt outcome, run ends without a final answer Client->>User: pending decisions surface as `interrupts` User->>Client: approve / reject / submit a value - Client->>Server: continuation request with the answers — a fresh run + Client->>Server: continuation request with the answers, a fresh run Server-->>Client: the agent picks up where it paused, final answer ``` Note that the pause spans **two runs**: the interrupted one ends, and the -continuation is a new run. One user-visible turn, two run lifecycles — see +continuation is a new run. One user-visible turn, two run lifecycles. See [Threads and runs](../chat/streaming#threads-and-runs). No database is required. The browser sends the full message history back on the @@ -59,21 +59,40 @@ Two kinds of interrupt show up in the `interrupts` array for you to resolve: | `kind` | You get a pause when | Guide | | --- | --- | --- | | `tool-approval` | A tool is marked `needsApproval` and the model calls it | [Tool Approval](./tool-approval) | -| `generic` | Your app ends a run to ask the user something that isn't a tool | [Generic Interrupts](./generic) | +| `generic` | Middleware requests typed client data at a lifecycle boundary | [Generic Interrupts](./generic) | -## Interrupts that aren't ours: `unbound` +## First-party generic interrupts + +For a generic interrupt that TanStack AI owns, define it once with +`defineInterrupt()`. Register the definition with both `chat({ interrupts })` +and `useChat({ interrupts })`. Middleware emits it through +`onInterruptBoundary`, and the client receives a typed bound item that it can +resolve or cancel. See [Generic Interrupts](./generic). To pick a phase, see +[Lifecycle Boundaries](./boundaries). To apply the answer, see +[Apply Answers](./apply-answers). + +## External generic interrupts An interrupt is a standard AG-UI object, and TanStack AI is not the only thing that can put one on a stream. A workflow engine pausing for a durable approval, or another agent framework sharing the same connection, emits the same envelope. -What makes a pause resumable *here* is a binding this library attaches to the -interrupt's metadata, under a key exported as `INTERRUPT_BINDING_METADATA_KEY`. -It records which run and generation the pause belongs to, so your answer can be -matched back to the paused step. +There are three cases: + +- A registered first-party generic interrupt has `kind: 'generic'`, a literal + `definitionId`, typed `payload`, and a typed `resolveInterrupt` method. +- An external generic interrupt with a valid binding also has `kind: 'generic'`. + Its response is `unknown`, but it has `resolveInterrupt`, `cancel`, and + `clearResolution`. It joins the root batch controls. +- An interrupt with a missing, malformed, or unsupported binding has + `kind: 'unbound'`. It remains visible, but has no controls. -When an interrupt arrives without one, you get it with `kind: 'unbound'` and -`canResolve: false`, and there is no `resolveInterrupt` to call: +The binding is stored in the interrupt metadata under +`INTERRUPT_BINDING_METADATA_KEY`. It records the interrupted run and generation. +The client uses it to send the answer to the matching paused step. + +Render unbound items as status information. Do not render a response form for +them: ```tsx import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' @@ -98,22 +117,29 @@ export function Pauses() { return ( <> {interrupts.map((interrupt) => { - // Someone else owns this pause: show it, but offer no way to answer it. if (interrupt.kind === 'unbound') { return (

- Paused elsewhere: {interrupt.message ?? interrupt.reason} + External pause: {interrupt.message ?? interrupt.reason}

) } if (interrupt.kind === 'generic') { return ( - +
+

{interrupt.message ?? interrupt.reason}

+ + + +
) } return ( @@ -131,13 +157,14 @@ export function Pauses() { ``` The library will not invent a binding to make these resolvable. Doing so would -render a form whose answer gets submitted against a run that has nothing pending -— failing only after the user has filled it in. `unbound` says plainly that the -pause belongs to something else, and unbound items never block you from +render a form whose answer is sent to a run that has nothing pending. The +submit would fail only after the user has filled it in. `unbound` says that +the pause belongs to something else. Unbound items never block you from resolving the ones that are yours. -If you emit your own pauses and want them resumable here, attach the binding -with `withInterruptBinding` rather than writing the metadata key by hand: +If an external producer wants the chat client to resume its pause, attach a +valid binding with `withInterruptBinding`. Do not write the metadata key by +hand. Use the exact interrupted run id and generation that own the pause: ```ts import { @@ -164,6 +191,8 @@ const descriptor = withInterruptBinding( v: INTERRUPT_BINDING_VERSION, kind: 'generic', interruptId: 'shipping-1', + interruptedRunId: 'run-42', + generation: 0, // The server checks the schema it hands out still matches the one it // validates against, so the hash is computed from the schema itself. responseSchemaHash: digestInterruptJson( @@ -173,9 +202,13 @@ const descriptor = withInterruptBinding( ) ``` -`v` is the binding's wire version. Readers reject a version they don't -recognise instead of guessing at the fields, which is what keeps another -producer's binding from being mistaken for one of ours. +The client treats this as an untyped generic interrupt. The example above can +stage a value, cancel it, or clear the draft. It can also join +`resolveInterrupts(...)` with tool approvals and first-party generic interrupts. + +`v` is the binding wire version. The client rejects unknown versions and bad +fields. Those interrupts become `unbound` rather than a form that cannot +resume the owner. ## What about client tools? @@ -201,5 +234,7 @@ the same `tool-approval` interrupt. | Approve or reject a single tool call | [Tool Approval](./tool-approval) | | Resolve several pending decisions at once | [Multiple Interrupts](./multiple) | | Ask the user something that isn't a tool | [Generic Interrupts](./generic) | +| Pick `beforeModel`, `afterModel`, `beforeTools`, or `afterTools` | [Lifecycle Boundaries](./boundaries) | +| Apply a generic answer to prompts or stop the run | [Apply Answers](./apply-answers) | | Run a tool in the browser | [Client Tools](../tools/client-tools) | | Move off the old `approval-requested` events | [Migration](./migration) | diff --git a/docs/interrupts/tool-approval.md b/docs/interrupts/tool-approval.md index 0437e40b60..667e563bfc 100644 --- a/docs/interrupts/tool-approval.md +++ b/docs/interrupts/tool-approval.md @@ -13,7 +13,7 @@ keywords: # Tool Approval -You have a tool that shouldn't run until a person says yes: transferring money, +You have a tool that must not run until a person says yes: transferring money, deleting a record, sending a message. You want the model to plan the call, then wait for a human to approve it before anything happens. diff --git a/docs/persistence/chat-persistence.md b/docs/persistence/chat-persistence.md index e366382763..0810c3b9c2 100644 --- a/docs/persistence/chat-persistence.md +++ b/docs/persistence/chat-persistence.md @@ -115,8 +115,13 @@ with `snapshotIntervalMs` (default `1000`). On **error**, the run is marked `failed`. On **abort**, the run is marked `aborted` with a `finishedAt`; `interrupted` is written only at an interrupt boundary, and it is not terminal. Resumes accepted in `onConfig` are **not** -consumed until a success boundary (interrupt or finish), so a failed run leaves -pending interrupts retryable with the same resume batch. +consumed until a success boundary (interrupt or finish). If a run fails before +the resume commit, every pending interrupt stays retryable. The same is true +when an atomic `commitBatch()` fails: the whole batch stays pending. + +The legacy sequential fallback is different. A write can fail after earlier +entries were committed. Reload the thread's current pending interrupt records +and submit only that remaining batch. Do not resend the original full batch. One abort does **not** terminalize: a plain client disconnect on a run that some other middleware has declared *detachable* (a durable event log plus a run @@ -149,22 +154,34 @@ stateDiagram-v2 ## Interrupts survive a restart -When a run pauses on an interrupt (a tool approval, a client-side tool, a -generic wait), the middleware records it. A later request on that thread must -carry a `resume` batch that answers the pending interrupts before new input is -accepted, otherwise it is rejected, which is why the example above forwards -`params.resume`. +When a run pauses on an interrupt (a tool approval, a client-side tool, or a +generic middleware request), the middleware records it. A later request on that +thread must carry a `resume` batch that answers the pending interrupts before +new input is accepted, otherwise it is rejected, which is why the example above +forwards `params.resume`. + +For a mixed batch, the persistence middleware validates every entry before it +continues. It commits all resolved and cancelled entries at one success +boundary. An `InterruptStore` can implement `commitBatch()` to make that write +atomic. Without it, the compatibility fallback writes entries one at a time and +is not atomic. Persistence is the **server-authoritative resume path**: the middleware validates the resume batch against pending interrupts, builds -`ChatResumeToolState` (approvals / client-tool results), and **clears** -`config.resume` so the chat engine skips its ephemeral reconstruction (which -needs client message history the persistence flow deliberately omits). Resumes -are committed (resolved/cancelled in the store) only once the run reaches a -successful interrupt or finish boundary. - -An interrupt record is born `pending` and only a commit moves it, which is why -a failed continuation leaves it answerable again: +`ChatResumeToolState` (approvals / client-tool results / generic requests), and +**clears** `config.resume` so the chat engine skips its ephemeral +reconstruction (which needs client message history the persistence flow +deliberately omits). Resumes are committed (resolved/cancelled in the store) +only once the run reaches a successful interrupt or finish boundary. + +`onInterruptResolution` still runs at the same moment as a request without +persistence: after init `onConfig`, before `onStart`. See +[Apply Answers](../interrupts/apply-answers). + +An interrupt record is born `pending` and only a commit moves it. With +`commitBatch()`, a failed commit leaves the full batch answerable again. With +the sequential fallback, reload first because some earlier entries may already +be resolved or cancelled: ```mermaid stateDiagram-v2 diff --git a/docs/persistence/store-reference.md b/docs/persistence/store-reference.md index fd0633f425..dc8cde2bb9 100644 --- a/docs/persistence/store-reference.md +++ b/docs/persistence/store-reference.md @@ -216,8 +216,13 @@ interface InterruptRecord { response?: unknown } +type InterruptCommitEntry = + | { interruptId: string; status: 'resolved'; response?: unknown } + | { interruptId: string; status: 'cancelled' } + interface InterruptStore { create(record: Omit): Promise + commitBatch?(entries: ReadonlyArray): Promise resolve(interruptId: string, response?: unknown): Promise cancel(interruptId: string): Promise get(interruptId: string): Promise @@ -228,6 +233,15 @@ interface InterruptStore { } ``` +`commitBatch` is optional. Use one database transaction for all entries when +you implement it. The legacy `resolve` and `cancel` fallback is sequential and +is not atomic. + +If you implement `commitBatch`, reject the whole batch (throw, write nothing) +when any entry has a duplicate `interruptId`, names an interrupt that does not +exist, or names an interrupt that is not `'pending'`. `resolve` and `cancel` +stay no-ops for a missing `interruptId`. + `create` accepts a record without `status`/`resolvedAt` so every interrupt is born `'pending'`; it is insert-if-absent, so a duplicate `create` never clobbers an already-resolved interrupt. The `list*` methods return records ordered by diff --git a/examples/ts-react-chat/README.md b/examples/ts-react-chat/README.md index ff61b66566..56e8d15ddd 100644 --- a/examples/ts-react-chat/README.md +++ b/examples/ts-react-chat/README.md @@ -74,6 +74,20 @@ The lazy tools are: `compareGuitars`, `calculateFinancing`, and `searchGuitars`. - In multi-turn conversations, previously discovered tools are usable immediately without re-discovery - If the LLM skips discovery, it gets an error and self-corrects +## Generic interrupt playground + +Open `/generic-interrupts` to try first-party generic interrupts at each +lifecycle phase: + +- `beforeModel` +- `afterModel` +- `beforeTools` +- `afterTools` + +Pick a resume policy (`continue`, `cancel`, or `stop`). Each pause shows two +typed cards (`reviewPlan` and `chooseAudience`). Resolve both, then watch the +selected policy. + ## ✨ Features ### AI Capabilities diff --git a/examples/ts-react-chat/src/components/Header.tsx b/examples/ts-react-chat/src/components/Header.tsx index 9b4d744f98..9b852ea7ba 100644 --- a/examples/ts-react-chat/src/components/Header.tsx +++ b/examples/ts-react-chat/src/components/Header.tsx @@ -353,6 +353,19 @@ export default function Header() { Interrupts Lab + setIsOpen(false)} + className="flex items-center gap-3 p-3 rounded-lg hover:bg-gray-800 transition-colors mb-2" + activeProps={{ + className: + 'flex items-center gap-3 p-3 rounded-lg bg-cyan-600 hover:bg-cyan-700 transition-colors mb-2', + }} + > + + Generic Interrupts + + setIsOpen(false)} diff --git a/examples/ts-react-chat/src/lib/generic-interrupt-playground.ts b/examples/ts-react-chat/src/lib/generic-interrupt-playground.ts new file mode 100644 index 0000000000..ca5f6069bc --- /dev/null +++ b/examples/ts-react-chat/src/lib/generic-interrupt-playground.ts @@ -0,0 +1,125 @@ +import { + defineInterrupt, + INTERRUPT_BOUNDARY_PHASES, + INTERRUPT_TOOL_RESUMES, + toolDefinition, +} from '@tanstack/ai' +import { z } from 'zod' +import type { InterruptBoundaryPhase, InterruptToolResume } from '@tanstack/ai' + +export const reviewPlan = defineInterrupt({ + id: 'review-plan', + payloadSchema: z.object({ + title: z.string(), + boundary: z.enum(INTERRUPT_BOUNDARY_PHASES), + }), + responseSchema: z.object({ + approved: z.boolean(), + note: z.string().min(1), + }), +}) + +export const AUDIENCE_OPTIONS = ['students', 'staff', 'mixed'] as const + +export const chooseAudience = defineInterrupt({ + id: 'choose-audience', + payloadSchema: z.object({ + question: z.string(), + options: z.array(z.enum(AUDIENCE_OPTIONS)), + }), + responseSchema: z.object({ + audience: z.enum(AUDIENCE_OPTIONS), + }), +}) + +export const playgroundInterrupts = [reviewPlan, chooseAudience] as const + +export const inspectPlan = toolDefinition({ + name: 'inspectPlan', + description: 'Inspect a published plan on the server.', + inputSchema: z.object({ + planId: z.string(), + }), + outputSchema: z.object({ + inspected: z.boolean(), + planId: z.string(), + }), +}) + +export interface PlaygroundScenario { + id: string + boundary: InterruptBoundaryPhase + title: string + blurb: string + message: string + needsTool: boolean +} + +export const playgroundScenarios: ReadonlyArray = [ + { + id: 'before-model', + boundary: 'beforeModel', + title: 'Before the model', + blurb: 'The run pauses before the first model call.', + message: 'Plan a one-hour visit for a school group.', + needsTool: false, + }, + { + id: 'after-model', + boundary: 'afterModel', + title: 'After the model', + blurb: 'The run pauses after the model writes a draft.', + message: 'Write a short welcome for new volunteers.', + needsTool: false, + }, + { + id: 'before-tools', + boundary: 'beforeTools', + title: 'Before tools', + blurb: 'The run pauses after the model asks for a tool, before it runs.', + message: 'Inspect plan PLAN-42.', + needsTool: true, + }, + { + id: 'after-tools', + boundary: 'afterTools', + title: 'After tools', + blurb: 'The run pauses after the tool result is ready.', + message: 'Inspect plan PLAN-42 and then summarize it.', + needsTool: true, + }, +] + +export function isPlaygroundBoundary( + value: unknown, +): value is InterruptBoundaryPhase { + return ( + typeof value === 'string' && + INTERRUPT_BOUNDARY_PHASES.some((boundary) => boundary === value) + ) +} + +export function isPlaygroundPolicy( + value: unknown, +): value is InterruptToolResume { + return ( + typeof value === 'string' && + INTERRUPT_TOOL_RESUMES.some((policy) => policy === value) + ) +} + +export function readPlaygroundForwarded(value: unknown): { + boundary: InterruptBoundaryPhase + policy: InterruptToolResume +} { + if (!value || typeof value !== 'object' || Array.isArray(value)) { + return { boundary: 'beforeModel', policy: 'continue' } + } + const record: Record = { ...value } + return { + boundary: isPlaygroundBoundary(record.boundary) + ? record.boundary + : 'beforeModel', + policy: isPlaygroundPolicy(record.policy) ? record.policy : 'continue', + } +} diff --git a/examples/ts-react-chat/src/routeTree.gen.ts b/examples/ts-react-chat/src/routeTree.gen.ts index 7dfa0251b7..70172efdc5 100644 --- a/examples/ts-react-chat/src/routeTree.gen.ts +++ b/examples/ts-react-chat/src/routeTree.gen.ts @@ -24,6 +24,7 @@ import { Route as Issue176ToolResultRouteImport } from './routes/issue-176-tool- import { Route as InterruptsRouteImport } from './routes/interrupts' import { Route as ImageToolReproRouteImport } from './routes/image-tool-repro' import { Route as ImageGenRouteImport } from './routes/image-gen' +import { Route as GenericInterruptsRouteImport } from './routes/generic-interrupts' import { Route as GenerationHooksRouteImport } from './routes/generation-hooks' import { Route as CapabilityDemoRouteImport } from './routes/capability-demo' import { Route as IndexRouteImport } from './routes/index' @@ -57,6 +58,7 @@ import { Route as ApiMcpAppsCallRouteImport } from './routes/api.mcp-apps-call' import { Route as ApiInterruptsRouteImport } from './routes/api.interrupts' import { Route as ApiImageToolReproRouteImport } from './routes/api.image-tool-repro' import { Route as ApiImageGenRouteImport } from './routes/api.image-gen' +import { Route as ApiGenericInterruptsRouteImport } from './routes/api.generic-interrupts' import { Route as ApiCapabilityDemoRouteImport } from './routes/api.capability-demo' import { Route as ApiArtifactsRouteImport } from './routes/api.artifacts' import { Route as ExampleGuitarsIndexRouteImport } from './routes/example.guitars/index' @@ -142,6 +144,11 @@ const ImageGenRoute = ImageGenRouteImport.update({ path: '/image-gen', getParentRoute: () => rootRouteImport, } as any) +const GenericInterruptsRoute = GenericInterruptsRouteImport.update({ + id: '/generic-interrupts', + path: '/generic-interrupts', + getParentRoute: () => rootRouteImport, +} as any) const GenerationHooksRoute = GenerationHooksRouteImport.update({ id: '/generation-hooks', path: '/generation-hooks', @@ -311,6 +318,11 @@ const ApiImageGenRoute = ApiImageGenRouteImport.update({ path: '/api/image-gen', getParentRoute: () => rootRouteImport, } as any) +const ApiGenericInterruptsRoute = ApiGenericInterruptsRouteImport.update({ + id: '/api/generic-interrupts', + path: '/api/generic-interrupts', + getParentRoute: () => rootRouteImport, +} as any) const ApiCapabilityDemoRoute = ApiCapabilityDemoRouteImport.update({ id: '/api/capability-demo', path: '/api/capability-demo', @@ -362,6 +374,7 @@ export interface FileRoutesByFullPath { '/': typeof IndexRoute '/capability-demo': typeof CapabilityDemoRoute '/generation-hooks': typeof GenerationHooksRoute + '/generic-interrupts': typeof GenericInterruptsRoute '/image-gen': typeof ImageGenRoute '/image-tool-repro': typeof ImageToolReproRoute '/interrupts': typeof InterruptsRoute @@ -379,6 +392,7 @@ export interface FileRoutesByFullPath { '/typesafe-tools': typeof TypesafeToolsRoute '/api/artifacts': typeof ApiArtifactsRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute + '/api/generic-interrupts': typeof ApiGenericInterruptsRoute '/api/image-gen': typeof ApiImageGenRoute '/api/image-tool-repro': typeof ApiImageToolReproRoute '/api/interrupts': typeof ApiInterruptsRoute @@ -421,6 +435,7 @@ export interface FileRoutesByTo { '/': typeof IndexRoute '/capability-demo': typeof CapabilityDemoRoute '/generation-hooks': typeof GenerationHooksRoute + '/generic-interrupts': typeof GenericInterruptsRoute '/image-gen': typeof ImageGenRoute '/image-tool-repro': typeof ImageToolReproRoute '/interrupts': typeof InterruptsRoute @@ -438,6 +453,7 @@ export interface FileRoutesByTo { '/typesafe-tools': typeof TypesafeToolsRoute '/api/artifacts': typeof ApiArtifactsRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute + '/api/generic-interrupts': typeof ApiGenericInterruptsRoute '/api/image-gen': typeof ApiImageGenRoute '/api/image-tool-repro': typeof ApiImageToolReproRoute '/api/interrupts': typeof ApiInterruptsRoute @@ -481,6 +497,7 @@ export interface FileRoutesById { '/': typeof IndexRoute '/capability-demo': typeof CapabilityDemoRoute '/generation-hooks': typeof GenerationHooksRoute + '/generic-interrupts': typeof GenericInterruptsRoute '/image-gen': typeof ImageGenRoute '/image-tool-repro': typeof ImageToolReproRoute '/interrupts': typeof InterruptsRoute @@ -498,6 +515,7 @@ export interface FileRoutesById { '/typesafe-tools': typeof TypesafeToolsRoute '/api/artifacts': typeof ApiArtifactsRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute + '/api/generic-interrupts': typeof ApiGenericInterruptsRoute '/api/image-gen': typeof ApiImageGenRoute '/api/image-tool-repro': typeof ApiImageToolReproRoute '/api/interrupts': typeof ApiInterruptsRoute @@ -542,6 +560,7 @@ export interface FileRouteTypes { | '/' | '/capability-demo' | '/generation-hooks' + | '/generic-interrupts' | '/image-gen' | '/image-tool-repro' | '/interrupts' @@ -559,6 +578,7 @@ export interface FileRouteTypes { | '/typesafe-tools' | '/api/artifacts' | '/api/capability-demo' + | '/api/generic-interrupts' | '/api/image-gen' | '/api/image-tool-repro' | '/api/interrupts' @@ -601,6 +621,7 @@ export interface FileRouteTypes { | '/' | '/capability-demo' | '/generation-hooks' + | '/generic-interrupts' | '/image-gen' | '/image-tool-repro' | '/interrupts' @@ -618,6 +639,7 @@ export interface FileRouteTypes { | '/typesafe-tools' | '/api/artifacts' | '/api/capability-demo' + | '/api/generic-interrupts' | '/api/image-gen' | '/api/image-tool-repro' | '/api/interrupts' @@ -660,6 +682,7 @@ export interface FileRouteTypes { | '/' | '/capability-demo' | '/generation-hooks' + | '/generic-interrupts' | '/image-gen' | '/image-tool-repro' | '/interrupts' @@ -677,6 +700,7 @@ export interface FileRouteTypes { | '/typesafe-tools' | '/api/artifacts' | '/api/capability-demo' + | '/api/generic-interrupts' | '/api/image-gen' | '/api/image-tool-repro' | '/api/interrupts' @@ -720,6 +744,7 @@ export interface RootRouteChildren { IndexRoute: typeof IndexRoute CapabilityDemoRoute: typeof CapabilityDemoRoute GenerationHooksRoute: typeof GenerationHooksRoute + GenericInterruptsRoute: typeof GenericInterruptsRoute ImageGenRoute: typeof ImageGenRoute ImageToolReproRoute: typeof ImageToolReproRoute InterruptsRoute: typeof InterruptsRoute @@ -737,6 +762,7 @@ export interface RootRouteChildren { TypesafeToolsRoute: typeof TypesafeToolsRoute ApiArtifactsRoute: typeof ApiArtifactsRoute ApiCapabilityDemoRoute: typeof ApiCapabilityDemoRoute + ApiGenericInterruptsRoute: typeof ApiGenericInterruptsRoute ApiImageGenRoute: typeof ApiImageGenRoute ApiImageToolReproRoute: typeof ApiImageToolReproRoute ApiInterruptsRoute: typeof ApiInterruptsRoute @@ -882,6 +908,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ImageGenRouteImport parentRoute: typeof rootRouteImport } + '/generic-interrupts': { + id: '/generic-interrupts' + path: '/generic-interrupts' + fullPath: '/generic-interrupts' + preLoaderRoute: typeof GenericInterruptsRouteImport + parentRoute: typeof rootRouteImport + } '/generation-hooks': { id: '/generation-hooks' path: '/generation-hooks' @@ -1113,6 +1146,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiImageGenRouteImport parentRoute: typeof rootRouteImport } + '/api/generic-interrupts': { + id: '/api/generic-interrupts' + path: '/api/generic-interrupts' + fullPath: '/api/generic-interrupts' + preLoaderRoute: typeof ApiGenericInterruptsRouteImport + parentRoute: typeof rootRouteImport + } '/api/capability-demo': { id: '/api/capability-demo' path: '/api/capability-demo' @@ -1194,6 +1234,7 @@ const rootRouteChildren: RootRouteChildren = { IndexRoute: IndexRoute, CapabilityDemoRoute: CapabilityDemoRoute, GenerationHooksRoute: GenerationHooksRoute, + GenericInterruptsRoute: GenericInterruptsRoute, ImageGenRoute: ImageGenRoute, ImageToolReproRoute: ImageToolReproRoute, InterruptsRoute: InterruptsRoute, @@ -1211,6 +1252,7 @@ const rootRouteChildren: RootRouteChildren = { TypesafeToolsRoute: TypesafeToolsRoute, ApiArtifactsRoute: ApiArtifactsRoute, ApiCapabilityDemoRoute: ApiCapabilityDemoRoute, + ApiGenericInterruptsRoute: ApiGenericInterruptsRoute, ApiImageGenRoute: ApiImageGenRoute, ApiImageToolReproRoute: ApiImageToolReproRoute, ApiInterruptsRoute: ApiInterruptsRoute, diff --git a/examples/ts-react-chat/src/routes/api.generic-interrupts.ts b/examples/ts-react-chat/src/routes/api.generic-interrupts.ts new file mode 100644 index 0000000000..43e811b1ba --- /dev/null +++ b/examples/ts-react-chat/src/routes/api.generic-interrupts.ts @@ -0,0 +1,134 @@ +import { createFileRoute } from '@tanstack/react-router' +import { + chat, + chatParamsFromRequestBody, + maxIterations, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { createOpenaiChat } from '@tanstack/ai-openai' +import { + AUDIENCE_OPTIONS, + chooseAudience, + inspectPlan, + playgroundInterrupts, + readPlaygroundForwarded, + reviewPlan, +} from '@/lib/generic-interrupt-playground' +import type { + ChatMiddleware, + InterruptBoundaryPhase, + InterruptToolResume, +} from '@tanstack/ai' + +const SYSTEM_PROMPT = + 'You are a planning assistant. When the user asks you to inspect a plan, ' + + 'call inspectPlan with the plan id from their message. Keep spoken replies short.' + +function createLifecycleMiddleware( + boundary: InterruptBoundaryPhase, + policy: InterruptToolResume, +): ChatMiddleware { + return { + name: 'generic-interrupt-playground', + onInterruptBoundary(ctx) { + if (ctx.phase !== boundary || ctx.parentRunId) return + return { + interrupts: [ + reviewPlan.interrupt({ + key: `${boundary}-review`, + reason: 'review_required', + message: `Review the plan at ${ctx.phase}.`, + payload: { + title: 'Playground review plan', + boundary: ctx.phase, + }, + }), + chooseAudience.interrupt({ + key: `${boundary}-audience`, + reason: 'audience_required', + message: 'Pick who this reply is for.', + payload: { + question: 'Who should the next reply speak to?', + options: Array.from(AUDIENCE_OPTIONS), + }, + }), + ], + } + }, + onInterruptResolution(_ctx, resolutions) { + for (const resolution of resolutions.for(reviewPlan)) { + if (resolution.status === 'resolved' && !resolution.response.approved) { + return { toolResume: 'stop' } + } + } + return { toolResume: policy } + }, + } +} + +async function handle(request: Request): Promise { + const apiKey = process.env.OPENAI_API_KEY + if (!apiKey) { + return new Response( + JSON.stringify({ + error: 'Set OPENAI_API_KEY in examples/ts-react-chat/.env', + }), + { status: 500, headers: { 'content-type': 'application/json' } }, + ) + } + + let params + try { + params = await chatParamsFromRequestBody(await request.json()) + } catch (error) { + return new Response( + error instanceof Error ? error.message : 'Bad request', + { status: 400 }, + ) + } + + const { boundary, policy } = readPlaygroundForwarded(params.forwardedProps) + const isResume = (params.resume?.length ?? 0) > 0 + const needsTool = boundary === 'beforeTools' || boundary === 'afterTools' + const tools = needsTool + ? [ + inspectPlan.server(async ({ planId }) => ({ + inspected: true, + planId, + })), + ] + : [] + const abortController = new AbortController() + + const stream = chat({ + adapter: createOpenaiChat('gpt-5.5', apiKey), + messages: params.messages, + tools, + systemPrompts: [SYSTEM_PROMPT], + agentLoopStrategy: maxIterations(8), + threadId: params.threadId, + runId: params.runId, + ...(params.parentRunId ? { parentRunId: params.parentRunId } : {}), + ...(params.resume ? { resume: params.resume } : {}), + interrupts: playgroundInterrupts, + middleware: [createLifecycleMiddleware(boundary, policy)], + ...(!isResume && needsTool + ? { + modelOptions: { + tool_choice: { type: 'function', name: inspectPlan.name }, + }, + } + : {}), + abortController, + }) + + return toServerSentEventsResponse(stream, { abortController }) +} + +export const Route = createFileRoute('/api/generic-interrupts')({ + server: { + handlers: { + POST: ({ request }) => handle(request), + }, + }, +}) diff --git a/examples/ts-react-chat/src/routes/generic-interrupts.tsx b/examples/ts-react-chat/src/routes/generic-interrupts.tsx new file mode 100644 index 0000000000..fe3185c756 --- /dev/null +++ b/examples/ts-react-chat/src/routes/generic-interrupts.tsx @@ -0,0 +1,449 @@ +import { useEffect, useRef, useState } from 'react' +import { createFileRoute } from '@tanstack/react-router' +import { INTERRUPT_TOOL_RESUMES } from '@tanstack/ai' +import { + fetchServerSentEvents, + useChat, + type GenericInterrupt, + type UseChatReturn, +} from '@tanstack/ai-react' +import { Check, RotateCcw, Sparkles, Trash2 } from 'lucide-react' +import { + AUDIENCE_OPTIONS, + chooseAudience, + inspectPlan, + playgroundInterrupts, + playgroundScenarios, + reviewPlan, +} from '@/lib/generic-interrupt-playground' +import type { InterruptToolResume } from '@tanstack/ai' +import type { PlaygroundScenario } from '@/lib/generic-interrupt-playground' + +export const Route = createFileRoute('/generic-interrupts')({ + component: GenericInterruptPlayground, +}) + +const clientTools = [inspectPlan.client()] as const +const connection = fetchServerSentEvents('/api/generic-interrupts') + +type PlaygroundChat = UseChatReturn< + typeof clientTools, + undefined, + typeof playgroundInterrupts +> + +const POLICY_HELP: Record = { + continue: 'The run continues after you resolve the review.', + cancel: 'Pending tools are cancelled. Then the run continues.', + stop: 'The run ends after you resolve the review.', +} + +function GenericInterruptPlayground() { + const [threadId, setThreadId] = useState(() => crypto.randomUUID()) + const [active, setActive] = useState(null) + const [policy, setPolicy] = useState('continue') + const [pending, setPending] = useState(null) + const [decisions, setDecisions] = useState>([]) + const record = (message: string) => + setDecisions((prev) => [message, ...prev].slice(0, 8)) + + const chat = useChat({ + threadId, + connection, + tools: clientTools, + interrupts: playgroundInterrupts, + forwardedProps: { + boundary: active?.boundary ?? 'beforeModel', + policy, + }, + }) + + useEffect(() => { + if (pending === null) return + void chat.sendMessage(pending) + setPending(null) + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [pending]) + + const runScenario = (scenario: PlaygroundScenario) => { + setDecisions([]) + setActive(scenario) + setThreadId(crypto.randomUUID()) + setPending(scenario.message) + } + + return ( +
+
+ + +
+
+ {chat.isLoading || chat.resuming + ? `Running ${active?.boundary ?? 'scenario'} with policy ${policy}...` + : active + ? `Active boundary: ${active.boundary}. Policy: ${policy}.` + : 'Pick a boundary on the left to start.'} +
+ + + + {chat.interruptErrors.length > 0 ? ( +
+ {chat.interruptErrors.map((error) => ( +
{error.message}
+ ))} + +
+ ) : null} + + {chat.interrupts.filter(isReviewPlan).map((interrupt) => ( + + ))} + + {chat.interrupts.filter(isChooseAudience).map((interrupt) => ( + + ))} + + {decisions.length > 0 ? ( +
+

+ Your decisions +

+
    + {decisions.map((decision, index) => ( +
  • {decision}
  • + ))} +
+
+ ) : null} +
+
+
+ ) +} + +function isReviewPlan( + interrupt: PlaygroundChat['interrupts'][number], +): interrupt is GenericInterrupt { + return ( + interrupt.kind === 'generic' && + 'definitionId' in interrupt && + interrupt.definitionId === reviewPlan.id + ) +} + +function isChooseAudience( + interrupt: PlaygroundChat['interrupts'][number], +): interrupt is GenericInterrupt { + return ( + interrupt.kind === 'generic' && + 'definitionId' in interrupt && + interrupt.definitionId === chooseAudience.id + ) +} + +function Transcript({ chat }: { chat: PlaygroundChat }) { + if (chat.messages.length === 0) { + return ( +
+ Pick a boundary on the left to start. +
+ ) + } + return ( +
+ {chat.messages.map((message) => ( +
+ + {message.role}:{' '} + + {message.parts.map((part, index) => { + if (part.type === 'text') { + return {part.content} + } + if (part.type === 'tool-call') { + return ( +
+ {part.name}({JSON.stringify(part.input ?? part.arguments)}) + {part.output !== undefined + ? ` → ${JSON.stringify(part.output)}` + : ` · ${part.state}`} +
+ ) + } + if (part.type === 'tool-result') { + const body = + typeof part.content === 'string' + ? part.content + : JSON.stringify(part.content) + return ( +
+ {part.error ?? body} +
+ ) + } + return null + })} +
+ ))} +
+ ) +} + +function ReviewCard({ + interrupt, + disabled, + record, +}: { + interrupt: GenericInterrupt + disabled: boolean + record: (message: string) => void +}) { + const [note, setNote] = useState('Looks good') + const payload = interrupt.payload + const cardRef = useRef(null) + + useEffect(() => { + cardRef.current?.scrollIntoView({ block: 'nearest', behavior: 'smooth' }) + }, []) + + return ( +
+
+

+ {payload?.boundary ?? interrupt.reason} +

+

+ {payload?.title ?? 'Review plan'} +

+

+ {interrupt.message ?? 'Review this step, then continue.'} +

+
+ +
+ + + +
+ {interrupt.errors.map((error) => ( +

+ {error.message} +

+ ))} +
+ ) +} + +function AudienceCard({ + interrupt, + disabled, + record, +}: { + interrupt: GenericInterrupt + disabled: boolean + record: (message: string) => void +}) { + const payload = interrupt.payload + const options = payload?.options ?? AUDIENCE_OPTIONS + const cardRef = useRef(null) + + useEffect(() => { + cardRef.current?.scrollIntoView({ block: 'nearest', behavior: 'smooth' }) + }, []) + + return ( +
+
+

+ choose-audience +

+

+ {payload?.question ?? 'Pick an audience'} +

+

+ {interrupt.message ?? 'This card uses a different response type.'} +

+
+
+ {options.map((audience) => ( + + ))} + +
+ {interrupt.errors.map((error) => ( +

+ {error.message} +

+ ))} +
+ ) +} diff --git a/examples/ts-react-chat/src/routes/index.tsx b/examples/ts-react-chat/src/routes/index.tsx index be30c2f919..3c6591b6a0 100644 --- a/examples/ts-react-chat/src/routes/index.tsx +++ b/examples/ts-react-chat/src/routes/index.tsx @@ -13,6 +13,7 @@ import { Music, PauseCircle, Send, + Sparkles, Square, Video, X, @@ -216,6 +217,13 @@ function Messages({ Interrupts Lab + + + Generic Interrupts + Interrupts Lab + + + Generic Interrupts + = any, TSchema extends SchemaInput | undefined = undefined, TContext = InferredClientContext, + const TInterrupts extends ReadonlyArray< + InterruptDefinition + > = readonly [], >( options: InjectChatOptions< TTools, TSchema, - TContext - > = {} as InjectChatOptions, -): InjectChatResult { + TContext, + TInterrupts + > = {} as InjectChatOptions, +): InjectChatResult { assertInInjectionContext(injectChat) type Partial = DeepPartial>> @@ -74,7 +79,7 @@ export function injectChat< const sessionGenerating = signal(false) const queue = signal>([]) const runId = signal(null) - const interruptState = signal>({ + const interruptState = signal>({ interrupts: EMPTY_INTERRUPTS, pendingInterrupts: EMPTY_INTERRUPTS, interruptErrors: EMPTY_INTERRUPT_ERRORS, @@ -97,7 +102,7 @@ export function injectChat< ? { connection: options.connection } : { fetcher: options.fetcher } - const client = new ChatClient({ + const client = new ChatClient({ devtoolsBridgeFactory: createChatDevtoolsBridge, ...transport, id: clientId, @@ -136,6 +141,9 @@ export function injectChat< options.onInterruptStateChange?.(nextInterruptState) }, tools: options.tools, + ...(options.interrupts !== undefined && { + interrupts: options.interrupts, + }), onCustomEvent: (eventType, data, context) => options.onCustomEvent?.(eventType, data, context), ...(options.streamProcessor !== undefined && { @@ -290,7 +298,11 @@ export function injectChat< const interruptErrors = computed(() => interruptState().interruptErrors) const resuming = computed(() => interruptState().resuming) const resolveInterrupts = ( - resolution: boolean | ((interrupt: ChatInterrupt) => undefined), + resolution: + | boolean + | (( + interrupt: ResolvableChatInterrupt, + ) => undefined), ) => { if (typeof resolution === 'boolean') { client.resolveInterrupts(resolution) @@ -335,5 +347,5 @@ export function injectChat< resumeInterruptsUnsafe, partial, final, - } as unknown as InjectChatResult + } as unknown as InjectChatResult } diff --git a/packages/ai-angular/src/types.ts b/packages/ai-angular/src/types.ts index c05948eafb..007cc40d0b 100644 --- a/packages/ai-angular/src/types.ts +++ b/packages/ai-angular/src/types.ts @@ -1,5 +1,6 @@ import type { AnyClientTool, + InterruptDefinition, InferSchemaType, ModelMessage, RunAgentResumeItem, @@ -10,7 +11,7 @@ import type { BoundInterrupts, ChatClientOptions, ChatClientState, - ChatInterrupt, + ResolvableChatInterrupt, ChatInterruptState, ChatRequestBody, ChatResumeState, @@ -67,8 +68,10 @@ export type InjectChatOptions< TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, TContext = InferredClientContext, + TInterrupts extends ReadonlyArray> = + readonly [], > = DistributedOmit< - ChatClientOptions, + ChatClientOptions, | 'onMessagesChange' | 'onLoadingChange' | 'onErrorChange' @@ -106,9 +109,12 @@ export type InjectChatOptions< export type InjectChatResult< TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, + TInterrupts extends ReadonlyArray> = + readonly [], > = BaseInjectChatResult< TTools, - TSchema extends SchemaInput ? InferSchemaType : unknown + TSchema extends SchemaInput ? InferSchemaType : unknown, + TInterrupts > & (TSchema extends SchemaInput ? { @@ -122,6 +128,8 @@ export type InjectChatResult< interface BaseInjectChatResult< TTools extends ReadonlyArray = any, TData = unknown, + TInterrupts extends ReadonlyArray> = + readonly [], > { /** Current messages in the conversation. */ messages: Signal>> @@ -162,16 +170,22 @@ interface BaseInjectChatResult< */ runId: Signal /** Immutable bound interrupts for the current interrupted run. */ - interrupts: Signal> + interrupts: Signal> /** @deprecated Use `interrupts`. */ - pendingInterrupts: Signal> + pendingInterrupts: Signal> /** Batch-level interrupt errors. */ - interruptErrors: Signal['interruptErrors']> + interruptErrors: Signal< + ChatInterruptState['interruptErrors'] + > /** Whether the client is submitting an interrupt batch. */ resuming: Signal resolveInterrupts: { (approved: boolean): void - (resolver: (interrupt: ChatInterrupt) => undefined): void + ( + resolver: ( + interrupt: ResolvableChatInterrupt, + ) => undefined, + ): void } cancelInterrupts: () => void retryInterrupts: () => void diff --git a/packages/ai-angular/tests/inject-chat-types.test.ts b/packages/ai-angular/tests/inject-chat-types.test.ts index 7b46167f8b..2c192d5695 100644 --- a/packages/ai-angular/tests/inject-chat-types.test.ts +++ b/packages/ai-angular/tests/inject-chat-types.test.ts @@ -6,10 +6,10 @@ import { describe, expectTypeOf, it } from 'vitest' import { z } from 'zod' -import { toolDefinition } from '@tanstack/ai' +import { defineInterrupt, toolDefinition } from '@tanstack/ai' import { clientTools } from '@tanstack/ai-client' import type { AnyClientTool } from '@tanstack/ai' -import type { injectChat } from '../src/inject-chat' +import { injectChat } from '../src/inject-chat' import type { Signal } from '@angular/core' import type { DeepPartial, InjectChatResult } from '../src/types' @@ -191,3 +191,86 @@ describe('injectChat() interrupt types', () => { void check }) }) + +describe('injectChat() registered generic interrupt types', () => { + it('keeps registered and external generic interrupts distinct', () => { + const reviewPlan = defineInterrupt({ + id: 'review-plan', + payloadSchema: z.object({ title: z.string() }), + responseSchema: z.string().transform((value) => Number(value)), + }) + const acknowledge = defineInterrupt({ + id: 'acknowledge', + responseSchema: z.object({ accepted: z.boolean() }), + }) + + const check = () => { + const chat = injectChat({ + connection: { connect: async function* () {} }, + interrupts: [reviewPlan, acknowledge], + }) + type Interrupt = ReturnType[number] + type Review = Extract + type External = Extract< + Exclude, + { kind: 'generic' } + > + type Unbound = Extract + type CallbackInterrupt = typeof chat.resolveInterrupts extends { + (resolver: (interrupt: infer TInterrupt) => undefined): void + } + ? TInterrupt + : never + + expectTypeOf().toEqualTypeOf< + { title: string } | undefined + >() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + expectTypeOf().toEqualTypeOf() + expectTypeOf< + Extract + >().toEqualTypeOf() + + const resolveReview = (review: Review) => { + review.resolveInterrupt('42') + // @ts-expect-error The transformed response still accepts its input type. + review.resolveInterrupt(42) + } + void resolveReview + chat.resolveInterrupts((interrupt) => { + interrupt.cancel() + return undefined + }) + + const existingTool = toolDefinition({ + name: 'angular-unregistered-tool', + description: 'A tool without an interrupt registry', + needsApproval: true, + }).client() + const withoutRegistry = injectChat({ + connection: { connect: async function* () {} }, + tools: clientTools(existingTool), + }) + type WithoutRegistry = ReturnType< + typeof withoutRegistry.interrupts + >[number] + type ExistingToolInterrupt = Extract< + WithoutRegistry, + { kind: 'tool-approval' } + > + type UnregisteredGeneric = Extract + expectTypeOf< + ExistingToolInterrupt['toolName'] + >().toEqualTypeOf<'angular-unregistered-tool'>() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + } + void check + }) +}) diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index 5f1dcb06ff..de333544fe 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -18,6 +18,7 @@ import { InterruptManager } from './interrupt-manager' import type { AnyClientTool, ContentPart, + InterruptDefinition, InterruptSubmissionError, ModelMessage, RunAgentResumeItem, @@ -42,8 +43,8 @@ import type { ChatClientOptions, ChatClientState, ChatFetcher, - ChatInterrupt, ChatInterruptState, + ResolvableChatInterrupt, ChatPendingInterrupt, ChatResumeSnapshot, ChatResumeState, @@ -66,8 +67,24 @@ interface InternalQueuedMessage extends QueuedMessage { body?: Record } +function assertUniqueInterruptDefinitions( + interrupts: + | ReadonlyArray> + | undefined, +): void { + const ids = new Set() + for (const interrupt of interrupts ?? []) { + if (ids.has(interrupt.id)) { + throw new Error(`Duplicate interrupt definition id: ${interrupt.id}`) + } + ids.add(interrupt.id) + } +} + type ChatClientUpdateOptionsWithoutContext< TTools extends ReadonlyArray, + TInterrupts extends ReadonlyArray> = + readonly [], > = { connection?: ConnectionAdapter fetcher?: ChatFetcher @@ -75,6 +92,7 @@ type ChatClientUpdateOptionsWithoutContext< body?: Record forwardedProps?: Record tools?: TTools + interrupts?: TInterrupts queue?: QueueOption onResponse?: (response?: Response) => void | Promise onChunk?: (chunk: StreamChunk) => void @@ -86,7 +104,7 @@ type ChatClientUpdateOptionsWithoutContext< onQueueChange?: (queue: Array) => void onResumeStateChange?: ( resumeState: ChatResumeState | null, - pendingInterrupts: BoundInterrupts, + pendingInterrupts: BoundInterrupts, ) => void /** * Fires whenever the id of the run in flight changes: the new id when a run @@ -271,6 +289,8 @@ const REJOIN_REBUILD_TRIGGERS = new Set([ export class ChatClient< TTools extends ReadonlyArray = any, TContext = unknown, + TInterrupts extends ReadonlyArray> = + any, > { private readonly processor: StreamProcessor private connection: SubscribeConnectionAdapter @@ -292,7 +312,7 @@ export class ChatClient< // run is rejoined at most once even when both the sync read and the async // hydrate surface the same resume pointer. private rejoinedRunId: string | null = null - private readonly interruptManager: InterruptManager + private readonly interruptManager: InterruptManager private activeInterruptSubmission: InterruptManagerSubmission | undefined private interruptSubmissionFailure: | { errors: ReadonlyArray } @@ -397,10 +417,12 @@ export class ChatClient< onQueueChange: (queue: Array) => void onResumeStateChange: ( resumeState: ChatResumeState | null, - pendingInterrupts: BoundInterrupts, + pendingInterrupts: BoundInterrupts, ) => void onRunIdChange: (runId: string | null) => void - onInterruptStateChange: (state: ChatInterruptState) => void + onInterruptStateChange: ( + state: ChatInterruptState, + ) => void onCustomEvent: ( eventType: string, data: unknown, @@ -409,7 +431,8 @@ export class ChatClient< } } - constructor(options: ChatClientOptions) { + constructor(options: ChatClientOptions) { + assertUniqueInterruptDefinitions(options.interrupts) this.threadId = options.threadId || this.generateUniqueId('thread') // The instance/devtools id defaults to the threadId (the chat's identity), // falling back to a generated id only when neither is set. `id` overrides it @@ -486,8 +509,11 @@ export class ChatClient< }, } - this.interruptManager = new InterruptManager({ + this.interruptManager = new InterruptManager({ ...(options.tools !== undefined ? { tools: options.tools } : {}), + ...(options.interrupts !== undefined + ? { interrupts: options.interrupts } + : {}), submit: (submission) => this.submitInterruptBatch(submission), onChange: () => this.notifyResumeStateChange(), }) @@ -1166,25 +1192,37 @@ export class ChatClient< this.callbacksRef.current.onRunIdChange(runId) } - getInterruptState(): ChatInterruptState { + getInterruptState(): ChatInterruptState { return this.interruptManager.getState() } - getInterrupts(): BoundInterrupts { - return this.interruptManager.getInterrupts() + getInterrupts(): BoundInterrupts { + return this.interruptManager.getInterrupts() as BoundInterrupts< + TTools, + TInterrupts + > } /** @deprecated Use getInterrupts(). */ - getPendingInterrupts(): BoundInterrupts { - return this.interruptManager.getInterrupts() + getPendingInterrupts(): BoundInterrupts { + return this.interruptManager.getInterrupts() as BoundInterrupts< + TTools, + TInterrupts + > } resolveInterrupts(approved: boolean): void resolveInterrupts( - resolver: (interrupt: ChatInterrupt) => undefined, + resolver: ( + interrupt: ResolvableChatInterrupt, + ) => undefined, ): void resolveInterrupts( - resolution: boolean | ((interrupt: ChatInterrupt) => undefined), + resolution: + | boolean + | (( + interrupt: ResolvableChatInterrupt, + ) => undefined), ): void { // Branch so TypeScript can select the InterruptManager.resolve overloads. if (typeof resolution === 'boolean') { @@ -2616,6 +2654,11 @@ export class ChatClient< * a text-only response has nothing to auto-send. */ private shouldAutoSend(): boolean { + // A pending interrupt owns the next send. Auto-continuing after a + // completed server tool would start a sibling run and hide the card. + if (this.lastResume) return false + if (this.activeInterruptSubmission) return false + if (this.interruptManager.getInterrupts().length > 0) return false const messages = this.processor.getMessages() const lastAssistant = messages.findLast( (m: UIMessage) => m.role === 'assistant', diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 69af126b24..a7c4fbedcb 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -430,7 +430,10 @@ async function fetchThreadHydration( const data = (await response.json()) as { messages?: Array activeRun?: { runId?: unknown } | null - interrupts?: { runId?: unknown; pending?: unknown } | null + interrupts?: { + runId?: unknown + pending?: unknown + } | null } const activeRun = data.activeRun && typeof data.activeRun.runId === 'string' @@ -830,7 +833,10 @@ export interface ChatHydrationResult { * so a reload (or another device) re-prompts the approval from the server. The * client restores them exactly as a persisted resume snapshot would. */ - interrupts: { runId: string; pending: Array } | null + interrupts: { + runId: string + pending: Array + } | null } /** @@ -2009,6 +2015,12 @@ 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 7063e93845..d5e617f379 100644 --- a/packages/ai-client/src/index.ts +++ b/packages/ai-client/src/index.ts @@ -38,6 +38,9 @@ export type { ChatInterrupt, ChatInterruptState, GenericAGUIInterrupt, + GenericInterrupt, + RegisteredGenericInterrupt, + ResolvableChatInterrupt, UnboundInterrupt, InterruptItemStatus, ToolApprovalInterrupt, diff --git a/packages/ai-client/src/interrupt-manager.ts b/packages/ai-client/src/interrupt-manager.ts index 1524ad8bd3..fbfef831c3 100644 --- a/packages/ai-client/src/interrupt-manager.ts +++ b/packages/ai-client/src/interrupt-manager.ts @@ -5,15 +5,19 @@ import { canonicalizeInterruptResolutions, cloneAndDeepFreezeJson, digestInterruptJson, + genericInterruptContinuationFromDescriptor, + hashInterruptDefinitionSchema, hashSchemaInput, isStandardSchema, normalizeApprovalSchema, + wrapGenericInterruptContinuation, } from '@tanstack/ai/client' import type { AnyClientTool, BatchInterruptError, Interrupt, InterruptBinding, + InterruptDefinition, InterruptSubmissionError, ItemInterruptError, RunAgentResumeItem, @@ -25,6 +29,7 @@ import type { ChatInterruptState, GenericAGUIInterrupt, InterruptItemStatus, + ResolvableChatInterrupt, UnboundInterrupt, } from './types' @@ -46,8 +51,11 @@ export interface InterruptManagerSubmission { export interface InterruptManagerOptions< TTools extends ReadonlyArray, + TInterrupts extends ReadonlyArray> = + readonly [], > { tools?: TTools + interrupts?: TInterrupts submit: (submission: InterruptManagerSubmission) => Promise onChange?: () => void } @@ -71,6 +79,11 @@ interface RuntimeInterrupt { error?: ItemInterruptError resolution?: RunAgentResumeItem tool?: AnyClientTool + definition?: InterruptDefinition + /** Validated display payload for a registered first-party generic item. */ + payload?: unknown + /** This binding is valid and can participate in this chat resume batch. */ + resumable: boolean validationGeneration: number } @@ -93,6 +106,33 @@ interface RuntimeInterruptCheckpoint { validationGeneration: number } +function isClientOwnedInterrupt(item: RuntimeInterrupt): boolean { + return item.resumable +} + +function resolutionWithContinuation( + item: RuntimeInterrupt, + resolution: RunAgentResumeItem, +): RunAgentResumeItem { + const continuation = genericInterruptContinuationFromDescriptor( + item.descriptor, + ) + if (!continuation) return resolution + return { + ...resolution, + metadata: wrapGenericInterruptContinuation(continuation), + } +} + +function isRootResolvableInterrupt< + TTools extends ReadonlyArray, + TInterrupts extends ReadonlyArray>, +>( + interrupt: ChatInterrupt, +): interrupt is ResolvableChatInterrupt { + return 'cancel' in interrupt && 'clearResolution' in interrupt +} + const itemErrorCodes = new Set([ 'invalid-payload', 'invalid-edited-args', @@ -169,7 +209,6 @@ function isBindingBase(value: UnknownObject): boolean { 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'])))) @@ -181,21 +220,61 @@ function readBinding(value: unknown): InterruptBinding | undefined { const expiresAt = typeof value['expiresAt'] === 'string' ? value['expiresAt'] : undefined if (value['kind'] === 'generic') { + if ( + value['responseSchemaHash'] !== undefined && + typeof value['responseSchemaHash'] !== 'string' + ) { + return undefined + } + const firstPartyFields = [ + value['definitionId'], + value['key'], + value['batchIndex'], + value['payloadSchemaHash'], + ] + const hasFirstPartyFields = firstPartyFields.some( + (field) => field !== undefined, + ) + if ( + hasFirstPartyFields && + (typeof value['definitionId'] !== 'string' || + typeof value['key'] !== 'string' || + typeof value['batchIndex'] !== 'number' || + !Number.isInteger(value['batchIndex']) || + value['batchIndex'] < 0 || + (value['payloadSchemaHash'] !== undefined && + typeof value['payloadSchemaHash'] !== 'string')) + ) { + return undefined + } return { v: INTERRUPT_BINDING_VERSION, kind: 'generic', interruptId: String(value['interruptId']), interruptedRunId: String(value['interruptedRunId']), generation: Number(value['generation']), - responseSchemaHash: String(value['responseSchemaHash']), + ...(typeof value['responseSchemaHash'] === 'string' + ? { responseSchemaHash: value['responseSchemaHash'] } + : {}), ...(expiresAt !== undefined ? { expiresAt } : {}), + ...(hasFirstPartyFields + ? { + definitionId: String(value['definitionId']), + key: String(value['key']), + batchIndex: Number(value['batchIndex']), + ...(typeof value['payloadSchemaHash'] === 'string' + ? { payloadSchemaHash: value['payloadSchemaHash'] } + : {}), + } + : {}), } } if ( value['kind'] === 'client-tool-execution' && typeof value['toolName'] === 'string' && typeof value['toolCallId'] === 'string' && - typeof value['outputSchemaHash'] === 'string' + typeof value['outputSchemaHash'] === 'string' && + typeof value['responseSchemaHash'] === 'string' ) { return { v: INTERRUPT_BINDING_VERSION, @@ -216,6 +295,7 @@ function readBinding(value: unknown): InterruptBinding | undefined { typeof value['toolCallId'] === 'string' && typeof value['inputSchemaHash'] === 'string' && typeof value['approvalSchemaHash'] === 'string' && + typeof value['responseSchemaHash'] === 'string' && 'originalArgs' in value ) { return { @@ -244,6 +324,43 @@ function getDescriptorBinding( return readBinding(candidate) } +function hasReservedFirstPartyBindingMarker(interrupt: Interrupt): boolean { + if (!isUnknownObject(interrupt.metadata)) return false + const binding = interrupt.metadata[INTERRUPT_BINDING_METADATA_KEY] + if (!isUnknownObject(binding)) return false + if ( + binding['v'] !== undefined && + binding['v'] !== INTERRUPT_BINDING_VERSION + ) { + return false + } + return ( + binding['kind'] === 'generic' || + binding['kind'] === 'tool-approval' || + binding['kind'] === 'client-tool-execution' || + 'definitionId' in binding || + 'key' in binding || + 'batchIndex' in binding || + 'payloadSchemaHash' in binding + ) +} + +function hasFirstPartyGenericMarker(interrupt: Interrupt): boolean { + if (!isUnknownObject(interrupt.metadata)) return false + const binding = interrupt.metadata[INTERRUPT_BINDING_METADATA_KEY] + if (!isUnknownObject(binding) || binding['kind'] !== 'generic') return false + return ( + 'definitionId' in binding || + 'key' in binding || + 'batchIndex' in binding || + 'payloadSchemaHash' in binding + ) +} + +function getInterruptPayload(interrupt: Interrupt): unknown { + return interrupt.metadata?.['tanstack:interruptPayload'] +} + /** * Only used to route *legacy* (pre-binding) descriptors, which have no binding * to classify off. Current descriptors are classified by their binding alone. @@ -264,6 +381,17 @@ function responseSchemaHash(interrupt: Interrupt): string | undefined { } } +function definitionSchemaHash( + schema: InterruptDefinition['responseSchema'] | undefined, +): string | undefined { + if (schema === undefined) return undefined + try { + return hashInterruptDefinitionSchema(schema) + } catch { + return undefined + } +} + function isPromiseLike(value: unknown): value is PromiseLike { return ( value !== null && @@ -431,16 +559,15 @@ function genericBinding( hydration: InterruptManagerHydration, candidate: InterruptBinding | undefined, ): InterruptBinding { + const schemaHash = + responseSchemaHash(interrupt) ?? candidate?.responseSchemaHash return cloneAndDeepFreezeJson({ v: INTERRUPT_BINDING_VERSION, kind: 'generic', interruptId: interrupt.id, interruptedRunId: hydration.interruptedRunId, generation: hydration.generation, - responseSchemaHash: - responseSchemaHash(interrupt) ?? - candidate?.responseSchemaHash ?? - 'invalid', + ...(schemaHash !== undefined ? { responseSchemaHash: schemaHash } : {}), ...(interrupt.expiresAt !== undefined ? { expiresAt: interrupt.expiresAt } : {}), @@ -450,9 +577,7 @@ function genericBinding( function baseSnapshot( item: RuntimeInterrupt, hydration: InterruptManagerHydration, - cancel: () => void, - clearResolution: () => void, -): BoundInterruptBase { +): Omit { const descriptor = cloneAndDeepFreezeJson(item.descriptor) const errors: ReadonlyArray = item.error === undefined @@ -482,21 +607,22 @@ function baseSnapshot( errors, ...(error !== undefined ? { error } : {}), canResolve: item.canResolve, - cancel, - clearResolution, } } export class InterruptManager< TTools extends ReadonlyArray = ReadonlyArray, + TInterrupts extends ReadonlyArray> = + readonly [], > { private hydration: InterruptManagerHydration | undefined private items: Array = [] - private snapshot: ReadonlyArray> = Object.freeze([]) + private snapshot: ReadonlyArray> = + Object.freeze([]) private rootErrors: ReadonlyArray = Object.freeze([]) private submissionRootErrors: ReadonlyArray = Object.freeze([]) - private state: ChatInterruptState = Object.freeze({ + private state: ChatInterruptState = Object.freeze({ interrupts: this.snapshot, pendingInterrupts: this.snapshot, interruptErrors: this.rootErrors, @@ -506,9 +632,26 @@ export class InterruptManager< private retrySubmission: InterruptManagerSubmission | undefined private resuming = false private tools: TTools | undefined + private readonly interruptDefinitions: ReadonlyMap< + string, + InterruptDefinition + > - constructor(private readonly options: InterruptManagerOptions) { + constructor( + private readonly options: InterruptManagerOptions, + ) { this.tools = options.tools + const definitions = new Map< + string, + InterruptDefinition + >() + for (const definition of options.interrupts ?? []) { + if (definitions.has(definition.id)) { + throw new Error(`Duplicate interrupt definition id: ${definition.id}`) + } + definitions.set(definition.id, definition) + } + this.interruptDefinitions = definitions } updateTools(tools: TTools): void { @@ -522,8 +665,23 @@ export class InterruptManager< generation: hydration.generation, interrupts: cloneAndDeepFreezeJson(hydration.interrupts), } + const firstPartyIndexes = new Map() + for (const interrupt of hydration.interrupts) { + const binding = getDescriptorBinding(interrupt) + if ( + binding?.kind === 'generic' && + binding.definitionId !== undefined && + binding.key !== undefined && + binding.batchIndex !== undefined + ) { + firstPartyIndexes.set( + binding.batchIndex, + (firstPartyIndexes.get(binding.batchIndex) ?? 0) + 1, + ) + } + } this.items = hydration.interrupts.map((interrupt) => - this.hydrateInterrupt(interrupt, hydration), + this.hydrateInterrupt(interrupt, hydration, firstPartyIndexes), ) this.rootErrors = Object.freeze([]) this.submissionRootErrors = Object.freeze([]) @@ -532,11 +690,11 @@ export class InterruptManager< this.publish() } - getInterrupts(): BoundInterrupts { + getInterrupts(): BoundInterrupts { return this.snapshot } - getState(): ChatInterruptState { + getState(): ChatInterruptState { return this.state } @@ -544,6 +702,12 @@ export class InterruptManager< return this.hydration?.interrupts ?? Object.freeze([]) } + hasValidatedFirstPartyGenericBatch(): boolean { + return this.items.some( + (item) => item.kind === 'generic' && item.definition !== undefined, + ) + } + reset(options?: { preserveRootErrors?: boolean }): void { this.hydration = undefined this.items = [] @@ -572,9 +736,15 @@ export class InterruptManager< } resolve(approved: boolean): void - resolve(resolver: (interrupt: ChatInterrupt) => undefined): void resolve( - resolution: boolean | ((interrupt: ChatInterrupt) => unknown), + resolver: ( + interrupt: ResolvableChatInterrupt, + ) => undefined, + ): void + resolve( + resolution: + | boolean + | ((interrupt: ResolvableChatInterrupt) => unknown), ): void { this.assertRootMutable() if (typeof resolution === 'boolean') { @@ -588,11 +758,14 @@ export class InterruptManager< this.assertRootMutable() this.invalidateRetry() for (const item of this.items) { + if (!isClientOwnedInterrupt(item)) continue item.validationGeneration++ - item.resolution = Object.freeze({ - interruptId: item.descriptor.id, - status: 'cancelled', - }) + item.resolution = Object.freeze( + resolutionWithContinuation(item, { + interruptId: item.descriptor.id, + status: 'cancelled', + }), + ) item.status = 'staged' item.error = undefined } @@ -648,9 +821,12 @@ export class InterruptManager< private hydrateInterrupt( descriptor: Interrupt, hydration: InterruptManagerHydration, + firstPartyIndexes: ReadonlyMap, ): RuntimeInterrupt { const interrupt = cloneAndDeepFreezeJson(descriptor) const candidate = getDescriptorBinding(interrupt) + const legacyResumable = + candidate === undefined && isLegacyInterruptMetadata(interrupt) // No binding we understand, and nothing else identifying the descriptor as // ours, means this interrupt was not produced by this package's resume @@ -668,18 +844,42 @@ export class InterruptManager< // Pre-binding TanStack descriptors are still ours: they carry the legacy // `metadata.kind` marker, so they keep hydrating through the generic path // below. - if (candidate === undefined && !isLegacyInterruptMetadata(interrupt)) { + if (candidate === undefined && !legacyResumable) { + if (hasReservedFirstPartyBindingMarker(interrupt)) { + return { + descriptor: interrupt, + binding: genericBinding(interrupt, hydration, undefined), + kind: 'generic', + status: 'error', + canResolve: false, + resumable: false, + error: this.itemError( + interrupt.id, + 'stale', + 'The interrupt binding is invalid or incomplete.', + ), + validationGeneration: 0, + } + } return { descriptor: interrupt, binding: undefined, kind: 'unbound', status: 'pending', canResolve: false, + resumable: false, validationGeneration: 0, } } const correlated = + candidate !== undefined && + candidate.interruptId === interrupt.id && + candidate.interruptedRunId === hydration.interruptedRunId && + candidate.generation === hydration.generation && + responseSchemaHash(interrupt) === candidate.responseSchemaHash + + const structurallyCorrelated = candidate !== undefined && candidate.interruptId === interrupt.id && candidate.interruptedRunId === hydration.interruptedRunId && @@ -687,7 +887,32 @@ export class InterruptManager< candidate.responseSchemaHash === (responseSchemaHash(interrupt) ?? candidate.responseSchemaHash) - if (correlated && candidate.kind === 'tool-approval') { + if ( + candidate !== undefined && + hasFirstPartyGenericMarker(interrupt) && + (!structurallyCorrelated || + candidate.kind !== 'generic' || + candidate.definitionId === undefined || + candidate.key === undefined || + candidate.batchIndex === undefined) + ) { + return { + descriptor: interrupt, + binding: genericBinding(interrupt, hydration, candidate), + kind: 'generic', + status: 'error', + canResolve: false, + resumable: false, + error: this.itemError( + interrupt.id, + 'stale', + 'The interrupt binding does not match this interrupted run.', + ), + validationGeneration: 0, + } + } + + if (structurallyCorrelated && candidate.kind === 'tool-approval') { const tool = this.tools?.find( (configured) => configured.name === candidate.toolName, ) @@ -714,6 +939,7 @@ export class InterruptManager< kind: 'tool-approval', status: 'pending', canResolve: true, + resumable: true, tool, validationGeneration: 0, } @@ -724,7 +950,7 @@ export class InterruptManager< } } - if (correlated && candidate.kind === 'client-tool-execution') { + if (structurallyCorrelated && candidate.kind === 'client-tool-execution') { const tool = this.tools?.find( (configured) => configured.name === candidate.toolName, ) @@ -740,31 +966,112 @@ export class InterruptManager< kind: 'client-tool-execution', status: 'pending', canResolve: true, + resumable: true, tool, validationGeneration: 0, } } } + if ( + candidate !== undefined && + candidate.kind === 'generic' && + candidate.definitionId !== undefined && + candidate.key !== undefined && + candidate.batchIndex !== undefined && + candidate.key.length > 0 && + firstPartyIndexes.get(candidate.batchIndex) !== 1 + ) { + return { + descriptor: interrupt, + binding: cloneAndDeepFreezeJson(candidate), + kind: 'generic', + status: 'error', + canResolve: false, + resumable: false, + error: this.itemError( + interrupt.id, + 'stale', + 'Generic interrupt batch contains a duplicate batchIndex.', + ), + validationGeneration: 0, + } + } + + if ( + correlated && + candidate.kind === 'generic' && + candidate.definitionId !== undefined && + candidate.key !== undefined && + candidate.batchIndex !== undefined && + candidate.key.length > 0 && + firstPartyIndexes.get(candidate.batchIndex) === 1 + ) { + const definition = this.interruptDefinitions.get(candidate.definitionId) + if ( + definition !== undefined && + definitionSchemaHash(definition.responseSchema) === + candidate.responseSchemaHash && + (definition.payloadSchema === undefined + ? candidate.payloadSchemaHash === undefined + : candidate.payloadSchemaHash === + definitionSchemaHash(definition.payloadSchema)) + ) { + const rawPayload = getInterruptPayload(interrupt) + // First-party display payloads are parsed by definition.interrupt() + // before the server emits them. Re-validating here would feed schema + // output back through an input schema and reject transforms such as + // z.string().transform(Number). The checks above still bind this value + // to the exact descriptor, run, generation, definition, and schemas. + return { + descriptor: interrupt, + binding: cloneAndDeepFreezeJson(candidate), + definition, + kind: 'generic', + status: 'pending', + canResolve: true, + resumable: true, + ...(rawPayload === undefined + ? {} + : { payload: cloneAndDeepFreezeJson(rawPayload) }), + validationGeneration: 0, + } + } + } + + const resumable = + legacyResumable || + (candidate !== undefined && + candidate.interruptId === interrupt.id && + candidate.interruptedRunId === hydration.interruptedRunId && + candidate.generation === hydration.generation && + (candidate.kind !== 'generic' || + responseSchemaHash(interrupt) === undefined || + candidate.responseSchemaHash === responseSchemaHash(interrupt))) return { descriptor: interrupt, binding: genericBinding(interrupt, hydration, candidate), kind: 'generic', status: 'pending', - // The library no longer validates the wire response schema, so a generic - // item is always resolvable. The application validates the value itself. - canResolve: true, + // A valid raw binding is an explicit request to use this resume path, + // even when this client has no registered first-party definition. Keep + // it untyped, but preserve its existing generic controls. Missing, + // malformed, and unsupported bindings remain display-only. + canResolve: resumable, + resumable, validationGeneration: 0, } } private buildSnapshot( transaction?: TransactionToken, - ): BoundInterrupts { + ): BoundInterrupts { const hydration = this.requireHydration() - // `client-tool-execution` items stay in `this.items` (they gate batch - // submission and are resolved internally via auto-execution / addToolResult), - // but they are never surfaced as public bound interrupts. + // `client-tool-execution` items stay in `this.items` (they usually gate + // batch submission and are resolved internally via auto-execution / + // addToolResult), but they are never surfaced as public bound interrupts. + // A mixed generic batch is the exception: those client tools wait for + // `toolResume` and must not block submit. // // Items with status `submitting` are also omitted: the resume stream is // already in flight, so Approve/Deny is not actionable. Keeping them in @@ -776,12 +1083,7 @@ export class InterruptManager< item.kind !== 'client-tool-execution' && item.status !== 'submitting', ) .map((item) => { - const base = baseSnapshot( - item, - hydration, - () => this.cancelItem(item.descriptor.id, transaction), - () => this.clearItem(item.descriptor.id, transaction), - ) + const base = baseSnapshot(item, hydration) // Not ours to resume: expose the descriptor so a UI can show the run // is paused, with no `resolveInterrupt` to call. if (item.kind === 'unbound' || item.binding === undefined) { @@ -804,6 +1106,9 @@ export class InterruptManager< toolName: item.binding.toolName, toolCallId: item.binding.toolCallId, originalArgs: cloneAndDeepFreezeJson(item.binding.originalArgs), + cancel: () => this.cancelItem(item.descriptor.id, transaction), + clearResolution: () => + this.clearItem(item.descriptor.id, transaction), resolveInterrupt: (approved: boolean, options?: unknown) => { const details = isUnknownObject(options) ? options : undefined this.resolveItem( @@ -832,15 +1137,33 @@ export class InterruptManager< interruptId: item.descriptor.id, interruptedRunId: hydration.interruptedRunId, generation: hydration.generation, - responseSchemaHash: - typeof item.binding.responseSchemaHash === 'string' - ? item.binding.responseSchemaHash - : 'none', + ...(typeof item.binding.responseSchemaHash === 'string' + ? { responseSchemaHash: item.binding.responseSchemaHash } + : {}), }) + if (item.definition !== undefined && item.binding.kind === 'generic') { + const snapshot = { + ...base, + kind: 'generic', + definitionId: item.definition.id, + key: item.binding.key ?? '', + payload: item.payload, + binding: boundGeneric, + cancel: () => this.cancelItem(item.descriptor.id, transaction), + clearResolution: () => + this.clearItem(item.descriptor.id, transaction), + resolveInterrupt: (response: unknown) => + this.resolveItem(item.descriptor.id, response, transaction), + } + return Object.freeze(snapshot) + } const snapshot: GenericAGUIInterrupt = { ...base, kind: 'generic', binding: boundGeneric, + cancel: () => this.cancelItem(item.descriptor.id, transaction), + clearResolution: () => + this.clearItem(item.descriptor.id, transaction), resolveInterrupt: (payload) => this.resolveItem(item.descriptor.id, payload, transaction), } @@ -851,7 +1174,7 @@ export class InterruptManager< // 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 + return Object.freeze(next) as BoundInterrupts } private publish(): void { @@ -929,7 +1252,9 @@ export class InterruptManager< const item = this.findItem(interruptId) this.invalidateRetry() item.validationGeneration++ - item.resolution = Object.freeze({ interruptId, status: 'cancelled' }) + item.resolution = Object.freeze( + resolutionWithContinuation(item, { interruptId, status: 'cancelled' }), + ) item.status = 'staged' item.error = undefined if (!transaction) { @@ -954,7 +1279,12 @@ export class InterruptManager< // owns them. Including them in the completeness gate would deadlock the // batch, so the run's own interrupts could never be answered once a // foreign one shared the stream. - const ours = this.items.filter((item) => item.kind !== 'unbound') + const hasGeneric = this.items.some((item) => item.kind === 'generic') + const ours = this.items.filter( + (item) => + isClientOwnedInterrupt(item) && + !(hasGeneric && item.kind === 'client-tool-execution'), + ) if ( ours.length === 0 || ours.some( @@ -1019,11 +1349,13 @@ export class InterruptManager< if (!transaction) this.publish() return } - item.resolution = cloneAndDeepFreezeJson({ - interruptId: item.descriptor.id, - status: 'resolved', - payload: result.payload, - }) + item.resolution = cloneAndDeepFreezeJson( + resolutionWithContinuation(item, { + interruptId: item.descriptor.id, + status: 'resolved', + payload: result.payload, + }), + ) item.status = 'staged' item.error = undefined if (!transaction) { @@ -1037,11 +1369,17 @@ export class InterruptManager< payload: unknown, ): ValidationResult | Promise { if (item.kind === 'generic') { - return validateWithSchema( - item.descriptor.responseSchema, + const validation = validateWithSchema( + item.definition?.responseSchema ?? item.descriptor.responseSchema, payload, 'invalid-payload', ) + if (item.definition === undefined) return validation + const preserveInput = (result: ValidationResult): ValidationResult => + 'valid' in result ? { valid: true, payload } : result + return isPromiseLike(validation) + ? Promise.resolve(validation).then(preserveInput) + : preserveInput(validation) } if (item.kind === 'client-tool-execution') { return validateWithSchema( @@ -1160,7 +1498,8 @@ export class InterruptManager< // addToolResult); they are transparent to the boolean shorthand. Eligibility // and resolution consider only the publicly resolvable items. const resolvable = this.items.filter( - (item) => item.kind !== 'client-tool-execution', + (item) => + isClientOwnedInterrupt(item) && item.kind !== 'client-tool-execution', ) const eligible = resolvable.every( (item) => @@ -1191,7 +1530,9 @@ export class InterruptManager< } private resolveTransaction( - resolver: (interrupt: ChatInterrupt) => unknown, + resolver: ( + interrupt: ResolvableChatInterrupt, + ) => unknown, ): void { const checkpoints = this.items.map((item) => ({ status: item.status, @@ -1201,7 +1542,7 @@ export class InterruptManager< })) const token: TransactionToken = { active: true } this.activeTransaction = token - const stable = this.buildSnapshot(token) + const stable = this.buildSnapshot(token).filter(isRootResolvableInterrupt) let failure: | { code: BatchInterruptError['code']; message: string } | undefined @@ -1227,7 +1568,8 @@ export class InterruptManager< // `client-tool-execution` items are resolved out-of-band (auto // execution / addToolResult), not by this synchronous resolver, so // they don't count against transaction completeness. `maybeSubmit` - // still gates the actual submission on them being resolved. + // still waits for them unless a generic interrupt shares the batch. + isClientOwnedInterrupt(item) && item.kind !== 'client-tool-execution' && (item.resolution === undefined || item.status !== 'staged'), ) @@ -1298,7 +1640,9 @@ export class InterruptManager< private submitBatch(submission: InterruptManagerSubmission): void { this.resuming = true this.retrySubmission = undefined - for (const item of this.items) item.status = 'submitting' + for (const item of this.items) { + if (isClientOwnedInterrupt(item)) item.status = 'submitting' + } this.publish() void this.performSubmission(submission) } @@ -1325,7 +1669,9 @@ export class InterruptManager< 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' + for (const item of this.items) { + if (isClientOwnedInterrupt(item)) item.status = 'error' + } return } @@ -1355,6 +1701,7 @@ export class InterruptManager< if (submissionError.scope === 'item') { const item = this.items.find( (candidate) => + isClientOwnedInterrupt(candidate) && candidate.descriptor.id === submissionError.interruptId, ) if (item) { @@ -1373,7 +1720,9 @@ export class InterruptManager< this.rootErrors = mergedBatchErrors.rootErrors this.submissionRootErrors = mergedBatchErrors.submissionRootErrors for (const item of this.items) { - if (item.status === 'submitting') item.status = 'error' + if (isClientOwnedInterrupt(item) && item.status === 'submitting') { + item.status = 'error' + } } this.retrySubmission = retryable && !nonRetryable ? submission : undefined } @@ -1394,7 +1743,9 @@ export class InterruptManager< source, retryable, interruptIds: Object.freeze( - this.items.map((item) => item.descriptor.id), + this.items + .filter(isClientOwnedInterrupt) + .map((item) => item.descriptor.id), ), threadId: hydration.threadId, interruptedRunId: hydration.interruptedRunId, diff --git a/packages/ai-client/src/types.ts b/packages/ai-client/src/types.ts index 75ff059763..625e94a194 100644 --- a/packages/ai-client/src/types.ts +++ b/packages/ai-client/src/types.ts @@ -9,6 +9,7 @@ import type { DocumentPart, ImagePart, InferSchemaType, + InterruptDefinition, InferToolInput, InferToolOutput, InputSchemaOf, @@ -86,6 +87,51 @@ export interface GenericAGUIInterrupt extends BoundInterruptBase { resolveInterrupt: (payload: unknown) => void } +type InterruptResponseInput = + TDefinition extends InterruptDefinition + ? InferSchemaType + : never + +type RegisteredGenericInterruptFor< + TDefinition extends InterruptDefinition, +> = + TDefinition extends InterruptDefinition< + infer TDefinitionId, + any, + any, + infer TPayload + > + ? BoundInterruptBase & { + readonly kind: 'generic' + readonly definitionId: TDefinitionId + readonly key: string + readonly payload: TPayload | undefined + readonly binding: Readonly< + Extract & { + definitionId: TDefinitionId + key: string + batchIndex: number + } + > + resolveInterrupt: ( + response: InterruptResponseInput, + ) => void + } + : never + +export type RegisteredGenericInterrupt< + TInterrupts extends ReadonlyArray>, +> = TInterrupts[number] extends infer TDefinition + ? TDefinition extends InterruptDefinition + ? RegisteredGenericInterruptFor + : never + : never + +/** A bound generic interrupt for one `defineInterrupt()` definition. */ +export type GenericInterrupt< + TDefinition extends InterruptDefinition, +> = RegisteredGenericInterruptFor + /** * An interrupt that arrived on the stream carrying no resume binding this * client understands — no `tanstack:interruptBinding`, or one written at a @@ -98,7 +144,10 @@ export interface GenericAGUIInterrupt extends BoundInterruptBase { * send an answer no one is waiting for. Render it, or route it to whatever * actually owns the pause. */ -export interface UnboundInterrupt extends BoundInterruptBase { +export interface UnboundInterrupt extends Omit< + BoundInterruptBase, + 'cancel' | 'clearResolution' +> { readonly kind: 'unbound' readonly binding?: undefined readonly canResolve: false @@ -188,18 +237,37 @@ type ApprovalInterrupts> = // union. export type ChatInterrupt< TTools extends ReadonlyArray = ReadonlyArray, -> = GenericAGUIInterrupt | UnboundInterrupt | ApprovalInterrupts + TInterrupts extends ReadonlyArray> = + readonly [], +> = + | GenericAGUIInterrupt + | RegisteredGenericInterrupt + | UnboundInterrupt + | ApprovalInterrupts + +export type ResolvableChatInterrupt< + TTools extends ReadonlyArray = ReadonlyArray, + TInterrupts extends ReadonlyArray> = + readonly [], +> = + | GenericAGUIInterrupt + | RegisteredGenericInterrupt + | ApprovalInterrupts export type BoundInterrupts< TTools extends ReadonlyArray = ReadonlyArray, -> = ReadonlyArray> + TInterrupts extends ReadonlyArray> = + readonly [], +> = ReadonlyArray> export interface ChatInterruptState< TTools extends ReadonlyArray = ReadonlyArray, + TInterrupts extends ReadonlyArray> = + readonly [], > { - readonly interrupts: BoundInterrupts + readonly interrupts: BoundInterrupts /** @deprecated Use `interrupts`. Same snapshot today. */ - readonly pendingInterrupts: BoundInterrupts + readonly pendingInterrupts: BoundInterrupts readonly interruptErrors: ReadonlyArray readonly resuming: boolean } @@ -712,6 +780,8 @@ export type ClientContextOptionFromTools = [ export interface ChatClientBaseOptions< TTools extends ReadonlyArray = any, TContext = unknown, + TInterrupts extends ReadonlyArray> = + readonly [], > { /** * Initial messages to populate the chat @@ -871,7 +941,7 @@ export interface ChatClientBaseOptions< */ onResumeStateChange?: ( resumeState: ChatResumeState | null, - pendingInterrupts: BoundInterrupts, + pendingInterrupts: BoundInterrupts, ) => void /** @@ -881,7 +951,9 @@ export interface ChatClientBaseOptions< onRunIdChange?: (runId: string | null) => void /** Callback when the immutable interrupt state snapshot changes. */ - onInterruptStateChange?: (state: ChatInterruptState) => void + onInterruptStateChange?: ( + state: ChatInterruptState, + ) => void /** * Callback when a custom event is received from a server-side tool. @@ -903,6 +975,9 @@ export interface ChatClientBaseOptions< */ tools?: TTools + /** First-party generic interrupts this client can type and resolve. */ + interrupts?: TInterrupts + /** * Devtools hook metadata for this client instance. */ @@ -937,7 +1012,12 @@ export interface ChatClientBaseOptions< export type ChatClientOptions< TTools extends ReadonlyArray = any, TContext = InferredClientContext, -> = DistributedOmit, 'context'> & + TInterrupts extends ReadonlyArray> = + readonly [], +> = DistributedOmit< + ChatClientBaseOptions, + 'context' +> & ClientContextOptionFromTools & ChatTransport @@ -986,9 +1066,12 @@ export function clientTools>( export function createChatClientOptions< const TTools extends ReadonlyArray, TContext = InferredClientContext, + const TInterrupts extends ReadonlyArray< + InterruptDefinition + > = readonly [], >( - options: ChatClientOptions, -): ChatClientOptions { + options: ChatClientOptions, +): ChatClientOptions { return options } diff --git a/packages/ai-client/tests/chat-client-interrupts.test.ts b/packages/ai-client/tests/chat-client-interrupts.test.ts index 89f27e086f..5fe406d221 100644 --- a/packages/ai-client/tests/chat-client-interrupts.test.ts +++ b/packages/ai-client/tests/chat-client-interrupts.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it, vi } from 'vitest' -import { chat } from '@tanstack/ai' +import { chat, createInterruptBinding, defineInterrupt } from '@tanstack/ai' import { EventType, canonicalInterruptJson, @@ -25,7 +25,7 @@ import type { } 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 { ResolvableChatInterrupt } from '../src/types' import type { ConnectConnectionAdapter, RunAgentInputContext, @@ -150,6 +150,58 @@ describe('InterruptManager foreign-interrupt handling', () => { expect(item?.canResolve).toBe(false) }) + it('keeps a raw generic binding resolvable when run id and schema hash match', () => { + const responseSchema = { + type: 'object', + properties: { confirmed: { type: 'boolean' } }, + required: ['confirmed'], + } + const { manager } = createManager() + manager.hydrate({ + threadId: 'foreign-1', + interruptedRunId: 'run-foreign-1', + generation: 0, + interrupts: [ + { + id: 'ours', + reason: 'confirmation', + message: 'Confirm the shipment?', + responseSchema, + metadata: { + 'tanstack:interruptBinding': { + v: INTERRUPT_BINDING_VERSION, + kind: 'generic', + interruptId: 'ours', + interruptedRunId: 'run-foreign-1', + generation: 0, + responseSchemaHash: digestInterruptJson( + canonicalInterruptJson(responseSchema), + ), + }, + }, + }, + { + id: 'theirs', + reason: 'approval_requested', + message: 'Approve the deployment?', + metadata: { 'acme:workflowApproval': { stepId: 'deploy' } }, + }, + ], + }) + + const [ours, theirs] = manager.getInterrupts() + expect(ours).toMatchObject({ + kind: 'generic', + canResolve: true, + id: 'ours', + }) + expect(theirs).toMatchObject({ + kind: 'unbound', + canResolve: false, + id: 'theirs', + }) + }) + it('does not let an unbound interrupt block submission of the bound ones', async () => { const { manager, submit } = createManager() manager.hydrate({ @@ -452,13 +504,19 @@ describe('InterruptManager hydration', () => { }) it('resolves a generic item regardless of its wire response schema', () => { + const responseSchema = { + $schema: 'https://json-schema.org/draft/2019-09/schema', + type: 'object', + } const binding: InterruptBinding = { v: INTERRUPT_BINDING_VERSION, kind: 'generic', interruptId: 'generic-1', interruptedRunId: 'run-1', generation: 1, - responseSchemaHash: 'any-schema', + responseSchemaHash: digestInterruptJson( + canonicalInterruptJson(responseSchema), + ), } const { manager } = createManager() manager.hydrate({ @@ -470,24 +528,180 @@ describe('InterruptManager hydration', () => { // The library does not compile or validate the wire schema, so even a // schema in another dialect leaves the item resolvable. The // application validates the value itself before resolving. - responseSchema: { - $schema: 'https://json-schema.org/draft/2019-09/schema', - type: 'object', - }, + responseSchema, }), ], }) const item = manager.getInterrupts()[0] - expect(item?.kind).toBe('generic') - expect(item?.canResolve).toBe(true) - item?.cancel() + if (item?.kind !== 'generic') throw new Error('Expected generic interrupt') + expect(item.canResolve).toBe(true) + item.cancel() // Cancellation immediately submits the batch; submitting items are omitted // from the public interrupt list (not user-actionable while the resume // stream is in flight). expect(manager.getInterrupts()).toEqual([]) expect(manager.getResuming()).toBe(true) }) + + it('does not treat a generic binding as resumable when the schema hash does not match', () => { + const binding: InterruptBinding = { + v: INTERRUPT_BINDING_VERSION, + kind: 'generic', + interruptId: 'generic-1', + interruptedRunId: 'run-1', + generation: 1, + responseSchemaHash: 'other-hash', + } + const { manager } = createManager() + manager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + interrupts: [ + descriptor(binding, { + responseSchema: { type: 'object' }, + }), + ], + }) + + const item = manager.getInterrupts()[0] + expect(item).toMatchObject({ kind: 'generic', canResolve: false }) + }) + + it('records a stale error when two first-party items share a batchIndex', () => { + const review = defineInterrupt({ + id: 'approval', + responseSchema: z.object({ answer: z.number() }), + }) + const responseSchema = convertSchemaToJsonSchema(review.responseSchema) + const hash = digestInterruptJson(canonicalInterruptJson(responseSchema)) + const first: InterruptBinding = { + v: INTERRUPT_BINDING_VERSION, + kind: 'generic', + interruptId: 'generic-1', + interruptedRunId: 'run-1', + generation: 1, + definitionId: 'approval', + key: 'one', + batchIndex: 0, + responseSchemaHash: hash, + } + const second: InterruptBinding = { + ...first, + interruptId: 'generic-2', + key: 'two', + } + const manager = new InterruptManager({ + tools, + interrupts: [review], + submit: vi.fn(async () => undefined), + }) + manager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 1, + interrupts: [ + descriptor(first, { responseSchema }), + descriptor(second, { responseSchema }), + ], + }) + + for (const item of manager.getInterrupts()) { + expect(item).toMatchObject({ + kind: 'generic', + canResolve: false, + status: 'error', + }) + if (item.kind === 'generic') { + expect(item.errors[0]?.code).toBe('stale') + } + } + }) + + it('hydrates a first-party item from the producer binding and schema hash', () => { + const review = defineInterrupt({ + id: 'review-plan', + payloadSchema: z.object({ + title: z.string(), + boundary: z.enum([ + 'beforeModel', + 'afterModel', + 'beforeTools', + 'afterTools', + ]), + }), + responseSchema: z.object({ + approved: z.boolean(), + note: z.string(), + }), + }) + const request = review.interrupt({ + key: 'generic-before-model-review', + reason: 'review_required', + message: 'Review the plan at beforeModel', + payload: { + title: 'Middleware review plan', + boundary: 'beforeModel' as const, + }, + }) + const emission = createInterruptBinding(request, { batchIndex: 0 }) + const interruptId = 'interrupt-1' + const manager = new InterruptManager({ + tools, + interrupts: [review], + submit: vi.fn(async () => undefined), + }) + manager.hydrate({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 0, + interrupts: [ + { + id: interruptId, + reason: request.reason, + message: request.message, + ...(emission.descriptor.responseSchemaCanonicalJson !== undefined + ? { + responseSchema: JSON.parse( + emission.descriptor.responseSchemaCanonicalJson, + ), + } + : {}), + metadata: { + 'tanstack:interruptBinding': { + v: INTERRUPT_BINDING_VERSION, + kind: 'generic', + interruptId, + interruptedRunId: 'run-1', + generation: 0, + definitionId: emission.descriptor.definitionId, + key: emission.descriptor.key, + batchIndex: emission.descriptor.batchIndex, + ...(emission.descriptor.payloadSchemaHash !== undefined + ? { payloadSchemaHash: emission.descriptor.payloadSchemaHash } + : {}), + ...(emission.descriptor.responseSchemaHash !== undefined + ? { + responseSchemaHash: emission.descriptor.responseSchemaHash, + } + : {}), + }, + ...(emission.payload !== undefined + ? { 'tanstack:interruptPayload': emission.payload } + : {}), + }, + }, + ], + }) + + const item = manager.getInterrupts()[0] + expect(item).toMatchObject({ + kind: 'generic', + canResolve: true, + definitionId: 'review-plan', + }) + }) }) describe('InterruptManager transactions', () => { @@ -505,7 +719,10 @@ describe('InterruptManager transactions', () => { first.resolveInterrupt('first') expect(multi.submit).not.toHaveBeenCalled() expect(multi.manager.getInterrupts()[0]?.status).toBe('staged') - multi.manager.getInterrupts()[1]?.cancel() + const second = multi.manager.getInterrupts()[1] + if (second?.kind !== 'generic') + throw new Error('Expected generic interrupt') + second.cancel() expect(multi.submit).toHaveBeenCalledTimes(1) const single = createManager() @@ -602,7 +819,7 @@ describe('InterruptManager transactions', () => { throw new Error('transaction failed') }) Reflect.apply(manager.resolve, manager, [ - (item: ChatInterrupt) => { + (item: ResolvableChatInterrupt) => { item.cancel() return 'not undefined' }, @@ -743,7 +960,10 @@ describe('InterruptManager transactions', () => { manager.retry() expect(submit.mock.calls[1]?.[0]).toBe(firstSubmission) await settle() - manager.getInterrupts()[0]?.clearResolution() + const cleared = manager.getInterrupts()[0] + if (cleared && cleared.kind !== 'unbound') { + cleared.clearResolution() + } manager.retry() expect(submit).toHaveBeenCalledTimes(2) }) @@ -946,7 +1166,11 @@ describe('InterruptManager transactions', () => { } firstItem.resolveInterrupt('first answer') await settle() - manager.getInterrupts()[0]?.clearResolution() + const clearedItem = manager.getInterrupts()[0] + if (clearedItem?.kind !== 'generic') { + throw new Error('Expected generic interrupt') + } + clearedItem.clearResolution() const secondItem = manager.getInterrupts()[0] if (secondItem?.kind !== 'generic') { throw new Error('Expected generic interrupt') @@ -1181,6 +1405,313 @@ describe('ChatClient native interrupts', () => { expect(sentMessages[1]).toEqual(sentMessages[0]) }) + it('sends first-party continuation on resume metadata', async () => { + const approval = defineInterrupt({ + id: 'approval', + responseSchema: z.object({ answer: z.number() }), + }) + const contexts: Array = [] + const binding: InterruptBinding = { + v: INTERRUPT_BINDING_VERSION, + kind: 'generic', + interruptId: 'generic_review_one', + interruptedRunId: 'placeholder', + generation: 0, + definitionId: 'approval', + key: 'one', + batchIndex: 0, + responseSchemaHash: digestInterruptJson( + canonicalInterruptJson( + convertSchemaToJsonSchema(approval.responseSchema), + ), + ), + } + let call = 0 + const connection: ConnectConnectionAdapter = { + async *connect(_messages, _data, _signal, context) { + contexts.push(context) + 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 + binding.interruptId = `generic_${runId}_approval_one` + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + outcome: { + type: 'interrupt', + interrupts: [ + descriptor(binding, { + responseSchema: convertSchemaToJsonSchema( + approval.responseSchema, + ), + }), + ], + }, + } + return + } + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + outcome: { type: 'success' }, + } + }, + } + const client = new ChatClient({ + connection, + threadId: 'thread-1', + interrupts: [approval], + }) + + await client.sendMessage('start') + const item = client.getInterrupts()[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: binding.interruptId, + status: 'resolved', + payload: { answer: 42 }, + metadata: { + 'tanstack:interruptContinuation': { + v: 1, + definitionId: 'approval', + key: 'one', + batchIndex: 0, + reason: 'confirmation', + message: '', + responseSchemaHash: binding.responseSchemaHash, + }, + }, + }, + ], + }) + }) + + it('does not auto-send after a completed tool when a generic interrupt is pending', async () => { + const review = defineInterrupt({ + id: 'review-plan', + responseSchema: z.object({ approved: z.boolean() }), + }) + const contexts: Array = [] + let call = 0 + const connection: ConnectConnectionAdapter = { + async *connect( + _messages, + _data, + _signal, + context, + ): AsyncGenerator { + contexts.push(context) + 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) { + const interruptId = `generic_${runId}_review` + yield { + type: EventType.TOOL_CALL_START, + toolCallId: 'call-inspect', + toolCallName: 'inspectPlan', + toolName: 'inspectPlan', + timestamp: Date.now(), + } + yield { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call-inspect', + delta: '{"planId":"PLAN-42"}', + timestamp: Date.now(), + } + yield { + type: EventType.TOOL_CALL_END, + toolCallId: 'call-inspect', + timestamp: Date.now(), + } + yield { + type: EventType.TOOL_CALL_RESULT, + toolCallId: 'call-inspect', + messageId: 'tool-result-inspect', + content: '{"inspected":true,"planId":"PLAN-42"}', + timestamp: Date.now(), + } + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + outcome: { + type: 'interrupt', + interrupts: [ + descriptor( + { + v: INTERRUPT_BINDING_VERSION, + kind: 'generic', + interruptId, + interruptedRunId: runId, + generation: 0, + definitionId: 'review-plan', + key: 'afterTools-review', + batchIndex: 0, + responseSchemaHash: digestInterruptJson( + canonicalInterruptJson( + convertSchemaToJsonSchema(review.responseSchema), + ), + ), + }, + { reason: 'review_required' }, + ), + ], + }, + } + return + } + throw new Error(`unexpected extra connect call ${call}`) + }, + } + const client = new ChatClient({ + connection, + threadId: 'thread-1', + interrupts: [review], + }) + + await client.sendMessage('inspect') + await vi.waitFor(() => expect(client.getInterrupts()).toHaveLength(1)) + await Promise.resolve() + await Promise.resolve() + expect(call).toBe(1) + expect(contexts[0]?.resume).toBeUndefined() + expect(contexts[0]?.parentRunId).toBeUndefined() + }) + + it('does not auto-send after rejecting an afterTools generic interrupt', async () => { + const review = defineInterrupt({ + id: 'review-plan', + responseSchema: z.object({ approved: z.boolean() }), + }) + let call = 0 + const connection: ConnectConnectionAdapter = { + async *connect( + _messages, + _data, + _signal, + context, + ): AsyncGenerator { + 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) { + const interruptId = `generic_${runId}_review` + yield { + type: EventType.TOOL_CALL_START, + toolCallId: 'call-inspect', + toolCallName: 'inspectPlan', + toolName: 'inspectPlan', + timestamp: Date.now(), + } + yield { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call-inspect', + delta: '{"planId":"PLAN-42"}', + timestamp: Date.now(), + } + yield { + type: EventType.TOOL_CALL_END, + toolCallId: 'call-inspect', + timestamp: Date.now(), + } + yield { + type: EventType.TOOL_CALL_RESULT, + toolCallId: 'call-inspect', + messageId: 'tool-result-inspect', + content: '{"inspected":true,"planId":"PLAN-42"}', + timestamp: Date.now(), + } + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + outcome: { + type: 'interrupt', + interrupts: [ + descriptor( + { + v: INTERRUPT_BINDING_VERSION, + kind: 'generic', + interruptId, + interruptedRunId: runId, + generation: 0, + definitionId: 'review-plan', + key: 'afterTools-review', + batchIndex: 0, + responseSchemaHash: digestInterruptJson( + canonicalInterruptJson( + convertSchemaToJsonSchema(review.responseSchema), + ), + ), + }, + { reason: 'review_required' }, + ), + ], + }, + } + return + } + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + finishReason: 'stop', + outcome: { type: 'success' }, + } + }, + } + const client = new ChatClient({ + connection, + threadId: 'thread-1', + interrupts: [review], + }) + + await client.sendMessage('inspect') + const item = client.getInterrupts()[0] + if (item?.kind !== 'generic') throw new Error('Expected generic interrupt') + item.resolveInterrupt({ approved: false }) + await vi.waitFor(() => expect(call).toBe(2)) + await Promise.resolve() + await Promise.resolve() + expect(call).toBe(2) + expect(client.getInterrupts()).toHaveLength(0) + }) + it('resumes a hydrated ephemeral batch with full history in a fresh child run', async () => { const contexts: Array = [] const sentMessages: Array | Array> = [] diff --git a/packages/ai-client/tests/connection-adapters.test.ts b/packages/ai-client/tests/connection-adapters.test.ts index 4f1b94c1f2..9a4e00e537 100644 --- a/packages/ai-client/tests/connection-adapters.test.ts +++ b/packages/ai-client/tests/connection-adapters.test.ts @@ -3,6 +3,7 @@ import { EventType } from '@tanstack/ai/client' import { fetchHttpStream, fetchServerSentEvents, + fetcherToConnectionAdapter, normalizeConnectionAdapter, rpcStream, stream, @@ -26,7 +27,110 @@ describe('connection-adapters', () => { vi.clearAllMocks() }) + it('forwards resume on a fetcher adapter', async () => { + const fetcher = vi.fn(async function* () {}) + const adapter = fetcherToConnectionAdapter(fetcher) + const signal = new AbortController().signal + + for await (const _chunk of adapter.connect([], { source: 'test' }, signal, { + threadId: 'thread-1', + runId: 'resume-run', + parentRunId: 'interrupted-run', + resume: [ + { + interruptId: 'generic-1', + status: 'cancelled', + metadata: { + 'tanstack:interruptContinuation': { + v: 1, + definitionId: 'review', + key: 'one', + batchIndex: 0, + reason: 'review', + message: 'Review', + }, + }, + }, + ], + })) { + // Consume the terminal event. + } + + expect(fetcher).toHaveBeenCalledWith( + expect.objectContaining({ + threadId: 'thread-1', + runId: 'resume-run', + parentRunId: 'interrupted-run', + resume: [ + { + interruptId: 'generic-1', + status: 'cancelled', + metadata: { + 'tanstack:interruptContinuation': { + v: 1, + definitionId: 'review', + key: 'one', + batchIndex: 0, + reason: 'review', + message: 'Review', + }, + }, + }, + ], + }), + { signal }, + ) + }) + describe('fetchServerSentEvents', () => { + it('sends generic continuation on resume metadata, not state', async () => { + const mockResponse = { + ok: true, + body: { + getReader: () => ({ + read: vi.fn().mockResolvedValue({ done: true, value: undefined }), + releaseLock: vi.fn(), + }), + }, + } + fetchMock.mockResolvedValue(mockResponse as any) + const adapter = fetchServerSentEvents('/api/chat') + const resume = [ + { + interruptId: 'generic-1', + status: 'cancelled' as const, + metadata: { + 'tanstack:interruptContinuation': { + v: 1, + definitionId: 'review', + key: 'one', + batchIndex: 0, + reason: 'review', + message: 'Review', + }, + }, + }, + ] + + for await (const _chunk of adapter.connect( + [{ role: 'user', content: 'Resume' }], + undefined, + undefined, + { + threadId: 'thread-1', + runId: 'run-2', + resume, + }, + )) { + // Consume the empty stream. + } + + const request = fetchMock.mock.calls[0]?.[1] as RequestInit + const body = JSON.parse(String(request.body)) + expect(body.resume).toEqual(resume) + expect(body.state).toEqual({}) + }) + it('should handle SSE format with data: prefix', async () => { const mockReader = { read: vi diff --git a/packages/ai-client/tests/dispose-tail-leak.test.ts b/packages/ai-client/tests/dispose-tail-leak.test.ts index 0d1a91cbf0..062eff5ff8 100644 --- a/packages/ai-client/tests/dispose-tail-leak.test.ts +++ b/packages/ai-client/tests/dispose-tail-leak.test.ts @@ -32,7 +32,7 @@ function seededStore(): { */ function mountedChatClient( options: ConstructorParameters[0], -): ChatClient { +) { const client = new ChatClient(options) client.attach() return client diff --git a/packages/ai-client/tests/interrupts-types.test-d.ts b/packages/ai-client/tests/interrupts-types.test-d.ts index e025102521..edf3d9e877 100644 --- a/packages/ai-client/tests/interrupts-types.test-d.ts +++ b/packages/ai-client/tests/interrupts-types.test-d.ts @@ -1,5 +1,5 @@ import { expectTypeOf } from 'vitest' -import { toolDefinition } from '@tanstack/ai/client' +import { defineInterrupt, toolDefinition } from '@tanstack/ai/client' import { z } from 'zod' import type { JSONSchema, ServerTool } from '@tanstack/ai' import type { @@ -14,6 +14,9 @@ import type { BoundInterrupts, ChatClient, ChatInterrupt, + GenericAGUIInterrupt, + GenericInterrupt, + RegisteredGenericInterrupt, ToolApprovalInterrupt, } from '../src/index' @@ -55,6 +58,17 @@ type Tools = readonly [ ] type Interrupt = ChatInterrupt +declare const externalGeneric: GenericAGUIInterrupt +externalGeneric.cancel() +externalGeneric.clearResolution() +externalGeneric.resolveInterrupt({ accepted: true }) + +declare const unboundInterrupt: Extract +// @ts-expect-error Unbound interrupts are owned by another system. +unboundInterrupt.cancel() +// @ts-expect-error Unbound interrupts are owned by another system. +unboundInterrupt.clearResolution() + expectTypeOf>().toEqualTypeOf() expectTypeOf>().toEqualTypeOf() expectTypeOf>().toEqualTypeOf() @@ -162,3 +176,67 @@ client.unsafeResumeInterrupts([]) expectTypeOf>().toMatchTypeOf< ToolApprovalInterrupt >() + +const reviewPlan = defineInterrupt({ + id: 'review-plan', + payloadSchema: z.object({ title: z.string() }), + responseSchema: z.string().transform((value) => Number(value)), +}) + +const acknowledge = defineInterrupt({ + id: 'acknowledge', + responseSchema: z.object({ accepted: z.boolean() }), +}) + +const payloadOnlyReview = defineInterrupt({ + id: 'payload-only-review', + payloadSchema: z.object({ + title: z.string().transform((value) => value.toUpperCase()), + }), +}) + +type RegisteredInterrupts = readonly [ + typeof reviewPlan, + typeof acknowledge, + typeof payloadOnlyReview, +] +type RegisteredClientInterrupt = ChatInterrupt + +declare const reviewInterrupt: Extract< + RegisteredClientInterrupt, + { definitionId: 'review-plan' } +> +expectTypeOf(reviewInterrupt).toMatchTypeOf< + RegisteredGenericInterrupt +>() +expectTypeOf(reviewInterrupt).toEqualTypeOf< + GenericInterrupt +>() +expectTypeOf(reviewInterrupt.payload).toEqualTypeOf< + { title: string } | undefined +>() +reviewInterrupt.resolveInterrupt('42') +// @ts-expect-error response schema input is a string, not its transformed output +reviewInterrupt.resolveInterrupt(42) + +declare const acknowledgement: Extract< + RegisteredClientInterrupt, + { definitionId: 'acknowledge' } +> +expectTypeOf(acknowledgement.payload).toEqualTypeOf() +acknowledgement.resolveInterrupt({ accepted: true }) +// @ts-expect-error response input must match the registered schema +acknowledgement.resolveInterrupt({ accepted: 'yes' }) + +declare const payloadOnlyInterrupt: Extract< + RegisteredClientInterrupt, + { definitionId: 'payload-only-review' } +> +expectTypeOf(payloadOnlyReview.responseSchema).toEqualTypeOf() +expectTypeOf(payloadOnlyInterrupt.payload).toEqualTypeOf< + { title: string } | undefined +>() +expectTypeOf(payloadOnlyInterrupt.resolveInterrupt) + .parameter(0) + .toEqualTypeOf() +payloadOnlyInterrupt.resolveInterrupt({ accepted: true, comment: 'continue' }) diff --git a/packages/ai-client/tests/resume-snapshot.test.ts b/packages/ai-client/tests/resume-snapshot.test.ts index 4bc0ecb5ad..462c898e9a 100644 --- a/packages/ai-client/tests/resume-snapshot.test.ts +++ b/packages/ai-client/tests/resume-snapshot.test.ts @@ -48,7 +48,7 @@ function memoryAdapter(initial?: ChatPersistedState | Array): { */ function mountedChatClient( options: ConstructorParameters[0], -): ChatClient { +) { const client = new ChatClient(options) client.attach() return client diff --git a/packages/ai-persistence/src/index.ts b/packages/ai-persistence/src/index.ts index 9519ec6235..dece721265 100644 --- a/packages/ai-persistence/src/index.ts +++ b/packages/ai-persistence/src/index.ts @@ -19,6 +19,7 @@ export type { TerminalRunStatus, RunRecord, RunStore, + InterruptCommitEntry, InterruptRecord, InterruptStatus, InterruptStore, diff --git a/packages/ai-persistence/src/memory.ts b/packages/ai-persistence/src/memory.ts index 922f5344a7..ef28cb1beb 100644 --- a/packages/ai-persistence/src/memory.ts +++ b/packages/ai-persistence/src/memory.ts @@ -14,6 +14,7 @@ import type { BlobStore, GenerationRunRecord, GenerationRunStore, + InterruptCommitEntry, InterruptRecord, InterruptStore, MessageStore, @@ -188,6 +189,49 @@ class MemoryInterruptStore implements InterruptStore { } return Promise.resolve() } + async commitBatch( + entries: ReadonlyArray, + ): Promise { + const ids = new Set() + for (const entry of entries) { + if (ids.has(entry.interruptId)) { + throw new Error( + `Interrupt batch contains duplicate id: ${entry.interruptId}.`, + ) + } + ids.add(entry.interruptId) + const existing = this.interrupts.get(entry.interruptId) + if (!existing) { + throw new Error( + `Interrupt batch references missing id: ${entry.interruptId}.`, + ) + } + if (existing.status !== 'pending') { + throw new Error( + `Interrupt batch references non-pending id: ${entry.interruptId}.`, + ) + } + } + const resolvedAt = Date.now() + for (const entry of entries) { + const existing = this.interrupts.get(entry.interruptId) + if (!existing) continue + if (entry.status === 'resolved') { + this.interrupts.set(entry.interruptId, { + ...existing, + status: 'resolved', + resolvedAt, + response: entry.response, + }) + } else { + this.interrupts.set(entry.interruptId, { + ...existing, + status: 'cancelled', + resolvedAt, + }) + } + } + } get(interruptId: string): Promise { return Promise.resolve(this.interrupts.get(interruptId) ?? null) } diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index dbeeea2928..0a1fa29f65 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -1,9 +1,21 @@ import { defineChatMiddleware, getDetachableRun, + InterruptResumeValidationError, + readInterruptBinding, + validateInterruptResumeBatch, wasCancelRequested, } from '@tanstack/ai' -import { providePendingTurn } from '@tanstack/ai/adapter-internals' +import { + createInterruptBinding, + getGenericInterruptDefinitionRegistry, + providePendingTurn, + rehydrateInterruptRequest, +} from '@tanstack/ai/adapter-internals' +import type { + GenericInterruptRequest, + InterruptDefinition, +} from '@tanstack/ai/adapter-internals' import { base64ToUint8Array } from '@tanstack/ai-utils' import { InterruptsCapability, @@ -28,12 +40,15 @@ import type { GenerationFinishInfo, GenerationMiddleware, GenerationMiddlewareContext, + Interrupt, ModelMessage, + PendingInterruptResumeRecord, PersistedArtifactActivity, PersistedArtifactRef, PersistedArtifactRole, RunAgentResumeItem, StreamChunk, + Tool, ToolApprovalResolution, TokenUsage, } from '@tanstack/ai' @@ -43,6 +58,7 @@ import type { ArtifactRecord, BlobBody, ChatTranscriptStores, + InterruptCommitEntry, InterruptRecord, RunStore, } from './types' @@ -272,27 +288,142 @@ const runState = new WeakMap() const validResumeStatuses = new Set(['resolved', 'cancelled']) +function mergeMaps( + left?: ReadonlyMap, + right?: ReadonlyMap, +): Map | undefined { + if (!left && !right) return undefined + return new Map([...(left ?? []), ...(right ?? [])]) +} + +function mergeSets( + left?: ReadonlySet, + right?: ReadonlySet, +): Set | undefined { + if (!left && !right) return undefined + return new Set([...(left ?? []), ...(right ?? [])]) +} + +function mergeResumeToolState( + left: ChatResumeToolState | undefined, + right: ChatResumeToolState | undefined, +): ChatResumeToolState | undefined { + if (!left) return right + if (!right) return left + return { + approvals: mergeMaps(left.approvals, right.approvals), + clientToolResults: mergeMaps( + left.clientToolResults, + right.clientToolResults, + ), + genericInterrupts: mergeMaps( + left.genericInterrupts, + right.genericInterrupts, + ), + genericInterruptRequests: mergeMaps( + left.genericInterruptRequests, + right.genericInterruptRequests, + ), + deniedToolResults: mergeMaps( + left.deniedToolResults, + right.deniedToolResults, + ), + cancelledToolCallIds: mergeSets( + left.cancelledToolCallIds, + right.cancelledToolCallIds, + ), + } +} + +function rejectMixedRunPending( + pending: Array, + ctx: Pick, +): void { + const runIds = new Set(pending.map((interrupt) => interrupt.runId)) + if (runIds.size <= 1) return + throw new InterruptResumeValidationError([ + { + scope: 'batch', + threadId: ctx.threadId, + interruptedRunId: ctx.runId, + generation: 0, + interruptIds: pending.map((interrupt) => interrupt.interruptId), + code: 'stale', + message: 'Thread has pending interrupts from more than one run.', + source: 'server', + retryable: false, + }, + ]) +} + function validatePendingResumes( pending: Array, resume: Array | undefined, + ctx: Pick, ): Map { + const interruptedRunId = pending[0]?.runId ?? ctx.runId + const failure = ( + interruptId: string, + code: 'conflict' | 'unknown-interrupt', + message: string, + ): never => { + throw new InterruptResumeValidationError([ + { + scope: 'item', + threadId: ctx.threadId, + interruptedRunId, + generation: 0, + interruptId, + code, + message, + source: 'client', + retryable: false, + }, + { + scope: 'batch', + threadId: ctx.threadId, + interruptedRunId, + generation: 0, + interruptIds: pending.map((interrupt) => interrupt.interruptId), + code: code === 'conflict' ? 'conflict' : 'incomplete-batch', + message: + 'Resume entries must resolve or cancel the complete interrupt batch.', + source: 'client', + retryable: false, + }, + ]) + } const pendingInterruptIds = new Set( pending.map((interrupt) => interrupt.interruptId), ) - const resumeByInterruptId = new Map( - (resume ?? []).map((entry) => [entry.interruptId, entry]), - ) + const resumeByInterruptId = new Map() + for (const entry of resume ?? []) { + if (resumeByInterruptId.has(entry.interruptId)) { + return failure( + entry.interruptId, + 'conflict', + `Interrupt ${entry.interruptId} has duplicate resume entries.`, + ) + } + resumeByInterruptId.set(entry.interruptId, entry) + } if (pending.length === 0) { const staleEntry = resume?.[0] if (staleEntry) { - throw new Error( + return failure( + staleEntry.interruptId, + 'unknown-interrupt', `Resume entry references non-pending interrupt ${staleEntry.interruptId}.`, ) } return resumeByInterruptId } + const firstPending = pending[0] + if (firstPending === undefined) return resumeByInterruptId if (!resume || resume.length === 0) { - throw new Error( + return failure( + firstPending.interruptId, + 'unknown-interrupt', `Thread has pending interrupts; resume is required before accepting new input.`, ) } @@ -300,19 +431,25 @@ function validatePendingResumes( for (const interrupt of pending) { const entry = resumeByInterruptId.get(interrupt.interruptId) if (!entry) { - throw new Error( + return failure( + interrupt.interruptId, + 'unknown-interrupt', `Missing resume entry for pending interrupt ${interrupt.interruptId}.`, ) } if (!validResumeStatuses.has(entry.status)) { - throw new Error( + return failure( + interrupt.interruptId, + 'unknown-interrupt', `Invalid resume status for pending interrupt ${interrupt.interruptId}: ${entry.status}.`, ) } } for (const entry of resume) { if (!pendingInterruptIds.has(entry.interruptId)) { - throw new Error( + return failure( + entry.interruptId, + 'unknown-interrupt', `Resume entry references non-pending interrupt ${entry.interruptId}.`, ) } @@ -325,6 +462,27 @@ async function applyPendingResumes( resumeByInterruptId: Map, interrupts: NonNullable, ): Promise { + const entries: Array = [] + for (const interrupt of pending) { + const entry = resumeByInterruptId.get(interrupt.interruptId) + if (!entry) continue + if (entry.status === 'resolved') { + entries.push({ + interruptId: interrupt.interruptId, + status: 'resolved', + response: entry.payload, + }) + } else { + entries.push({ + interruptId: interrupt.interruptId, + status: 'cancelled', + }) + } + } + if (interrupts.commitBatch) { + await interrupts.commitBatch(entries) + return + } for (const interrupt of pending) { const entry = resumeByInterruptId.get(interrupt.interruptId) if (!entry) continue @@ -375,6 +533,267 @@ function interruptKind(interrupt: InterruptRecord): string | undefined { return metadata ? stringField(metadata, 'kind') : undefined } +function hasReservedInterruptBinding(payload: unknown): boolean { + const descriptor = objectValue(payload) + const metadata = objectValue(descriptor?.metadata) + return !!metadata && 'tanstack:interruptBinding' in metadata +} + +function isPersistedInterruptDescriptor( + value: unknown, +): value is Interrupt & { reason: string; message: string } { + const record = objectValue(value) + return ( + !!record && + typeof record.id === 'string' && + typeof record.reason === 'string' && + typeof record.message === 'string' + ) +} + +/** + * Does this pending record belong to the TanStack chat resume protocol? + * + * An external system can persist an AG-UI descriptor in the same durable + * thread. A descriptor without a TanStack binding or legacy tool marker stays + * pending for its owner, but it does not make this resume incomplete. Older + * opaque records remain owned because their provenance cannot be known. + */ +function isChatOwnedPendingInterrupt(interrupt: InterruptRecord): boolean { + const kind = interruptKind(interrupt) + return ( + !isPersistedInterruptDescriptor(interrupt.payload) || + stringField(interrupt.payload, 'toolCallId') !== undefined || + kind === 'approval' || + kind === 'client_tool' || + hasReservedInterruptBinding(interrupt.payload) + ) +} + +function durableGenericFailure( + ctx: Pick, + persisted: InterruptRecord, + message: string, +): InterruptResumeValidationError { + return new InterruptResumeValidationError([ + { + scope: 'item', + threadId: ctx.threadId, + interruptedRunId: persisted.runId || ctx.runId, + generation: 0, + interruptId: persisted.interruptId, + code: 'stale', + message, + source: 'server', + retryable: false, + }, + { + scope: 'batch', + threadId: ctx.threadId, + interruptedRunId: persisted.runId || ctx.runId, + generation: 0, + interruptIds: [persisted.interruptId], + code: 'item-validation-failed', + message: 'One or more persisted interrupt records are invalid.', + source: 'server', + retryable: false, + }, + ]) +} + +async function durableGenericResumeState( + ctx: ChatMiddlewareContext, + pending: Array, + resume: ReadonlyArray, + tools: Array, +): Promise { + const registry = getGenericInterruptDefinitionRegistry(ctx, { + optional: true, + }) + const records: Array = [] + + for (const persisted of pending) { + if (!isPersistedInterruptDescriptor(persisted.payload)) { + if (hasReservedInterruptBinding(persisted.payload)) { + throw durableGenericFailure( + ctx, + persisted, + `Persisted interrupt ${persisted.interruptId} has an invalid binding descriptor.`, + ) + } + continue + } + const descriptor = persisted.payload + const binding = readInterruptBinding(descriptor) + if (!binding) { + if (hasReservedInterruptBinding(descriptor)) { + throw durableGenericFailure( + ctx, + persisted, + `Persisted interrupt ${persisted.interruptId} has an invalid or incomplete binding.`, + ) + } + continue + } + if ( + descriptor.id !== persisted.interruptId || + binding.interruptId !== persisted.interruptId || + binding.interruptedRunId !== persisted.runId || + binding.generation !== 0 + ) { + throw durableGenericFailure( + ctx, + persisted, + `Persisted interrupt ${persisted.interruptId} has stale correlation metadata.`, + ) + } + if (binding.kind !== 'generic') { + records.push({ + interruptId: persisted.interruptId, + payload: descriptor, + binding, + }) + continue + } + if ( + !binding.definitionId || + !binding.key || + binding.batchIndex === undefined + ) { + records.push({ + interruptId: persisted.interruptId, + payload: descriptor, + binding, + }) + continue + } + if (!registry) { + throw durableGenericFailure( + ctx, + persisted, + `Persisted generic interrupt ${persisted.interruptId} cannot be restored because no interrupt registry is available.`, + ) + } + const definition = registry.definitions.get(binding.definitionId) + if (!definition) { + throw durableGenericFailure( + ctx, + persisted, + `Persisted generic interrupt definition ${binding.definitionId} is unavailable.`, + ) + } + const metadata = objectValue(descriptor.metadata) + const payload = metadata?.['tanstack:interruptPayload'] + let request: GenericInterruptRequest< + InterruptDefinition + > + try { + request = rehydrateInterruptRequest(definition, { + key: binding.key, + reason: descriptor.reason, + message: descriptor.message, + ...(descriptor.expiresAt !== undefined + ? { expiresAt: descriptor.expiresAt } + : {}), + ...(payload !== undefined ? { payload } : {}), + }) + } catch (error) { + throw durableGenericFailure( + ctx, + persisted, + `Persisted generic interrupt ${persisted.interruptId} is invalid: ${error instanceof Error ? error.message : String(error)}`, + ) + } + const emitted = createInterruptBinding(request, { + batchIndex: binding.batchIndex, + }) + if ( + emitted.descriptor.responseSchemaHash !== binding.responseSchemaHash || + emitted.descriptor.payloadSchemaHash !== binding.payloadSchemaHash || + binding.interruptId !== persisted.interruptId + ) { + throw durableGenericFailure( + ctx, + persisted, + `Persisted generic interrupt ${persisted.interruptId} is stale.`, + ) + } + records.push({ + interruptId: persisted.interruptId, + payload: descriptor, + binding, + genericRequest: request, + }) + } + + const firstRecord = records[0] + if (firstRecord === undefined) return undefined + const interruptedRunId = firstRecord.binding.interruptedRunId + const generation = firstRecord.binding.generation + const validated = await validateInterruptResumeBatch({ + threadId: ctx.threadId, + interruptedRunId, + generation, + pending: records, + resume: resume.filter((entry) => + records.some((record) => record.interruptId === entry.interruptId), + ), + tools, + }) + if (validated.errors.length > 0 || !validated.resumeToolState) { + throw new InterruptResumeValidationError(validated.errors) + } + type GenericRecord = PendingInterruptResumeRecord & { + binding: Extract< + PendingInterruptResumeRecord['binding'], + { kind: 'generic' } + > + genericRequest: GenericInterruptRequest< + InterruptDefinition + > + } + const isGenericRecord = ( + record: PendingInterruptResumeRecord, + ): record is GenericRecord => + record.binding.kind === 'generic' && record.genericRequest !== undefined + const genericRecords: Array<{ record: GenericRecord; batchIndex: number }> = + [] + const batchIndexes = new Set() + for (const record of records) { + if (!isGenericRecord(record)) continue + const batchIndex = record.binding.batchIndex + if (batchIndex === undefined || batchIndexes.has(batchIndex)) { + throw new InterruptResumeValidationError([ + { + scope: 'batch', + threadId: ctx.threadId, + interruptedRunId, + generation, + interruptIds: records.map((item) => item.interruptId), + code: 'stale', + message: + 'Persisted generic interrupts have duplicate or invalid batch indexes.', + source: 'server', + retryable: false, + }, + ]) + } + batchIndexes.add(batchIndex) + genericRecords.push({ record, batchIndex }) + } + genericRecords.sort((left, right) => left.batchIndex - right.batchIndex) + return { + ...validated.resumeToolState, + genericInterruptRequests: new Map( + genericRecords.flatMap(({ record }) => + record.genericRequest + ? [[record.interruptId, record.genericRequest] as const] + : [], + ), + ), + } +} + function resolvedApprovalDecision(entry: RunAgentResumeItem): boolean { if (entry.status === 'cancelled') return false const payload = objectValue(entry.payload) @@ -1480,11 +1899,15 @@ export function withPersistence( const pending = await persistence.stores.interrupts.listPending( ctx.threadId, ) - // Gate: a thread with pending interrupts must carry a resume batch that - // references them. + // Gate only records that this chat owns. A foreign AG-UI interrupt can + // share the durable thread, but its owner resolves it outside this + // resume protocol. Including it would deadlock this chat resume. + const ownedPending = pending.filter(isChatOwnedPendingInterrupt) + rejectMixedRunPending(ownedPending, ctx) const resumeByInterruptId = validatePendingResumes( - pending, + ownedPending, config.resume, + ctx, ) // Persistence is the server-authoritative resume path: translate the // persisted interrupts into the engine's resume tool state and CLEAR @@ -1493,18 +1916,29 @@ export function withPersistence( // persistence flow deliberately omits). if ((config.resume?.length ?? 0) > 0) { const resumeToolState = resumeToolStateFromPending( - pending, + ownedPending, resumeByInterruptId, ) + const genericResumeState = await durableGenericResumeState( + ctx, + ownedPending, + config.resume ?? [], + config.tools, + ) patch.resume = [] - if (resumeToolState) patch.resumeToolState = resumeToolState + if (resumeToolState || genericResumeState) { + patch.resumeToolState = mergeResumeToolState( + resumeToolState, + genericResumeState, + ) + } } // Defer marking these interrupts resolved/cancelled until the run // succeeds (see commitPendingResumes). Committing here would consume the // approval even if the run then failed, breaking a retry. const state = runState.get(ctx) - if (state && pending.length > 0) { - state.pendingResumes = { pending, resumeByInterruptId } + if (state && ownedPending.length > 0) { + state.pendingResumes = { pending: ownedPending, resumeByInterruptId } } } @@ -1634,17 +2068,25 @@ export function withPersistence( // resumes stay pending so a retry can re-apply them. Completing the run // or consuming approvals before the durable history lands leaves a // "finished" run whose transcript is missing the terminal turn. - await messageStore.saveThread( - ctx.threadId, - finishedTranscript( - ctx.messages, - info, - state?.streamingMessageId, - state?.streamingMessageCreatedAt, - ), - ) - await completeRun(runs, ctx.runId, info.usage) - await commitPendingResumes(state, persistence.stores.interrupts) + try { + await messageStore.saveThread( + ctx.threadId, + finishedTranscript( + ctx.messages, + info, + state?.streamingMessageId, + state?.streamingMessageCreatedAt, + ), + ) + await commitPendingResumes(state, persistence.stores.interrupts) + await completeRun(runs, ctx.runId, info.usage) + } catch (error) { + // Core has already selected its terminal hook. Persist the failed run + // here, so a failed transcript save or batch write does not leave an + // interrupted or completed run whose pending records need retrying. + await failRun(runs, ctx.runId, error) + throw error + } }, async onError(ctx: ChatMiddlewareContext, info: ErrorInfo) { diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index 7362141b57..05ee5a9ad9 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -221,6 +221,18 @@ export interface InterruptRecord { response?: unknown } +/** A terminal interrupt write for {@link InterruptStore.commitBatch}. */ +export type InterruptCommitEntry = + | { + interruptId: string + status: 'resolved' + response?: unknown + } + | { + interruptId: string + status: 'cancelled' + } + /** Durable store for human-in-the-loop interrupts. */ export interface InterruptStore { /** @@ -249,6 +261,19 @@ export interface InterruptStore { * `interruptId` does not exist. */ cancel: (interruptId: string) => Promise + /** + * Commit terminal writes for a validated resume batch. + * + * Optional. When present, `withPersistence` calls it once instead of + * calling `resolve` and `cancel` for each entry. Apply every entry or none. + * + * Reject the whole batch (throw, writing nothing) when any entry has a + * duplicate `interruptId`, references an `interruptId` that does not exist, + * or references an interrupt whose status is not `'pending'`. This is + * stricter than `resolve` / `cancel`, which are no-ops for a missing + * `interruptId`. + */ + commitBatch?: (entries: ReadonlyArray) => Promise /** Return the interrupt for `interruptId`, or `null` if none exists. */ get: (interruptId: string) => Promise /** diff --git a/packages/ai-persistence/tests/interrupts.test.ts b/packages/ai-persistence/tests/interrupts.test.ts index ac46b31d03..92f0668796 100644 --- a/packages/ai-persistence/tests/interrupts.test.ts +++ b/packages/ai-persistence/tests/interrupts.test.ts @@ -1,8 +1,19 @@ import { describe, expect, it, vi } from 'vitest' -import { EventType, chat, defineChatMiddleware } from '@tanstack/ai' -import type { AnyTextAdapter, StreamChunk, Tool } from '@tanstack/ai' +import { + EventType, + chat, + defineChatMiddleware, + defineInterrupt, +} from '@tanstack/ai' +import type { + AnyTextAdapter, + ChatResumeToolState, + StreamChunk, + Tool, +} from '@tanstack/ai' import { memoryPersistence } from '../src/memory' import { withPersistence } from '../src/middleware' +import type { InterruptStore } from '../src/types' function mockAdapter(iterations: Array>) { const calls: Array = [] @@ -31,6 +42,76 @@ async function collect(stream: AsyncIterable) { return out } +const coercedCountSchema = { + '~standard': { + version: 1, + vendor: 'test', + validate(value: unknown) { + if ( + !value || + typeof value !== 'object' || + Array.isArray(value) || + !('count' in value) || + (typeof value.count !== 'string' && typeof value.count !== 'number') + ) { + return { issues: [{ message: 'count is required' }] } + } + const count = Number(value.count) + return Number.isFinite(count) + ? { value: { count } } + : { issues: [{ message: 'count must be numeric' }] } + }, + jsonSchema: { + input() { + return { + type: 'object', + required: ['count'], + properties: { count: { type: 'string' } }, + } + }, + }, + }, +} as const + +const transformedDisplaySchema = { + '~standard': { + version: 1, + vendor: 'test', + validate(value: unknown) { + return typeof value === 'string' + ? { value: value.length } + : { issues: [{ message: 'display payload must be a string' }] } + }, + jsonSchema: { + input() { + return { type: 'string' } + }, + }, + }, +} as const + +function expectResumeError( + chunks: ReadonlyArray, + interruptId: string, +) { + const error = chunks.find((chunk) => chunk.type === EventType.RUN_ERROR) + expect(error?.['tanstack:interruptErrors']).toEqual( + expect.arrayContaining([ + expect.objectContaining({ scope: 'item', interruptId }), + ]), + ) +} + +function isInterruptTerminal(chunk: StreamChunk): chunk is StreamChunk & { + outcome: { type: 'interrupt'; interrupts: ReadonlyArray<{ id: string }> } +} { + return ( + chunk.type === EventType.RUN_FINISHED && + 'outcome' in chunk && + chunk.outcome?.type === 'interrupt' + ) +} + const interruptFinished = (runId = 'r1'): StreamChunk => ({ type: EventType.RUN_FINISHED, runId, @@ -238,8 +319,8 @@ describe('interrupt persistence', () => { }) const { adapter } = mockAdapter([[interruptFinished()]]) - await expect( - collect( + expectResumeError( + await collect( chat({ adapter, messages: [{ role: 'user', content: 'new input' }], @@ -248,7 +329,8 @@ describe('interrupt persistence', () => { middleware: [withPersistence(persistence)], }) as AsyncIterable, ), - ).rejects.toThrow(/pending interrupt/i) + 'interrupt-1', + ) expect(await persistence.stores.runs!.get('r2')).toBeNull() }) @@ -477,8 +559,8 @@ describe('interrupt persistence', () => { ) const continuation = mockAdapter([[text('SHOULD NOT RUN')]]) - await expect( - collect( + expectResumeError( + await collect( chat({ adapter: continuation.adapter, messages: [], @@ -488,10 +570,11 @@ describe('interrupt persistence', () => { middleware: [withPersistence(persistence)], }) as AsyncIterable, ), - ).rejects.toThrow(/pending interrupts.*resume is required/i) + 'interrupt-1', + ) - await expect( - collect( + expectResumeError( + await collect( chat({ adapter: continuation.adapter, messages: [], @@ -501,7 +584,8 @@ describe('interrupt persistence', () => { middleware: [withPersistence(persistence)], }) as AsyncIterable, ), - ).rejects.toThrow(/missing resume entry.*interrupt-1/i) + 'interrupt-1', + ) expect(continuation.calls).toHaveLength(0) expect(await persistence.stores.interrupts!.listPending('t1')).toHaveLength( @@ -524,8 +608,8 @@ describe('interrupt persistence', () => { expect(await persistence.stores.interrupts!.listPending('t1')).toEqual([]) const continuation = mockAdapter([[text('SHOULD NOT RUN')]]) - await expect( - collect( + expectResumeError( + await collect( chat({ adapter: continuation.adapter, messages: [], @@ -535,7 +619,8 @@ describe('interrupt persistence', () => { middleware: [withPersistence(persistence)], }) as AsyncIterable, ), - ).rejects.toThrow(/non-pending interrupt stale-interrupt/i) + 'stale-interrupt', + ) expect(continuation.calls).toHaveLength(0) }) @@ -551,8 +636,8 @@ describe('interrupt persistence', () => { }) const bad = mockAdapter([[runStarted(), interruptFinished()]]) - await expect( - collect( + expectResumeError( + await collect( chat({ adapter: bad.adapter, messages: [{ role: 'user', content: 'new input' }], @@ -562,7 +647,8 @@ describe('interrupt persistence', () => { middleware: [withPersistence(persistence)], }) as AsyncIterable, ), - ).rejects.toThrow(/missing resume entry.*interrupt-1/i) + 'interrupt-1', + ) const good = mockAdapter([[runStarted(), interruptFinished('r2')]]) await collect( @@ -590,8 +676,8 @@ describe('interrupt persistence', () => { }) const run = mockAdapter([[text('SHOULD NOT RUN')]]) - await expect( - collect( + expectResumeError( + await collect( chat({ adapter: run.adapter, messages: [{ role: 'user', content: 'new input' }], @@ -604,7 +690,8 @@ describe('interrupt persistence', () => { middleware: [withPersistence(persistence)], }) as AsyncIterable, ), - ).rejects.toThrow(/non-pending interrupt stale-interrupt/i) + 'stale-interrupt', + ) expect(run.calls).toHaveLength(0) expect(await persistence.stores.interrupts!.listPending('t1')).toHaveLength( @@ -975,4 +1062,691 @@ describe('interrupt persistence', () => { 'resolved', ) }) + + it('restores a registered generic interrupt after reload and gives its transformed response to the resolution hook', async () => { + const persistence = memoryPersistence() + const review = defineInterrupt({ + id: 'persisted-review', + payloadSchema: transformedDisplaySchema, + responseSchema: coercedCountSchema, + }) + const boundary = defineChatMiddleware({ + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterModel') return + return { + interrupts: [ + review.interrupt({ + key: 'one', + payload: 'Review this plan', + reason: 'review', + message: 'Review this plan', + }), + ], + } + }, + }) + const first = mockAdapter([[runStarted(), runFinished('r1')]]) + await collect( + chat({ + adapter: first.adapter, + interrupts: [review], + messages: [{ role: 'user', content: 'hi' }], + runId: 'r1', + threadId: 't1', + middleware: [boundary, withPersistence(persistence)], + }) as AsyncIterable, + ) + const interruptId = ( + await persistence.stores.interrupts!.listPending('t1') + )[0]?.interruptId + expect(interruptId).toBeDefined() + if (!interruptId) throw new Error('Expected a persisted generic interrupt') + + const observed: Array = [] + const resumed = mockAdapter([ + [runStarted(), text('continued'), runFinished('r1')], + ]) + await collect( + chat({ + adapter: resumed.adapter, + interrupts: [review], + messages: [], + runId: 'r1', + threadId: 't1', + resume: [ + { + interruptId, + status: 'resolved', + payload: { count: '2' }, + }, + ], + middleware: [ + defineChatMiddleware({ + onInterruptResolution(_ctx, resolutions) { + observed.push(...resolutions.for(review)) + return { toolResume: 'continue' } + }, + }), + withPersistence(persistence), + ], + }) as AsyncIterable, + ) + + expect(observed).toEqual([ + expect.objectContaining({ + status: 'resolved', + request: expect.objectContaining({ + payload: 'Review this plan'.length, + }), + response: { count: 2 }, + }), + ]) + expect( + (await persistence.stores.interrupts!.get(interruptId))?.status, + ).toBe('resolved') + }) + + it('resumes a registered generic record without blocking on a foreign persisted interrupt', async () => { + const persistence = memoryPersistence() + const review = defineInterrupt({ + id: 'mixed-persisted-review', + responseSchema: coercedCountSchema, + }) + const boundary = defineChatMiddleware({ + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterModel') return + return { + interrupts: [ + review.interrupt({ + key: 'review', + reason: 'review', + message: 'Review this plan', + }), + ], + } + }, + }) + const legacyInterruptId = 'legacy-external' + const initial = mockAdapter([ + [ + runStarted(), + { + type: EventType.RUN_FINISHED, + runId: 'r1', + threadId: 't1', + finishReason: 'stop', + timestamp: 1, + outcome: { + type: 'interrupt', + interrupts: [ + { + id: legacyInterruptId, + reason: 'legacy-review', + message: 'A legacy system needs a response', + }, + ], + }, + } satisfies StreamChunk, + ], + ]) + + await collect( + chat({ + adapter: initial.adapter, + interrupts: [review], + messages: [{ role: 'user', content: 'Start' }], + runId: 'r1', + threadId: 't1', + middleware: [boundary, withPersistence(persistence)], + }) as AsyncIterable, + ) + + const pending = await persistence.stores.interrupts!.listPending('t1') + expect(pending).toHaveLength(2) + const registered = pending.find( + (record) => record.interruptId !== legacyInterruptId, + )! + const legacy = pending.find( + (record) => record.interruptId === legacyInterruptId, + )! + expect(registered.payload).toMatchObject({ + metadata: { + 'tanstack:interruptBinding': { + definitionId: 'mixed-persisted-review', + key: 'review', + }, + }, + }) + expect(legacy.payload).toEqual({ + id: legacyInterruptId, + reason: 'legacy-review', + message: 'A legacy system needs a response', + }) + + const observed: Array = [] + const resumed = mockAdapter([ + [runStarted(), text('complete'), runFinished('r1')], + ]) + await collect( + chat({ + adapter: resumed.adapter, + interrupts: [review], + messages: [], + runId: 'r1', + threadId: 't1', + resume: [ + { + interruptId: registered.interruptId, + status: 'resolved', + payload: { count: '3' }, + }, + ], + middleware: [ + defineChatMiddleware({ + onInterruptResolution(_ctx, resolutions) { + observed.push(...resolutions.for(review)) + return { toolResume: 'continue' } + }, + }), + withPersistence(persistence), + ], + }) as AsyncIterable, + ) + + expect(observed).toEqual([ + expect.objectContaining({ + request: expect.objectContaining({ + definition: review, + key: 'review', + }), + status: 'resolved', + response: { count: 3 }, + }), + ]) + expect(resumed.calls).toHaveLength(1) + expect( + (await persistence.stores.interrupts!.get(registered.interruptId)) + ?.status, + ).toBe('resolved') + expect( + (await persistence.stores.interrupts!.get(legacyInterruptId))?.status, + ).toBe('pending') + }) + + it('restores two requests with the same definition and key by their unique instance ids', async () => { + const persistence = memoryPersistence() + const review = defineInterrupt({ + id: 'repeat-review', + responseSchema: coercedCountSchema, + }) + const firstChunks = await collect( + chat({ + adapter: mockAdapter([[runStarted(), runFinished('r1')]]).adapter, + interrupts: [review], + messages: [{ role: 'user', content: 'hi' }], + runId: 'r1', + threadId: 't1', + middleware: [ + defineChatMiddleware({ + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterModel') return + return { + interrupts: [ + review.interrupt({ + key: 'same', + reason: 'review', + message: 'First', + }), + review.interrupt({ + key: 'same', + reason: 'review', + message: 'Second', + }), + ], + } + }, + }), + withPersistence(persistence), + ], + }) as AsyncIterable, + ) + const terminal = firstChunks.find(isInterruptTerminal) + const ids = + terminal?.outcome.interrupts.map((interrupt) => interrupt.id) ?? [] + expect(ids).toHaveLength(2) + expect(new Set(ids).size).toBe(2) + + const observed: Array = [] + await collect( + chat({ + adapter: mockAdapter([ + [runStarted(), text('continued'), runFinished('r1')], + ]).adapter, + interrupts: [review], + messages: [], + runId: 'r1', + threadId: 't1', + resume: ids.map((interruptId, index) => ({ + interruptId, + status: 'resolved' as const, + payload: { count: String(index + 1) }, + })), + middleware: [ + defineChatMiddleware({ + onInterruptResolution(_ctx, resolutions) { + observed.push(...resolutions.for(review)) + return { toolResume: 'continue' } + }, + }), + withPersistence(persistence), + ], + }) as AsyncIterable, + ) + expect(observed).toHaveLength(2) + expect(await persistence.stores.interrupts!.listPending('t1')).toEqual([]) + }) + + it('returns a structured conflict error for duplicate durable resume entries', async () => { + const persistence = memoryPersistence() + await persistence.stores.interrupts!.create({ + interruptId: 'duplicate-me', + runId: 'r1', + threadId: 't1', + requestedAt: 1, + payload: {}, + }) + const { adapter } = mockAdapter([[text('SHOULD NOT RUN')]]) + + const chunks = await collect( + chat({ + adapter, + messages: [], + runId: 'r1', + threadId: 't1', + resume: [ + { interruptId: 'duplicate-me', status: 'resolved' }, + { interruptId: 'duplicate-me', status: 'cancelled' }, + ], + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ) + + const error = chunks.find((chunk) => chunk.type === EventType.RUN_ERROR) + expect(error?.['tanstack:interruptErrors']).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + scope: 'item', + interruptId: 'duplicate-me', + code: 'conflict', + }), + ]), + ) + expect(await persistence.stores.interrupts!.listPending('t1')).toHaveLength( + 1, + ) + }) + + it('rejects pending interrupts from more than one run on the same thread', async () => { + const persistence = memoryPersistence() + await persistence.stores.interrupts!.create({ + interruptId: 'from-run-1', + runId: 'run-1', + threadId: 't1', + requestedAt: 1, + payload: {}, + }) + await persistence.stores.interrupts!.create({ + interruptId: 'from-run-2', + runId: 'run-2', + threadId: 't1', + requestedAt: 2, + payload: {}, + }) + + const chunks = await collect( + chat({ + adapter: mockAdapter([[text('SHOULD NOT RUN')]]).adapter, + messages: [], + runId: 'run-3', + threadId: 't1', + resume: [ + { interruptId: 'from-run-1', status: 'resolved' }, + { interruptId: 'from-run-2', status: 'cancelled' }, + ], + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ) + + const error = chunks.find((chunk) => chunk.type === EventType.RUN_ERROR) + expect(error?.['tanstack:interruptErrors']).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + scope: 'batch', + code: 'stale', + message: 'Thread has pending interrupts from more than one run.', + }), + ]), + ) + expect(await persistence.stores.interrupts!.listPending('t1')).toHaveLength( + 2, + ) + }) + + it('keeps an opaque approval when a generic interrupt is also resumed', async () => { + const persistence = memoryPersistence() + const review = defineInterrupt({ + id: 'merge-review', + responseSchema: coercedCountSchema, + }) + const first = await collect( + chat({ + adapter: mockAdapter([[runStarted(), runFinished('r1')]]).adapter, + interrupts: [review], + messages: [{ role: 'user', content: 'hi' }], + runId: 'r1', + threadId: 't1', + middleware: [ + defineChatMiddleware({ + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterModel') return + return { + interrupts: [ + review.interrupt({ + key: 'review', + reason: 'review', + message: 'Review this plan', + }), + ], + } + }, + }), + withPersistence(persistence), + ], + }) as AsyncIterable, + ) + const terminal = first.find(isInterruptTerminal) + const genericId = terminal?.outcome.interrupts[0]?.id + if (!genericId) throw new Error('Expected a persisted generic interrupt') + + await persistence.stores.interrupts!.create({ + interruptId: 'approval-1', + runId: 'r1', + threadId: 't1', + requestedAt: 2, + payload: { toolCallId: 'tc1', metadata: { kind: 'approval' } }, + }) + + const resumeStates: Array = [] + await collect( + chat({ + adapter: mockAdapter([ + [runStarted(), text('continued'), runFinished('r1')], + ]).adapter, + interrupts: [review], + messages: [], + runId: 'r1', + threadId: 't1', + resume: [ + { + interruptId: 'approval-1', + status: 'resolved', + payload: { approved: true }, + }, + { + interruptId: genericId, + status: 'resolved', + payload: { count: '2' }, + }, + ], + middleware: [ + withPersistence(persistence), + defineChatMiddleware({ + name: 'observe-merged-resume-state', + onConfig(_ctx, config) { + resumeStates.push(config.resumeToolState) + }, + onInterruptResolution() { + return { toolResume: 'continue' } + }, + }), + ], + }) as AsyncIterable, + ) + + expect(resumeStates[0]?.approvals?.get('approval-1')).toBe(true) + expect(resumeStates[0]?.genericInterruptRequests?.has(genericId)).toBe(true) + }) + + it('commits a mixed resume batch once when the store supports commitBatch', async () => { + const persistence = memoryPersistence() + const store = persistence.stores.interrupts! + const defaultCommitBatch = store.commitBatch!.bind(store) + const resolve = vi.spyOn(store, 'resolve') + const cancel = vi.spyOn(store, 'cancel') + const commitBatch = vi.fn(defaultCommitBatch) + store.commitBatch = commitBatch + + await store.create({ + interruptId: 'resolved-1', + runId: 'r1', + threadId: 't1', + requestedAt: 1, + payload: {}, + }) + await store.create({ + interruptId: 'cancelled-1', + runId: 'r1', + threadId: 't1', + requestedAt: 2, + payload: {}, + }) + + const { adapter } = mockAdapter([[runStarted(), text('ok'), runFinished()]]) + await collect( + chat({ + adapter, + messages: [], + runId: 'r1', + threadId: 't1', + resume: [ + { + interruptId: 'resolved-1', + status: 'resolved', + payload: { answer: 'yes' }, + }, + { interruptId: 'cancelled-1', status: 'cancelled' }, + ], + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ) + + expect(commitBatch).toHaveBeenCalledTimes(1) + expect(commitBatch).toHaveBeenCalledWith([ + { + interruptId: 'resolved-1', + status: 'resolved', + response: { answer: 'yes' }, + }, + { interruptId: 'cancelled-1', status: 'cancelled' }, + ]) + expect(resolve).not.toHaveBeenCalled() + expect(cancel).not.toHaveBeenCalled() + expect((await store.get('resolved-1'))?.status).toBe('resolved') + expect((await store.get('cancelled-1'))?.status).toBe('cancelled') + }) + + it('uses legacy writes when an old interrupt store has no commitBatch method', async () => { + const persistence = memoryPersistence() + const base = persistence.stores.interrupts! + const resolve = vi.fn((interruptId: string, response?: unknown) => + base.resolve(interruptId, response), + ) + const cancel = vi.fn((interruptId: string) => base.cancel(interruptId)) + const legacyStore: InterruptStore = { + create: (record) => base.create(record), + resolve, + cancel, + 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), + } + persistence.stores.interrupts = legacyStore + + await legacyStore.create({ + interruptId: 'resolved-1', + runId: 'r1', + threadId: 't1', + requestedAt: 1, + payload: {}, + }) + await legacyStore.create({ + interruptId: 'cancelled-1', + runId: 'r1', + threadId: 't1', + requestedAt: 2, + payload: {}, + }) + + const { adapter } = mockAdapter([[runStarted(), text('ok'), runFinished()]]) + await collect( + chat({ + adapter, + messages: [], + runId: 'r1', + threadId: 't1', + resume: [ + { interruptId: 'resolved-1', status: 'resolved', payload: 'yes' }, + { interruptId: 'cancelled-1', status: 'cancelled' }, + ], + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ) + + expect(resolve).toHaveBeenCalledWith('resolved-1', 'yes') + expect(cancel).toHaveBeenCalledWith('cancelled-1') + }) + + it('preflights a memory interrupt batch before it changes any record', async () => { + const persistence = memoryPersistence() + const store = persistence.stores.interrupts! + await store.create({ + interruptId: 'pending-1', + runId: 'r1', + threadId: 't1', + requestedAt: 1, + payload: {}, + }) + await store.create({ + interruptId: 'pending-2', + runId: 'r1', + threadId: 't1', + requestedAt: 2, + payload: {}, + }) + await store.create({ + interruptId: 'terminal-1', + runId: 'r1', + threadId: 't1', + requestedAt: 3, + payload: {}, + }) + await store.resolve('terminal-1', { done: true }) + + await expect( + store.commitBatch!([ + { interruptId: 'pending-1', status: 'resolved', response: 'yes' }, + { interruptId: 'missing-1', status: 'cancelled' }, + ]), + ).rejects.toThrow('missing id: missing-1') + expect((await store.get('pending-1'))?.status).toBe('pending') + + await expect( + store.commitBatch!([ + { interruptId: 'pending-1', status: 'resolved', response: 'yes' }, + { interruptId: 'pending-1', status: 'cancelled' }, + ]), + ).rejects.toThrow('duplicate id: pending-1') + expect((await store.get('pending-1'))?.status).toBe('pending') + + await expect( + store.commitBatch!([ + { interruptId: 'pending-2', status: 'resolved', response: 'yes' }, + { interruptId: 'terminal-1', status: 'cancelled' }, + ]), + ).rejects.toThrow('non-pending id: terminal-1') + expect((await store.get('pending-2'))?.status).toBe('pending') + expect((await store.get('terminal-1'))?.status).toBe('resolved') + }) + + it('keeps a mixed resume batch retryable when commitBatch rejects', async () => { + const persistence = memoryPersistence() + const store = persistence.stores.interrupts! + const defaultCommitBatch = store.commitBatch!.bind(store) + const resolve = vi.spyOn(store, 'resolve') + const cancel = vi.spyOn(store, 'cancel') + const commitBatch = vi.fn(async () => { + throw new Error('batch write failed') + }) + store.commitBatch = commitBatch + await store.create({ + interruptId: 'resolved-1', + runId: 'r1', + threadId: 't1', + requestedAt: 1, + payload: {}, + }) + await store.create({ + interruptId: 'cancelled-1', + runId: 'r1', + threadId: 't1', + requestedAt: 2, + payload: {}, + }) + + const { adapter } = mockAdapter([[runStarted(), text('ok'), runFinished()]]) + await expect( + collect( + chat({ + adapter, + messages: [], + runId: 'r1', + threadId: 't1', + resume: [ + { interruptId: 'resolved-1', status: 'resolved', payload: 'yes' }, + { interruptId: 'cancelled-1', status: 'cancelled' }, + ], + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ), + ).rejects.toThrow('batch write failed') + expect(await store.listPending('t1')).toHaveLength(2) + expect((await store.get('resolved-1'))?.status).toBe('pending') + expect((await store.get('cancelled-1'))?.status).toBe('pending') + expect(resolve).not.toHaveBeenCalled() + expect(cancel).not.toHaveBeenCalled() + expect((await persistence.stores.runs!.get('r1'))?.status).toBe('failed') + + store.commitBatch = vi.fn(defaultCommitBatch) + const retry = mockAdapter([[runStarted(), text('retried'), runFinished()]]) + const chunks = await collect( + chat({ + adapter: retry.adapter, + messages: [], + runId: 'r1', + threadId: 't1', + resume: [ + { interruptId: 'resolved-1', status: 'resolved', payload: 'yes' }, + { interruptId: 'cancelled-1', status: 'cancelled' }, + ], + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ) + + expect(chunks).toContainEqual(expect.objectContaining({ delta: 'retried' })) + expect(await store.listPending('t1')).toEqual([]) + expect((await store.get('resolved-1'))?.status).toBe('resolved') + expect((await store.get('cancelled-1'))?.status).toBe('cancelled') + }) }) diff --git a/packages/ai-persistence/tests/persistence-fixtures.ts b/packages/ai-persistence/tests/persistence-fixtures.ts index ad434c461e..668f60d8f1 100644 --- a/packages/ai-persistence/tests/persistence-fixtures.ts +++ b/packages/ai-persistence/tests/persistence-fixtures.ts @@ -88,6 +88,7 @@ export function createInterruptStore(): InterruptStore { create: () => Promise.resolve(), resolve: () => Promise.resolve(), cancel: () => Promise.resolve(), + commitBatch: () => Promise.resolve(), get: () => Promise.resolve(null), list: () => Promise.resolve([]), listPending: () => Promise.resolve([]), diff --git a/packages/ai-persistence/tests/with-persistence.test.ts b/packages/ai-persistence/tests/with-persistence.test.ts index 338431c927..edec73f6a5 100644 --- a/packages/ai-persistence/tests/with-persistence.test.ts +++ b/packages/ai-persistence/tests/with-persistence.test.ts @@ -79,13 +79,6 @@ async function collect(stream: AsyncIterable) { return out } -async function expectCollectRejects( - stream: AsyncIterable, - pattern: RegExp, -) { - await expect(collect(stream)).rejects.toThrow(pattern) -} - function serverSearchTool(): Tool { return { name: 'search', @@ -522,7 +515,7 @@ describe('withPersistence (state-only)', () => { ) const next = mockAdapter([[ev.text('SHOULD NOT RUN')]]) - await expectCollectRejects( + const blockedChunks = await collect( chat({ adapter: next.adapter, messages: [{ role: 'user', content: 'new input' }], @@ -530,7 +523,22 @@ describe('withPersistence (state-only)', () => { threadId: 't1', middleware: [withPersistence(persistence)], }) as AsyncIterable, - /pending interrupts.*resume is required/i, + ) + const blockedError = blockedChunks.find( + (chunk) => chunk.type === EventType.RUN_ERROR, + ) + expect(blockedError?.['tanstack:interruptErrors']).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + scope: 'item', + interruptId: 'interrupt-1', + code: 'unknown-interrupt', + }), + expect.objectContaining({ + scope: 'batch', + code: 'incomplete-batch', + }), + ]), ) expect(next.calls.length).toBe(0) }) @@ -550,7 +558,7 @@ describe('withPersistence (state-only)', () => { ) const next = mockAdapter([[ev.text('SHOULD NOT RUN')]]) - await expectCollectRejects( + const mismatchChunks = await collect( chat({ adapter: next.adapter, messages: [{ role: 'user', content: 'new input' }], @@ -559,7 +567,22 @@ describe('withPersistence (state-only)', () => { resume: [{ interruptId: 'other-interrupt', status: 'resolved' }], middleware: [withPersistence(persistence)], }) as AsyncIterable, - /missing resume entry for pending interrupt interrupt-1/i, + ) + const mismatchError = mismatchChunks.find( + (chunk) => chunk.type === EventType.RUN_ERROR, + ) + expect(mismatchError?.['tanstack:interruptErrors']).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + scope: 'item', + interruptId: 'interrupt-1', + code: 'unknown-interrupt', + }), + expect.objectContaining({ + scope: 'batch', + code: 'incomplete-batch', + }), + ]), ) expect(next.calls.length).toBe(0) }) diff --git a/packages/ai-preact/src/types.ts b/packages/ai-preact/src/types.ts index 02cd74ad2a..928e681f6a 100644 --- a/packages/ai-preact/src/types.ts +++ b/packages/ai-preact/src/types.ts @@ -1,5 +1,6 @@ import type { AnyClientTool, + InterruptDefinition, ModelMessage, RunAgentResumeItem, SchemaInput, @@ -9,7 +10,7 @@ import type { BoundInterrupts, ChatClientOptions, ChatClientState, - ChatInterrupt, + ResolvableChatInterrupt, ChatInterruptState, ChatRequestBody, ChatResumeState, @@ -59,8 +60,10 @@ export type { export type UseChatOptions< TTools extends ReadonlyArray = any, TContext = InferredClientContext, + TInterrupts extends ReadonlyArray> = + readonly [], > = DistributedOmit< - ChatClientOptions, + ChatClientOptions, | 'onMessagesChange' | 'onLoadingChange' | 'onErrorChange' @@ -92,6 +95,8 @@ export type UseChatOptions< export interface UseChatReturn< TTools extends ReadonlyArray = any, + TInterrupts extends ReadonlyArray> = + readonly [], > { /** * Current messages in the conversation @@ -153,14 +158,18 @@ export interface UseChatReturn< * it, correlate a log line). */ runId: string | null - interrupts: BoundInterrupts + interrupts: BoundInterrupts /** @deprecated Use `interrupts`. */ - pendingInterrupts: BoundInterrupts - interruptErrors: ChatInterruptState['interruptErrors'] + pendingInterrupts: BoundInterrupts + interruptErrors: ChatInterruptState['interruptErrors'] resuming: boolean resolveInterrupts: { (approved: boolean): void - (resolver: (interrupt: ChatInterrupt) => undefined): void + ( + resolver: ( + interrupt: ResolvableChatInterrupt, + ) => undefined, + ): void } cancelInterrupts: () => void retryInterrupts: () => void diff --git a/packages/ai-preact/src/use-chat.ts b/packages/ai-preact/src/use-chat.ts index bd50d1180d..18042855bb 100644 --- a/packages/ai-preact/src/use-chat.ts +++ b/packages/ai-preact/src/use-chat.ts @@ -10,7 +10,7 @@ import { } from 'preact/hooks' import type { ChatClientState, - ChatInterrupt, + ResolvableChatInterrupt, ChatInterruptState, ChatResumeState, ConnectionStatus, @@ -20,6 +20,7 @@ import type { } from '@tanstack/ai-client' import type { AnyClientTool, + InterruptDefinition, ModelMessage, RunAgentResumeItem, } from '@tanstack/ai' @@ -37,7 +38,12 @@ const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) export function useChat< const TTools extends ReadonlyArray = any, TContext = InferredClientContext, ->(options: UseChatOptions): UseChatReturn { + const TInterrupts extends ReadonlyArray< + InterruptDefinition + > = readonly [], +>( + options: UseChatOptions, +): UseChatReturn { // The hook's identity is its `threadId` — also the persistence key, so a // reload with the same `threadId` restores the same conversation. `hookId` is // only a stable fallback for client-recreation keying when no `threadId` is @@ -58,7 +64,7 @@ export function useChat< const [queue, setQueue] = useState>([]) const [runId, setRunId] = useState(null) const [interruptState, setInterruptState] = useState< - ChatInterruptState + ChatInterruptState >(() => ({ interrupts: EMPTY_INTERRUPTS, pendingInterrupts: EMPTY_INTERRUPTS, @@ -79,7 +85,8 @@ export function useChat< client: ChatClient timeout: ReturnType } | null>(null) - const optionsRef = useRef>(options) + const optionsRef = + useRef>(options) optionsRef.current = options @@ -107,7 +114,7 @@ export function useChat< : { fetcher: initialOptions.fetcher } const instanceHolder: { - current: ChatClient | undefined + current: ChatClient | undefined } = { current: undefined } const getActiveInstance = () => { const currentInstance = instanceHolder.current @@ -117,7 +124,7 @@ export function useChat< return currentInstance } const pendingInitializationErrors: Array = [] - const instance = new ChatClient({ + const instance = new ChatClient({ devtoolsBridgeFactory: createChatDevtoolsBridge, ...transport, initialMessages: messagesToUse, @@ -174,6 +181,9 @@ export function useChat< ...(initialOptions.tools !== undefined && { tools: initialOptions.tools, }), + ...(initialOptions.interrupts !== undefined && { + interrupts: initialOptions.interrupts, + }), ...(options.streamProcessor !== undefined && { streamProcessor: options.streamProcessor, }), @@ -433,7 +443,11 @@ export function useChat< const resolveInterrupts = useCallback( ( - resolution: boolean | ((interrupt: ChatInterrupt) => undefined), + resolution: + | boolean + | (( + interrupt: ResolvableChatInterrupt, + ) => undefined), ) => { if (typeof resolution === 'boolean') { client.resolveInterrupts(resolution) diff --git a/packages/ai-preact/tests/use-chat-types.test.ts b/packages/ai-preact/tests/use-chat-types.test.ts index 0d65ab115c..44aa41bddb 100644 --- a/packages/ai-preact/tests/use-chat-types.test.ts +++ b/packages/ai-preact/tests/use-chat-types.test.ts @@ -4,7 +4,7 @@ */ import { describe, expectTypeOf, it } from 'vitest' -import { toolDefinition } from '@tanstack/ai' +import { defineInterrupt, toolDefinition } from '@tanstack/ai' import { clientTools } from '@tanstack/ai-client' import { useChat } from '../src/use-chat' import type { UseChatOptions, UseChatReturn } from '../src/types' @@ -15,6 +15,16 @@ type TestSchema = { readonly vendor: 'test' readonly types: { readonly input: T; readonly output: T } readonly validate: (value: unknown) => { readonly value: T } + readonly jsonSchema: { readonly input: () => Record } + } +} +type TransformSchema = { + readonly '~standard': { + readonly version: 1 + readonly vendor: 'test' + readonly types: { readonly input: string; readonly output: number } + readonly validate: (value: unknown) => { readonly value: number } + readonly jsonSchema: { readonly input: () => Record } } } @@ -132,6 +142,7 @@ describe('useChat() interrupt types', () => { output: { accountId: '' }, }, validate: () => ({ value: { accountId: '' } }), + jsonSchema: { input: () => ({ type: 'object' }) }, }, } const transfer = toolDefinition({ @@ -201,3 +212,102 @@ describe('useChat() interrupt types', () => { void check }) }) + +describe('useChat() registered generic interrupt types', () => { + it('keeps registered and external generic interrupts distinct', () => { + const payloadSchema: TestSchema<{ title: string }> = { + '~standard': { + version: 1, + vendor: 'test', + types: { input: { title: '' }, output: { title: '' } }, + validate: () => ({ value: { title: '' } }), + jsonSchema: { input: () => ({ type: 'object' }) }, + }, + } + const responseSchema: TransformSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { input: '', output: 0 }, + validate: () => ({ value: 0 }), + jsonSchema: { input: () => ({ type: 'string' }) }, + }, + } + const reviewPlan = defineInterrupt({ + id: 'review-plan', + payloadSchema, + responseSchema, + }) + const acknowledge = defineInterrupt({ + id: 'acknowledge', + responseSchema: payloadSchema, + }) + + const check = () => { + const chat = useChat({ + connection: { connect: async function* () {} }, + interrupts: [reviewPlan, acknowledge], + }) + type Interrupt = (typeof chat.interrupts)[number] + type Review = Extract + type External = Extract< + Exclude, + { kind: 'generic' } + > + type Unbound = Extract + type CallbackInterrupt = typeof chat.resolveInterrupts extends { + (resolver: (interrupt: infer TInterrupt) => undefined): void + } + ? TInterrupt + : never + + expectTypeOf().toEqualTypeOf< + { title: string } | undefined + >() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + expectTypeOf().toEqualTypeOf() + expectTypeOf< + Extract + >().toEqualTypeOf() + + const resolveReview = (review: Review) => { + review.resolveInterrupt('42') + // @ts-expect-error The transformed response still accepts its input type. + review.resolveInterrupt(42) + } + void resolveReview + chat.resolveInterrupts((interrupt) => { + interrupt.cancel() + return undefined + }) + + const existingTool = toolDefinition({ + name: 'preact-unregistered-tool', + description: 'A tool without an interrupt registry', + needsApproval: true, + }).client() + const withoutRegistry = useChat({ + connection: { connect: async function* () {} }, + tools: clientTools(existingTool), + }) + type WithoutRegistry = (typeof withoutRegistry.interrupts)[number] + type ExistingToolInterrupt = Extract< + WithoutRegistry, + { kind: 'tool-approval' } + > + type UnregisteredGeneric = Extract + expectTypeOf< + ExistingToolInterrupt['toolName'] + >().toEqualTypeOf<'preact-unregistered-tool'>() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + } + void check + }) +}) diff --git a/packages/ai-react/src/index.ts b/packages/ai-react/src/index.ts index 75bff7c3ce..4fbad82626 100644 --- a/packages/ai-react/src/index.ts +++ b/packages/ai-react/src/index.ts @@ -98,6 +98,8 @@ export { type FetchConnectionOptions, type XhrConnectionOptions, type InferChatMessages, + type GenericInterrupt, + type RegisteredGenericInterrupt, type GenerationClientState, type ImageGenerateInput, type AudioGenerateInput, diff --git a/packages/ai-react/src/types.ts b/packages/ai-react/src/types.ts index 02a413df0b..0dfccfad25 100644 --- a/packages/ai-react/src/types.ts +++ b/packages/ai-react/src/types.ts @@ -1,5 +1,6 @@ import type { AnyClientTool, + InterruptDefinition, InferSchemaType, ModelMessage, RunAgentResumeItem, @@ -10,7 +11,7 @@ import type { BoundInterrupts, ChatClientOptions, ChatClientState, - ChatInterrupt, + ResolvableChatInterrupt, ChatInterruptState, ChatRequestBody, ChatResumeState, @@ -83,8 +84,10 @@ export type UseChatOptions< TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, TContext = InferredClientContext, + TInterrupts extends ReadonlyArray> = + readonly [], > = DistributedOmit< - ChatClientOptions, + ChatClientOptions, | 'onMessagesChange' | 'onLoadingChange' | 'onErrorChange' @@ -126,9 +129,12 @@ export type UseChatOptions< export type UseChatReturn< TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, + TInterrupts extends ReadonlyArray> = + readonly [], > = BaseUseChatReturn< TTools, - TSchema extends SchemaInput ? InferSchemaType : unknown + TSchema extends SchemaInput ? InferSchemaType : unknown, + TInterrupts > & (TSchema extends SchemaInput ? { @@ -151,6 +157,8 @@ export type UseChatReturn< interface BaseUseChatReturn< TTools extends ReadonlyArray = any, TData = unknown, + TInterrupts extends ReadonlyArray> = + readonly [], > { /** * Current messages in the conversation. When `outputSchema` is supplied, @@ -218,14 +226,18 @@ interface BaseUseChatReturn< * it, correlate a log line). */ runId: string | null - interrupts: BoundInterrupts + interrupts: BoundInterrupts /** @deprecated Use `interrupts`. */ - pendingInterrupts: BoundInterrupts - interruptErrors: ChatInterruptState['interruptErrors'] + pendingInterrupts: BoundInterrupts + interruptErrors: ChatInterruptState['interruptErrors'] resuming: boolean resolveInterrupts: { (approved: boolean): void - (resolver: (interrupt: ChatInterrupt) => undefined): void + ( + resolver: ( + interrupt: ResolvableChatInterrupt, + ) => undefined, + ): void } cancelInterrupts: () => void retryInterrupts: () => void diff --git a/packages/ai-react/src/use-chat.ts b/packages/ai-react/src/use-chat.ts index 96a75eb9ca..a0b4430146 100644 --- a/packages/ai-react/src/use-chat.ts +++ b/packages/ai-react/src/use-chat.ts @@ -4,6 +4,7 @@ import { useCallback, useEffect, useId, useMemo, useRef, useState } from 'react' import type { AnyClientTool, InferSchemaType, + InterruptDefinition, ModelMessage, RunAgentResumeItem, SchemaInput, @@ -11,7 +12,7 @@ import type { } from '@tanstack/ai/client' import type { ChatClientState, - ChatInterrupt, + ResolvableChatInterrupt, ChatInterruptState, ChatResumeState, ConnectionStatus, @@ -36,9 +37,12 @@ export function useChat< const TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, TContext = InferredClientContext, + const TInterrupts extends ReadonlyArray< + InterruptDefinition + > = readonly [], >( - options: UseChatOptions, -): UseChatReturn { + options: UseChatOptions, +): UseChatReturn { // The hook's identity is its `threadId` — also the persistence key, so a // reload with the same `threadId` restores the same conversation. `hookId` is // only a stable fallback for React's client-recreation keying when no @@ -59,7 +63,7 @@ export function useChat< const [queue, setQueue] = useState>([]) const [runId, setRunId] = useState(null) const [interruptState, setInterruptState] = useState< - ChatInterruptState + ChatInterruptState >(() => ({ interrupts: EMPTY_INTERRUPTS, pendingInterrupts: EMPTY_INTERRUPTS, @@ -89,7 +93,8 @@ export function useChat< messagesRef.current = messages // Track current options in a ref to avoid recreating client when options change - const optionsRef = useRef>(options) + const optionsRef = + useRef>(options) optionsRef.current = options const syncResumeState = useCallback((target: ChatClient | null) => { @@ -113,7 +118,7 @@ export function useChat< : { fetcher: initialOptions.fetcher } const instanceHolder: { - current: ChatClient | undefined + current: ChatClient | undefined } = { current: undefined } const getActiveInstance = () => { const currentInstance = instanceHolder.current @@ -123,7 +128,7 @@ export function useChat< return currentInstance } const pendingInitializationErrors: Array = [] - const instance = new ChatClient({ + const instance = new ChatClient({ devtoolsBridgeFactory: createChatDevtoolsBridge, ...transport, initialMessages: messagesToUse, @@ -173,6 +178,9 @@ export function useChat< ...(initialOptions.tools !== undefined && { tools: initialOptions.tools, }), + ...(initialOptions.interrupts !== undefined && { + interrupts: initialOptions.interrupts, + }), onCustomEvent: (eventType, data, context) => { if (!getActiveInstance()) return optionsRef.current.onCustomEvent?.(eventType, data, context) @@ -458,7 +466,11 @@ export function useChat< const resolveInterrupts = useCallback( ( - resolution: boolean | ((interrupt: ChatInterrupt) => undefined), + resolution: + | boolean + | (( + interrupt: ResolvableChatInterrupt, + ) => undefined), ) => { if (typeof resolution === 'boolean') { client.resolveInterrupts(resolution) @@ -562,5 +574,5 @@ export function useChat< resumeInterrupts, partial, final, - } as unknown as UseChatReturn + } as unknown as UseChatReturn } diff --git a/packages/ai-react/tests/use-chat-types.test.ts b/packages/ai-react/tests/use-chat-types.test.ts index 7170b3b8d1..5f603973cc 100644 --- a/packages/ai-react/tests/use-chat-types.test.ts +++ b/packages/ai-react/tests/use-chat-types.test.ts @@ -7,7 +7,7 @@ */ import { describe, expectTypeOf, it } from 'vitest' -import { toolDefinition } from '@tanstack/ai' +import { defineInterrupt, toolDefinition } from '@tanstack/ai' import { clientTools } from '@tanstack/ai-client' import { useChat } from '../src/use-chat' import type { AnyClientTool } from '@tanstack/ai' @@ -23,6 +23,16 @@ type TestSchema = { readonly vendor: 'test' readonly types: { readonly input: T; readonly output: T } readonly validate: (value: unknown) => { readonly value: T } + readonly jsonSchema: { readonly input: () => Record } + } +} +type TransformSchema = { + readonly '~standard': { + readonly version: 1 + readonly vendor: 'test' + readonly types: { readonly input: string; readonly output: number } + readonly validate: (value: unknown) => { readonly value: number } + readonly jsonSchema: { readonly input: () => Record } } } @@ -245,6 +255,7 @@ describe('useChat() interrupt types', () => { output: { accountId: '' }, }, validate: () => ({ value: { accountId: '' } }), + jsonSchema: { input: () => ({ type: 'object' }) }, }, } const transfer = toolDefinition({ @@ -314,3 +325,102 @@ describe('useChat() interrupt types', () => { void check }) }) + +describe('useChat() registered generic interrupt types', () => { + it('keeps registered and external generic interrupts distinct', () => { + const payloadSchema: TestSchema<{ title: string }> = { + '~standard': { + version: 1, + vendor: 'test', + types: { input: { title: '' }, output: { title: '' } }, + validate: () => ({ value: { title: '' } }), + jsonSchema: { input: () => ({ type: 'object' }) }, + }, + } + const responseSchema: TransformSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { input: '', output: 0 }, + validate: () => ({ value: 0 }), + jsonSchema: { input: () => ({ type: 'string' }) }, + }, + } + const reviewPlan = defineInterrupt({ + id: 'review-plan', + payloadSchema, + responseSchema, + }) + const acknowledge = defineInterrupt({ + id: 'acknowledge', + responseSchema: payloadSchema, + }) + + const check = () => { + const chat = useChat({ + connection: { connect: async function* () {} }, + interrupts: [reviewPlan, acknowledge], + }) + type Interrupt = (typeof chat.interrupts)[number] + type Review = Extract + type External = Extract< + Exclude, + { kind: 'generic' } + > + type Unbound = Extract + type CallbackInterrupt = typeof chat.resolveInterrupts extends { + (resolver: (interrupt: infer TInterrupt) => undefined): void + } + ? TInterrupt + : never + + expectTypeOf().toEqualTypeOf< + { title: string } | undefined + >() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + expectTypeOf().toEqualTypeOf() + expectTypeOf< + Extract + >().toEqualTypeOf() + + const resolveReview = (review: Review) => { + review.resolveInterrupt('42') + // @ts-expect-error The transformed response still accepts its input type. + review.resolveInterrupt(42) + } + void resolveReview + chat.resolveInterrupts((interrupt) => { + interrupt.cancel() + return undefined + }) + + const existingTool = toolDefinition({ + name: 'react-unregistered-tool', + description: 'A tool without an interrupt registry', + needsApproval: true, + }).client() + const withoutRegistry = useChat({ + connection: { connect: async function* () {} }, + tools: clientTools(existingTool), + }) + type WithoutRegistry = (typeof withoutRegistry.interrupts)[number] + type ExistingToolInterrupt = Extract< + WithoutRegistry, + { kind: 'tool-approval' } + > + type UnregisteredGeneric = Extract + expectTypeOf< + ExistingToolInterrupt['toolName'] + >().toEqualTypeOf<'react-unregistered-tool'>() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + } + void check + }) +}) diff --git a/packages/ai-solid/src/types.ts b/packages/ai-solid/src/types.ts index 8eb5c69584..73002892fd 100644 --- a/packages/ai-solid/src/types.ts +++ b/packages/ai-solid/src/types.ts @@ -1,5 +1,6 @@ import type { AnyClientTool, + InterruptDefinition, InferSchemaType, ModelMessage, RunAgentResumeItem, @@ -10,7 +11,7 @@ import type { BoundInterrupts, ChatClientOptions, ChatClientState, - ChatInterrupt, + ResolvableChatInterrupt, ChatInterruptState, ChatRequestBody, ChatResumeState, @@ -79,8 +80,10 @@ export type UseChatOptions< TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, TContext = InferredClientContext, + TInterrupts extends ReadonlyArray> = + readonly [], > = DistributedOmit< - ChatClientOptions, + ChatClientOptions, | 'onMessagesChange' | 'onLoadingChange' | 'onErrorChange' @@ -115,9 +118,12 @@ export type UseChatOptions< export type UseChatReturn< TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, + TInterrupts extends ReadonlyArray> = + readonly [], > = BaseUseChatReturn< TTools, - TSchema extends SchemaInput ? InferSchemaType : unknown + TSchema extends SchemaInput ? InferSchemaType : unknown, + TInterrupts > & (TSchema extends SchemaInput ? { @@ -137,6 +143,8 @@ export type UseChatReturn< interface BaseUseChatReturn< TTools extends ReadonlyArray = any, TData = unknown, + TInterrupts extends ReadonlyArray> = + readonly [], > { /** * Current messages in the conversation. When `outputSchema` is supplied, @@ -204,14 +212,20 @@ interface BaseUseChatReturn< * it, correlate a log line). */ runId: Accessor - interrupts: Accessor> + interrupts: Accessor> /** @deprecated Use `interrupts`. */ - pendingInterrupts: Accessor> - interruptErrors: Accessor['interruptErrors']> + pendingInterrupts: Accessor> + interruptErrors: Accessor< + ChatInterruptState['interruptErrors'] + > resuming: Accessor resolveInterrupts: { (approved: boolean): void - (resolver: (interrupt: ChatInterrupt) => undefined): void + ( + resolver: ( + interrupt: ResolvableChatInterrupt, + ) => undefined, + ): void } cancelInterrupts: () => void retryInterrupts: () => void diff --git a/packages/ai-solid/src/use-chat.ts b/packages/ai-solid/src/use-chat.ts index 6ac73e332a..bc89683849 100644 --- a/packages/ai-solid/src/use-chat.ts +++ b/packages/ai-solid/src/use-chat.ts @@ -11,7 +11,7 @@ import { ChatClient } from '@tanstack/ai-client' import { createChatDevtoolsBridge } from '@tanstack/ai-client/devtools' import type { ChatClientState, - ChatInterrupt, + ResolvableChatInterrupt, ChatInterruptState, ChatResumeState, ConnectionStatus, @@ -22,6 +22,7 @@ import type { } from '@tanstack/ai-client' import type { AnyClientTool, + InterruptDefinition, InferSchemaType, ModelMessage, RunAgentResumeItem, @@ -43,13 +44,17 @@ export function useChat< const TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, TContext = InferredClientContext, + const TInterrupts extends ReadonlyArray< + InterruptDefinition + > = readonly [], >( - options: UseChatOptions = {} as UseChatOptions< + options: UseChatOptions< TTools, TSchema, - TContext - >, -): UseChatReturn { + TContext, + TInterrupts + > = {} as UseChatOptions, +): UseChatReturn { // The hook's identity is its `threadId` — also the persistence key, so a // reload with the same `threadId` restores the same conversation. `hookId` is // only a stable fallback for client-recreation keying when no `threadId` is @@ -70,7 +75,7 @@ export function useChat< const [queue, setQueue] = createSignal>([]) const [runId, setRunId] = createSignal(null) const [interruptState, setInterruptState] = createSignal< - ChatInterruptState + ChatInterruptState >({ interrupts: EMPTY_INTERRUPTS, pendingInterrupts: EMPTY_INTERRUPTS, @@ -104,7 +109,7 @@ export function useChat< const transport = options.connection ? { connection: options.connection } : { fetcher: options.fetcher } - return new ChatClient({ + return new ChatClient({ devtoolsBridgeFactory: createChatDevtoolsBridge, ...transport, ...(options.initialMessages !== undefined && { @@ -139,6 +144,9 @@ export function useChat< options.onError?.(err) }, tools: options.tools, + ...(options.interrupts !== undefined && { + interrupts: options.interrupts, + }), onCustomEvent: (eventType, data, context) => options.onCustomEvent?.(eventType, data, context), ...(options.streamProcessor !== undefined && { @@ -323,7 +331,11 @@ export function useChat< } const resolveInterrupts = ( - resolution: boolean | ((interrupt: ChatInterrupt) => undefined), + resolution: + | boolean + | (( + interrupt: ResolvableChatInterrupt, + ) => undefined), ) => { if (typeof resolution === 'boolean') { client().resolveInterrupts(resolution) @@ -421,5 +433,5 @@ export function useChat< resumeInterrupts, partial, final, - } as unknown as UseChatReturn + } as unknown as UseChatReturn } diff --git a/packages/ai-solid/tests/use-chat-types.test.ts b/packages/ai-solid/tests/use-chat-types.test.ts index 79e25700c4..3ff3cb2ab6 100644 --- a/packages/ai-solid/tests/use-chat-types.test.ts +++ b/packages/ai-solid/tests/use-chat-types.test.ts @@ -4,7 +4,7 @@ */ import { describe, expectTypeOf, it } from 'vitest' -import { toolDefinition } from '@tanstack/ai' +import { defineInterrupt, toolDefinition } from '@tanstack/ai' import { clientTools } from '@tanstack/ai-client' import { useChat } from '../src/use-chat' import type { AnyClientTool } from '@tanstack/ai' @@ -22,6 +22,16 @@ type TestSchema = { readonly vendor: 'test' readonly types: { readonly input: T; readonly output: T } readonly validate: (value: unknown) => { readonly value: T } + readonly jsonSchema: { readonly input: () => Record } + } +} +type TransformSchema = { + readonly '~standard': { + readonly version: 1 + readonly vendor: 'test' + readonly types: { readonly input: string; readonly output: number } + readonly validate: (value: unknown) => { readonly value: number } + readonly jsonSchema: { readonly input: () => Record } } } @@ -200,6 +210,7 @@ describe('useChat() interrupt types', () => { output: { accountId: '' }, }, validate: () => ({ value: { accountId: '' } }), + jsonSchema: { input: () => ({ type: 'object' }) }, }, } const transfer = toolDefinition({ @@ -271,3 +282,104 @@ describe('useChat() interrupt types', () => { void check }) }) + +describe('useChat() registered generic interrupt types', () => { + it('keeps registered and external generic interrupts distinct', () => { + const payloadSchema: TestSchema<{ title: string }> = { + '~standard': { + version: 1, + vendor: 'test', + types: { input: { title: '' }, output: { title: '' } }, + validate: () => ({ value: { title: '' } }), + jsonSchema: { input: () => ({ type: 'object' }) }, + }, + } + const responseSchema: TransformSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { input: '', output: 0 }, + validate: () => ({ value: 0 }), + jsonSchema: { input: () => ({ type: 'string' }) }, + }, + } + const reviewPlan = defineInterrupt({ + id: 'review-plan', + payloadSchema, + responseSchema, + }) + const acknowledge = defineInterrupt({ + id: 'acknowledge', + responseSchema: payloadSchema, + }) + + const check = () => { + const chat = useChat({ + connection: { connect: async function* () {} }, + interrupts: [reviewPlan, acknowledge], + }) + type Interrupt = ReturnType[number] + type Review = Extract + type External = Extract< + Exclude, + { kind: 'generic' } + > + type Unbound = Extract + type CallbackInterrupt = typeof chat.resolveInterrupts extends { + (resolver: (interrupt: infer TInterrupt) => undefined): void + } + ? TInterrupt + : never + + expectTypeOf().toEqualTypeOf< + { title: string } | undefined + >() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + expectTypeOf().toEqualTypeOf() + expectTypeOf< + Extract + >().toEqualTypeOf() + + const resolveReview = (review: Review) => { + review.resolveInterrupt('42') + // @ts-expect-error The transformed response still accepts its input type. + review.resolveInterrupt(42) + } + void resolveReview + chat.resolveInterrupts((interrupt) => { + interrupt.cancel() + return undefined + }) + + const existingTool = toolDefinition({ + name: 'solid-unregistered-tool', + description: 'A tool without an interrupt registry', + needsApproval: true, + }).client() + const withoutRegistry = useChat({ + connection: { connect: async function* () {} }, + tools: clientTools(existingTool), + }) + type WithoutRegistry = ReturnType< + typeof withoutRegistry.interrupts + >[number] + type ExistingToolInterrupt = Extract< + WithoutRegistry, + { kind: 'tool-approval' } + > + type UnregisteredGeneric = Extract + expectTypeOf< + ExistingToolInterrupt['toolName'] + >().toEqualTypeOf<'solid-unregistered-tool'>() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + } + void check + }) +}) diff --git a/packages/ai-svelte/src/create-chat.svelte.ts b/packages/ai-svelte/src/create-chat.svelte.ts index 873010503f..db361f14df 100644 --- a/packages/ai-svelte/src/create-chat.svelte.ts +++ b/packages/ai-svelte/src/create-chat.svelte.ts @@ -3,7 +3,7 @@ import { createChatDevtoolsBridge } from '@tanstack/ai-client/devtools' import { onMount } from 'svelte' import type { ChatClientState, - ChatInterrupt, + ResolvableChatInterrupt, ChatInterruptState, ChatResumeState, ConnectionStatus, @@ -14,6 +14,7 @@ import type { } from '@tanstack/ai-client' import type { AnyClientTool, + InterruptDefinition, InferSchemaType, ModelMessage, RunAgentResumeItem, @@ -65,9 +66,12 @@ export function createChat< const TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, TContext = InferredClientContext, + const TInterrupts extends ReadonlyArray< + InterruptDefinition + > = readonly [], >( - options: CreateChatOptions, -): CreateChatReturn { + options: CreateChatOptions, +): CreateChatReturn { // Create reactive state using Svelte 5 runes let messages = $state>>(options.initialMessages || []) let isLoading = $state(false) @@ -78,7 +82,7 @@ export function createChat< let sessionGenerating = $state(false) let queue = $state>([]) let runId = $state(null) - let interruptState = $state.raw>({ + let interruptState = $state.raw>({ interrupts: EMPTY_INTERRUPTS, pendingInterrupts: EMPTY_INTERRUPTS, interruptErrors: EMPTY_INTERRUPT_ERRORS, @@ -109,7 +113,7 @@ export function createChat< // The hook's identity is its `threadId`, which ChatClient also uses as the // persistence key — no separate `id`. When no `threadId` is given the client // generates one, so an ephemeral chat still works but is not restored on reload. - const client = new ChatClient({ + const client = new ChatClient({ devtoolsBridgeFactory: createChatDevtoolsBridge, ...transport, ...(options.initialMessages !== undefined && { @@ -144,6 +148,9 @@ export function createChat< options.onError?.(err) }, tools: options.tools, + ...(options.interrupts !== undefined && { + interrupts: options.interrupts, + }), ...(options.onCustomEvent !== undefined && { onCustomEvent: options.onCustomEvent, }), @@ -304,7 +311,11 @@ export function createChat< } const resolveInterrupts = ( - resolution: boolean | ((interrupt: ChatInterrupt) => undefined), + resolution: + | boolean + | (( + interrupt: ResolvableChatInterrupt, + ) => undefined), ) => { if (typeof resolution === 'boolean') { client.resolveInterrupts(resolution) @@ -447,5 +458,5 @@ export function createChat< updateBody, updateForwardedProps, updateContext, - } as unknown as CreateChatReturn + } as unknown as CreateChatReturn } diff --git a/packages/ai-svelte/src/types.ts b/packages/ai-svelte/src/types.ts index 124a4f0b0e..7e76002fb6 100644 --- a/packages/ai-svelte/src/types.ts +++ b/packages/ai-svelte/src/types.ts @@ -1,5 +1,6 @@ import type { AnyClientTool, + InterruptDefinition, InferSchemaType, ModelMessage, RunAgentResumeItem, @@ -10,7 +11,7 @@ import type { BoundInterrupts, ChatClientOptions, ChatClientState, - ChatInterrupt, + ResolvableChatInterrupt, ChatInterruptState, ChatRequestBody, ChatResumeState, @@ -78,8 +79,10 @@ export type CreateChatOptions< TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, TContext = InferredClientContext, + TInterrupts extends ReadonlyArray> = + readonly [], > = DistributedOmit< - ChatClientOptions, + ChatClientOptions, | 'onMessagesChange' | 'onLoadingChange' | 'onErrorChange' @@ -116,10 +119,13 @@ export type CreateChatReturn< TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, TContext = unknown, + TInterrupts extends ReadonlyArray> = + readonly [], > = BaseCreateChatReturn< TTools, TSchema extends SchemaInput ? InferSchemaType : unknown, - TContext + TContext, + TInterrupts > & (TSchema extends SchemaInput ? { @@ -141,6 +147,8 @@ interface BaseCreateChatReturn< TTools extends ReadonlyArray = any, TData = unknown, TContext = unknown, + TInterrupts extends ReadonlyArray> = + readonly [], > { /** * Current messages in the conversation (reactive getter). When @@ -205,14 +213,21 @@ interface BaseCreateChatReturn< * it, correlate a log line). */ readonly runId: string | null - readonly interrupts: BoundInterrupts + readonly interrupts: BoundInterrupts /** @deprecated Use `interrupts`. */ - readonly pendingInterrupts: BoundInterrupts - readonly interruptErrors: ChatInterruptState['interruptErrors'] + readonly pendingInterrupts: BoundInterrupts + readonly interruptErrors: ChatInterruptState< + TTools, + TInterrupts + >['interruptErrors'] readonly resuming: boolean resolveInterrupts: { (approved: boolean): void - (resolver: (interrupt: ChatInterrupt) => undefined): void + ( + resolver: ( + interrupt: ResolvableChatInterrupt, + ) => undefined, + ): void } cancelInterrupts: () => void retryInterrupts: () => void diff --git a/packages/ai-svelte/tests/create-chat-types.test.ts b/packages/ai-svelte/tests/create-chat-types.test.ts index 697a060180..41d906a9ef 100644 --- a/packages/ai-svelte/tests/create-chat-types.test.ts +++ b/packages/ai-svelte/tests/create-chat-types.test.ts @@ -4,15 +4,12 @@ */ import { describe, expectTypeOf, it } from 'vitest' -import { toolDefinition } from '@tanstack/ai' +import { defineInterrupt, toolDefinition } from '@tanstack/ai' import { clientTools } from '@tanstack/ai-client' import { createChat } from '../src/create-chat.svelte' import type { AnyClientTool } from '@tanstack/ai' import type { StructuredOutputPart } from '@tanstack/ai-client' -import type { - StandardJSONSchemaV1, - StandardSchemaV1, -} from '@standard-schema/spec' +import type { StandardJSONSchemaV1 } from '@standard-schema/spec' import type { CreateChatOptions, CreateChatReturn, @@ -22,6 +19,24 @@ import type { 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 } + readonly jsonSchema: { readonly input: () => Record } + } +} +type TransformSchema = { + readonly '~standard': { + readonly version: 1 + readonly vendor: 'test' + readonly types: { readonly input: string; readonly output: number } + readonly validate: (value: unknown) => { readonly value: number } + readonly jsonSchema: { readonly input: () => Record } + } +} describe('createChat() return type (svelte)', () => { describe('with outputSchema', () => { @@ -185,10 +200,7 @@ describe('createChat() interrupt types', () => { validate: () => ({ value: { reason: '' } }), }, } - const outputSchema: StandardSchemaV1< - { accountId: string }, - { accountId: string } - > = { + const outputSchema: TestSchema<{ accountId: string }> = { '~standard': { version: 1 as const, vendor: 'test', @@ -197,6 +209,7 @@ describe('createChat() interrupt types', () => { output: { accountId: '' }, }, validate: () => ({ value: { accountId: '' } }), + jsonSchema: { input: () => ({ type: 'object' }) }, }, } const transfer = toolDefinition({ @@ -266,3 +279,102 @@ describe('createChat() interrupt types', () => { void check }) }) + +describe('createChat() registered generic interrupt types', () => { + it('keeps registered and external generic interrupts distinct', () => { + const payloadSchema: TestSchema<{ title: string }> = { + '~standard': { + version: 1, + vendor: 'test', + types: { input: { title: '' }, output: { title: '' } }, + validate: () => ({ value: { title: '' } }), + jsonSchema: { input: () => ({ type: 'object' }) }, + }, + } + const responseSchema: TransformSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { input: '', output: 0 }, + validate: () => ({ value: 0 }), + jsonSchema: { input: () => ({ type: 'string' }) }, + }, + } + const reviewPlan = defineInterrupt({ + id: 'review-plan', + payloadSchema, + responseSchema, + }) + const acknowledge = defineInterrupt({ + id: 'acknowledge', + responseSchema: payloadSchema, + }) + + const check = () => { + const chat = createChat({ + connection: { connect: async function* () {} }, + interrupts: [reviewPlan, acknowledge], + }) + type Interrupt = (typeof chat.interrupts)[number] + type Review = Extract + type External = Extract< + Exclude, + { kind: 'generic' } + > + type Unbound = Extract + type CallbackInterrupt = typeof chat.resolveInterrupts extends { + (resolver: (interrupt: infer TInterrupt) => undefined): void + } + ? TInterrupt + : never + + expectTypeOf().toEqualTypeOf< + { title: string } | undefined + >() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + expectTypeOf().toEqualTypeOf() + expectTypeOf< + Extract + >().toEqualTypeOf() + + const resolveReview = (review: Review) => { + review.resolveInterrupt('42') + // @ts-expect-error The transformed response still accepts its input type. + review.resolveInterrupt(42) + } + void resolveReview + chat.resolveInterrupts((interrupt) => { + interrupt.cancel() + return undefined + }) + + const existingTool = toolDefinition({ + name: 'svelte-unregistered-tool', + description: 'A tool without an interrupt registry', + needsApproval: true, + }).client() + const withoutRegistry = createChat({ + connection: { connect: async function* () {} }, + tools: clientTools(existingTool), + }) + type WithoutRegistry = (typeof withoutRegistry.interrupts)[number] + type ExistingToolInterrupt = Extract< + WithoutRegistry, + { kind: 'tool-approval' } + > + type UnregisteredGeneric = Extract + expectTypeOf< + ExistingToolInterrupt['toolName'] + >().toEqualTypeOf<'svelte-unregistered-tool'>() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + } + void check + }) +}) diff --git a/packages/ai-vue/src/types.ts b/packages/ai-vue/src/types.ts index e3c4e9b8b7..35a4ee29fc 100644 --- a/packages/ai-vue/src/types.ts +++ b/packages/ai-vue/src/types.ts @@ -1,5 +1,6 @@ import type { AnyClientTool, + InterruptDefinition, InferSchemaType, ModelMessage, RunAgentResumeItem, @@ -10,7 +11,7 @@ import type { BoundInterrupts, ChatClientOptions, ChatClientState, - ChatInterrupt, + ResolvableChatInterrupt, ChatInterruptState, ChatRequestBody, ChatResumeState, @@ -80,8 +81,10 @@ export type UseChatOptions< TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, TContext = InferredClientContext, + TInterrupts extends ReadonlyArray> = + readonly [], > = DistributedOmit< - ChatClientOptions, + ChatClientOptions, | 'onMessagesChange' | 'onLoadingChange' | 'onErrorChange' @@ -117,9 +120,12 @@ export type UseChatOptions< export type UseChatReturn< TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, + TInterrupts extends ReadonlyArray> = + readonly [], > = BaseUseChatReturn< TTools, - TSchema extends SchemaInput ? InferSchemaType : unknown + TSchema extends SchemaInput ? InferSchemaType : unknown, + TInterrupts > & (TSchema extends SchemaInput ? { @@ -140,6 +146,8 @@ export type UseChatReturn< interface BaseUseChatReturn< TTools extends ReadonlyArray = any, TData = unknown, + TInterrupts extends ReadonlyArray> = + readonly [], > { /** * Current messages in the conversation. When `outputSchema` is supplied, @@ -203,16 +211,20 @@ interface BaseUseChatReturn< * it, correlate a log line). */ runId: DeepReadonly> - interrupts: DeepReadonly>> + interrupts: Readonly>> /** @deprecated Use `interrupts`. */ - pendingInterrupts: DeepReadonly>> + pendingInterrupts: Readonly>> interruptErrors: DeepReadonly< - ShallowRef['interruptErrors']> + ShallowRef['interruptErrors']> > resuming: DeepReadonly> resolveInterrupts: { (approved: boolean): void - (resolver: (interrupt: ChatInterrupt) => undefined): void + ( + resolver: ( + interrupt: ResolvableChatInterrupt, + ) => undefined, + ): void } cancelInterrupts: () => void retryInterrupts: () => void diff --git a/packages/ai-vue/src/use-chat.ts b/packages/ai-vue/src/use-chat.ts index fe5c9e4514..b76cbbbde7 100644 --- a/packages/ai-vue/src/use-chat.ts +++ b/packages/ai-vue/src/use-chat.ts @@ -10,6 +10,7 @@ import { } from 'vue' import type { AnyClientTool, + InterruptDefinition, InferSchemaType, ModelMessage, RunAgentResumeItem, @@ -18,7 +19,7 @@ import type { } from '@tanstack/ai' import type { ChatClientState, - ChatInterrupt, + ResolvableChatInterrupt, ChatInterruptState, ChatResumeState, ConnectionStatus, @@ -42,13 +43,17 @@ export function useChat< const TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, TContext = InferredClientContext, + const TInterrupts extends ReadonlyArray< + InterruptDefinition + > = readonly [], >( - options: UseChatOptions = {} as UseChatOptions< + options: UseChatOptions< TTools, TSchema, - TContext - >, -): UseChatReturn { + TContext, + TInterrupts + > = {} as UseChatOptions, +): UseChatReturn { const messages = shallowRef>>( options.initialMessages || [], ) @@ -60,7 +65,7 @@ export function useChat< const sessionGenerating = shallowRef(false) const queue = shallowRef>([]) const runId = shallowRef(null) - const interruptState = shallowRef>({ + const interruptState = shallowRef>({ interrupts: EMPTY_INTERRUPTS, pendingInterrupts: EMPTY_INTERRUPTS, interruptErrors: EMPTY_INTERRUPT_ERRORS, @@ -95,7 +100,7 @@ export function useChat< // The hook's identity is its `threadId`, which ChatClient also uses as the // persistence key — no separate `id`. When no `threadId` is given the client // generates one, so an ephemeral chat still works but is not restored on reload. - const client = new ChatClient({ + const client = new ChatClient({ devtoolsBridgeFactory: createChatDevtoolsBridge, ...transport, ...(options.initialMessages !== undefined && { @@ -130,6 +135,9 @@ export function useChat< options.onError?.(err) }, tools: options.tools, + ...(options.interrupts !== undefined && { + interrupts: options.interrupts, + }), onCustomEvent: (eventType, data, context) => options.onCustomEvent?.(eventType, data, context), ...(options.streamProcessor !== undefined && { @@ -320,7 +328,11 @@ export function useChat< const resuming = computed(() => interruptState.value.resuming) const resolveInterrupts = ( - resolution: boolean | ((interrupt: ChatInterrupt) => undefined), + resolution: + | boolean + | (( + interrupt: ResolvableChatInterrupt, + ) => undefined), ) => { if (typeof resolution === 'boolean') { client.resolveInterrupts(resolution) @@ -417,5 +429,5 @@ export function useChat< resumeInterrupts, partial: readonly(partial), final: readonly(final), - } as unknown as UseChatReturn + } as unknown as UseChatReturn } diff --git a/packages/ai-vue/tests/use-chat-types.test.ts b/packages/ai-vue/tests/use-chat-types.test.ts index ce3eb0c54a..4eee1f087d 100644 --- a/packages/ai-vue/tests/use-chat-types.test.ts +++ b/packages/ai-vue/tests/use-chat-types.test.ts @@ -4,7 +4,7 @@ */ import { describe, expectTypeOf, it } from 'vitest' -import { toolDefinition } from '@tanstack/ai' +import { defineInterrupt, toolDefinition } from '@tanstack/ai' import { clientTools } from '@tanstack/ai-client' import { useChat } from '../src/use-chat' import type { AnyClientTool } from '@tanstack/ai' @@ -16,6 +16,26 @@ import type { } from '@standard-schema/spec' import type { DeepPartial, UseChatOptions, UseChatReturn } from '../src/types' +type TransformSchema = { + readonly '~standard': { + readonly version: 1 + readonly vendor: 'test' + readonly types: { readonly input: string; readonly output: number } + readonly validate: (value: unknown) => { readonly value: number } + readonly jsonSchema: { readonly input: () => Record } + } +} + +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 } + readonly jsonSchema: { readonly input: () => Record } + } +} + type Person = { name: string; age: number; email: string } type PersonSchema = StandardJSONSchemaV1 type NoTools = ReadonlyArray @@ -261,3 +281,104 @@ describe('useChat() interrupt types', () => { void check }) }) + +describe('useChat() registered generic interrupt types', () => { + it('keeps registered and external generic interrupts distinct', () => { + const payloadSchema: TestSchema<{ title: string }> = { + '~standard': { + version: 1, + vendor: 'test', + types: { input: { title: '' }, output: { title: '' } }, + validate: () => ({ value: { title: '' } }), + jsonSchema: { input: () => ({ type: 'object' }) }, + }, + } + const responseSchema: TransformSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { input: '', output: 0 }, + validate: () => ({ value: 0 }), + jsonSchema: { input: () => ({ type: 'string' }) }, + }, + } + const reviewPlan = defineInterrupt({ + id: 'review-plan', + payloadSchema, + responseSchema, + }) + const acknowledge = defineInterrupt({ + id: 'acknowledge', + responseSchema: payloadSchema, + }) + + const check = () => { + const chat = useChat({ + connection: { connect: async function* () {} }, + interrupts: [reviewPlan, acknowledge], + }) + type Interrupt = (typeof chat.interrupts.value)[number] + type Review = Extract + type External = Extract< + Exclude, + { kind: 'generic' } + > + type Unbound = Extract + type CallbackInterrupt = typeof chat.resolveInterrupts extends { + (resolver: (interrupt: infer TInterrupt) => undefined): void + } + ? TInterrupt + : never + + expectTypeOf().toEqualTypeOf< + { title: string } | undefined + >() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + expectTypeOf().toEqualTypeOf() + expectTypeOf< + Extract + >().toEqualTypeOf() + + const resolveReview = (review: Review) => { + review.resolveInterrupt('42') + // @ts-expect-error The transformed response still accepts its input type. + review.resolveInterrupt(42) + } + void resolveReview + chat.resolveInterrupts((interrupt) => { + interrupt.cancel() + return undefined + }) + // @ts-expect-error The public ref is readonly. + chat.interrupts.value = [] + + const existingTool = toolDefinition({ + name: 'vue-unregistered-tool', + description: 'A tool without an interrupt registry', + needsApproval: true, + }).client() + const withoutRegistry = useChat({ + connection: { connect: async function* () {} }, + tools: clientTools(existingTool), + }) + type WithoutRegistry = (typeof withoutRegistry.interrupts.value)[number] + type ExistingToolInterrupt = Extract< + WithoutRegistry, + { kind: 'tool-approval' } + > + type UnregisteredGeneric = Extract + expectTypeOf< + ExistingToolInterrupt['toolName'] + >().toEqualTypeOf<'vue-unregistered-tool'>() + expectTypeOf< + Parameters[0] + >().toEqualTypeOf() + } + void check + }) +}) diff --git a/packages/ai/skills/ai-core/tool-calling/SKILL.md b/packages/ai/skills/ai-core/tool-calling/SKILL.md index 67316ef08f..c482172381 100644 --- a/packages/ai/skills/ai-core/tool-calling/SKILL.md +++ b/packages/ai/skills/ai-core/tool-calling/SKILL.md @@ -4,7 +4,8 @@ description: > Isomorphic tool system: toolDefinition() with Zod schemas, .server() and .client() implementations, passing tools to both chat() on server and useChat/clientTools on client, tool approval - flows with needsApproval and bound interrupts (resolveInterrupt), lazy tool + flows with needsApproval and bound interrupts (resolveInterrupt), generic + middleware interrupts with defineInterrupt(), lazy tool discovery with lazy:true, rendering ToolCallPart and ToolResultPart in UI. type: sub-skill @@ -133,6 +134,58 @@ function ChatPage() { ## Core Patterns +### Generic middleware interrupts + +Use `defineInterrupt()` when middleware needs typed data from the client. This +does not replace `needsApproval`. Tool approval asks whether a tool can run. +Generic interrupts ask for application data at a chat lifecycle boundary. + +Define the interrupt once. Register it with both `chat({ interrupts })` and +`useChat({ interrupts })`. Emit it only from `onInterruptBoundary`, then read +the typed result in `onInterruptResolution`. + +```typescript +import { defineInterrupt, type ChatMiddleware } from '@tanstack/ai' +import { z } from 'zod' + +const reviewPlan = defineInterrupt({ + id: 'review-plan', + payloadSchema: z.object({ title: z.string() }), + responseSchema: z.object({ approved: z.boolean() }), +}) + +const reviewMiddleware: ChatMiddleware = { + onInterruptBoundary(ctx) { + if (ctx.phase !== 'beforeTools') return + return { + interrupts: [ + reviewPlan.interrupt({ + key: 'release-plan', + reason: 'review-required', + message: 'Approve this plan?', + payload: { title: 'Release plan' }, + }), + ], + } + }, + onInterruptResolution(_ctx, resumedInterrupts) { + for (const result of resumedInterrupts.for(reviewPlan)) { + if (result.status === 'resolved' && !result.response.approved) { + return { toolResume: 'stop' } + } + } + }, +} +``` + +Several middleware can request generic interrupts at one boundary. They share +one AG-UI interrupt batch with tool approvals. A continuation starts only after +the client resolves or cancels every bound item. `stop` is more restrictive than +`cancel`, which is more restrictive than `continue`. + +Do not emit raw AG-UI interrupt events from middleware. Use the boundary hook +so the engine creates one terminal event and persistence records the batch. + ### Pattern 1: Server-Only Tool Define with `toolDefinition()`, implement with `.server()`, pass to `chat({ tools })`. diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 15fb215ba0..019ae5512b 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -18,6 +18,15 @@ import { validateInterruptResumeBatch, } from '../../interrupt-resume' import { INTERRUPT_BINDING_VERSION } from '../../interrupts' +import { + INTERRUPT_PAYLOAD_METADATA_KEY, + createInterruptBinding, +} from '../../interrupt-definition' +import { readGenericInterruptContinuation } from '../../generic-interrupt-continuation' +import type { + GenericInterruptRequest, + InterruptDefinition, +} from '../../interrupt-definition' import { canonicalInterruptJson, digestInterruptJson, @@ -99,10 +108,13 @@ import type { ChatMiddleware, ChatMiddlewareConfig, ChatMiddlewareContext, + ChatResumeGenericResolution, ChatResumeToolState, + InterruptResolutionCollection, SandboxFileHookEvent, StructuredOutputMiddlewareConfig, } from './middleware/types' +import { provideGenericInterruptDefinitionRegistry } from './middleware/generic-interrupts' import type { CheckCoverage } from './middleware/builder' import type { SystemPrompt } from '../../system-prompts' import type { InternalLogger } from '../../logger/internal-logger' @@ -214,7 +226,6 @@ function normalizePublicInterruptBinding( typeof binding.generation !== 'number' || !Number.isInteger(binding.generation) || binding.generation < 0 || - typeof binding.responseSchemaHash !== 'string' || (binding.expiresAt !== undefined && typeof binding.expiresAt !== 'string') ) { return undefined @@ -224,15 +235,58 @@ function normalizePublicInterruptBinding( interruptId: binding.interruptId, interruptedRunId: binding.interruptedRunId, generation: binding.generation, - responseSchemaHash: binding.responseSchemaHash, + ...(typeof binding.responseSchemaHash === 'string' + ? { responseSchemaHash: binding.responseSchemaHash } + : {}), ...(typeof binding.expiresAt === 'string' ? { expiresAt: binding.expiresAt } : {}), } if (binding.kind === 'generic') { + if ( + binding.responseSchemaHash !== undefined && + typeof binding.responseSchemaHash !== 'string' + ) { + return undefined + } + const hasFirstPartyFields = [ + binding.definitionId, + binding.key, + binding.batchIndex, + binding.payloadSchemaHash, + ].some((field) => field !== undefined) + if ( + hasFirstPartyFields && + (typeof binding.definitionId !== 'string' || + typeof binding.key !== 'string' || + typeof binding.batchIndex !== 'number' || + !Number.isInteger(binding.batchIndex) || + binding.batchIndex < 0 || + (binding.payloadSchemaHash !== undefined && + typeof binding.payloadSchemaHash !== 'string')) + ) { + return undefined + } + if ( + typeof binding.definitionId === 'string' && + typeof binding.key === 'string' && + typeof binding.batchIndex === 'number' + ) { + return { + kind: binding.kind, + ...base, + definitionId: binding.definitionId, + key: binding.key, + batchIndex: binding.batchIndex, + ...(typeof binding.payloadSchemaHash === 'string' + ? { payloadSchemaHash: binding.payloadSchemaHash } + : {}), + } + } return { kind: binding.kind, ...base } } if ( + typeof binding.responseSchemaHash !== 'string' || typeof binding.toolName !== 'string' || typeof binding.toolCallId !== 'string' ) { @@ -245,6 +299,7 @@ function normalizePublicInterruptBinding( return { kind: binding.kind, ...base, + responseSchemaHash: binding.responseSchemaHash, toolName: binding.toolName, toolCallId: binding.toolCallId, outputSchemaHash: binding.outputSchemaHash, @@ -259,6 +314,7 @@ function normalizePublicInterruptBinding( return { kind: binding.kind, ...base, + responseSchemaHash: binding.responseSchemaHash, toolName: binding.toolName, toolCallId: binding.toolCallId, originalArgs: binding.originalArgs, @@ -305,33 +361,133 @@ type InferredContext = [ ? unknown : ContextFromInputs -type RequiredContextFromInputs = [ - ContextFromInputs, +type RegistryInterrupt< + TInterrupts extends ReadonlyArray>, +> = [TInterrupts[number]] extends [never] ? never : TInterrupts[number] + +type DuplicateInterruptDefinitionId< + TInterrupts extends ReadonlyArray>, + TSeenIds extends string = never, +> = TInterrupts extends readonly [infer THead, ...infer TTail] + ? THead extends InterruptDefinition + ? string extends TId + ? TTail extends ReadonlyArray> + ? DuplicateInterruptDefinitionId + : never + : TId extends TSeenIds + ? TId + : TTail extends ReadonlyArray> + ? DuplicateInterruptDefinitionId + : never + : never + : never + +type CheckUniqueInterruptDefinitions< + TInterrupts extends ReadonlyArray>, +> = [DuplicateInterruptDefinitionId] extends [never] + ? unknown + : { + readonly '✖ Duplicate interrupt definition id in chat({ interrupts }).': never + } + +type InlineChatContext = MergeContext< + ContextFromArray>, + TContext +> + +type RegistryChatMiddleware< + TContext, + TInterrupts extends ReadonlyArray>, +> = ChatMiddleware> + +type MiddlewareInterruptDefinitions = + TMiddleware extends ReadonlyArray + ? TMiddlewareItem extends ChatMiddleware + ? TDefinitions + : never + : never + +type IsAny = 0 extends 1 & TValue ? true : false + +type CheckInterruptRegistry< + TInterrupts extends ReadonlyArray>, + TMiddleware, +> = + IsAny> extends true + ? unknown + : [MiddlewareInterruptDefinitions] extends [never] + ? unknown + : [ + Exclude< + MiddlewareInterruptDefinitions, + RegistryInterrupt + >, + ] extends [never] + ? unknown + : { + readonly '✖ Middleware emits an interrupt definition that is not registered in chat({ interrupts }).': never + } + +type RuntimeContextOption = [ + MergeContext, TContext>, ] extends [never] - ? { context?: unknown } - : undefined extends ContextFromInputs - ? { context?: ContextFromInputs } - : { context: ContextFromInputs } + ? { context?: TContext } + : undefined extends MergeContext< + ContextFromInputs, + TContext + > + ? { + context?: MergeContext, TContext> + } + : { + context: MergeContext, TContext> + } + +type ExactMiddlewareOption< + TTools, + TContext, + TInterrupts extends ReadonlyArray>, + TMiddleware extends Array | undefined, +> = [TMiddleware] extends [undefined] + ? Array< + RegistryChatMiddleware< + InlineChatContext, + NoInfer + > + > + : TMiddleware & + (TMiddleware extends Array< + RegistryChatMiddleware< + InlineChatContext>, + NoInfer + > + > + ? Array< + RegistryChatMiddleware< + InlineChatContext, + NoInfer + > + > + : CheckInterruptRegistry) & + CheckCoverage>> type TextActivityOptionsWithContext< TAdapter extends AnyTextAdapter, TSchema extends SchemaInput | undefined, TStream extends boolean, TTools extends TextActivityOptions['tools'], - TMiddleware extends TextActivityOptions< - TAdapter, - TSchema, - TStream, - any - >['middleware'], + TInterrupts extends ReadonlyArray> = + [], + TContext = unknown, + TMiddleware extends Array | undefined = undefined, > = Omit< TextActivityOptions, - 'tools' | 'middleware' | 'context' + 'tools' | 'middleware' | 'context' | 'interrupts' > & { tools?: TTools - middleware?: TMiddleware & - CheckCoverage>> -} & RequiredContextFromInputs + interrupts?: TInterrupts & CheckUniqueInterruptDefinitions + middleware?: ExactMiddlewareOption +} & RuntimeContextOption // =========================== // Activity Options Type @@ -489,6 +645,11 @@ export interface TextActivityOptions< * ``` */ middleware?: Array> + /** + * First-party generic interrupt definitions for this chat call. + * Register the same definitions on the client to type payloads and answers. + */ + interrupts?: ReadonlyArray> /** * Runtime context value passed to middleware hooks and server tools. */ @@ -529,28 +690,21 @@ export function createChatOptions< TStream, any >['tools'] = TextActivityOptions['tools'], - const TMiddleware extends TextActivityOptions< - TAdapter, - TSchema, - TStream, - any - >['middleware'] = TextActivityOptions< - TAdapter, - TSchema, - TStream, - any - >['middleware'], + const TInterrupts extends ReadonlyArray< + InterruptDefinition + > = [], + TContext = unknown, + const TMiddleware extends Array | undefined = undefined, >( options: TextActivityOptionsWithContext< TAdapter, TSchema, TStream, TTools, + TInterrupts, + TContext, TMiddleware >, - // Preserve the concrete `tools` tuple on the returned options (so a later - // `chat({ ...opts })` still narrows tool-call events to the tool names) - // while threading the inferred runtime context like the bare options type. ): Omit< TextActivityOptions< TAdapter, @@ -558,8 +712,12 @@ export function createChatOptions< TStream, InferredContext >, - 'tools' -> & { tools?: TTools } { + 'tools' | 'middleware' | 'interrupts' +> & { + tools?: TTools + interrupts?: TInterrupts + middleware?: ExactMiddlewareOption +} { return options } @@ -623,7 +781,7 @@ interface TextEngineConfig< adapter: TAdapter systemPrompts?: Array params: TParams - middleware?: Array> + middleware?: Array context?: TContext /** * If set, after the agent loop finishes the engine runs a @@ -705,12 +863,18 @@ class TextEngine< >, > { private readonly adapter: TAdapter + private readonly interruptDefinitions: ReadonlyMap< + string, + InterruptDefinition + > private params: TParams private systemPrompts: Array private tools: Array private readonly loopStrategy: AgentLoopStrategy private toolCallManager: ToolCallManager, TContext> private readonly lazyToolManager: LazyToolManager + /** A public interruption terminal must always have this run's start event. */ + private hasPublicRunStarted = false private readonly initialMessageCount: number private readonly requestId: string private readonly streamId: string @@ -740,6 +904,9 @@ class TextEngine< private eventToolNames?: Array private finishedEvent: RunFinishedEvent | null = null private deferredToolCallRunFinishedChunks: Array = [] + /** The model terminal is held until afterModel can choose an interrupt. */ + private deferredModelRunFinishedChunks: Array = [] + private earlyTermination = false private toolPhase: ToolPhaseResult = 'continue' private cyclePhase: CyclePhase = 'processText' @@ -750,6 +917,14 @@ class TextEngine< private readonly resumeClientToolResults = new Map() private readonly resumeDeniedToolResults = new Map() private readonly resumeCancelledToolCallIds = new Set() + private readonly resumeGenericInterrupts = new Map< + string, + ChatResumeGenericResolution + >() + private readonly resumeGenericInterruptRequests = new Map< + string, + GenericInterruptRequest> + >() // AG-UI protocol IDs private readonly threadId: string @@ -757,7 +932,10 @@ class TextEngine< private readonly parentRunIdOverride?: string // Middleware support - private readonly middlewareRunner: MiddlewareRunner + private readonly middlewareRunner: MiddlewareRunner< + TContext, + InterruptDefinition + > private readonly middlewareCtx: ChatMiddlewareContext private readonly sandboxFileQueue: Array = [] private readonly deferredPromises: Array> = [] @@ -815,6 +993,15 @@ class TextEngine< ) { this.logger = logger this.adapter = config.adapter + this.interruptDefinitions = new Map( + ( + ( + config.params as TParams & { + interrupts?: ReadonlyArray> + } + ).interrupts ?? [] + ).map((definition) => [definition.id, definition]), + ) this.finalStructuredOutput = config.finalStructuredOutput this.params = config.params this.systemPrompts = config.params.systemPrompts || [] @@ -866,7 +1053,9 @@ class TextEngine< // handleStreamChunk processes raw chunks BEFORE middleware, so internal // state management sees extended fields (finishReason, delta, toolCallName, etc.). // The strip middleware ensures the yielded public stream is AG-UI spec-compliant. - const allMiddleware: Array> = [ + const allMiddleware: Array< + ChatMiddleware> + > = [ devtoolsMiddleware(), ...(config.middleware || []), stripToSpecMiddleware(), @@ -945,6 +1134,10 @@ class TextEngine< }, }) + provideGenericInterruptDefinitionRegistry(this.middlewareCtx, { + definitions: this.interruptDefinitions, + }) + // Provide the internal SandboxRuntime capability so harness adapters and // sandbox middleware can emit file events. The sink logs, fans the event // out through the middleware `onFile*` hooks (fire-and-forget), and queues @@ -1036,10 +1229,25 @@ class TextEngine< ) this.applyMiddlewareConfig(transformedConfig) await this.applyEphemeralInterruptResume(transformedConfig) + await this.applyDurableGenericInterruptResolution() // Run onStart (devtools middleware emits text:request:started and initial messages here) await this.middlewareRunner.runOnStart(this.middlewareCtx) + if (this.earlyTermination) { + yield* this.emitSuccessfulEarlyTermination() + if (!this.terminalHookCalled) { + this.terminalHookCalled = true + await this.middlewareRunner.runOnFinish(this.middlewareCtx, { + finishReason: this.lastFinishReason, + duration: Date.now() - this.streamStartTime, + content: this.accumulatedContent, + usage: this.finishedEvent?.usage, + }) + } + return + } + const pendingPhase = yield* this.checkForPendingToolCalls() if (pendingPhase === 'wait') { return @@ -1083,7 +1291,35 @@ class TextEngine< ) this.applyMiddlewareConfig(iterTransformedConfig) + if ( + yield* this.emitBoundaryInterrupts( + 'beforeModel', + this.createSyntheticFinishedEvent(), + ) + ) { + this.setToolPhase('wait') + return + } + yield* this.streamModelResponse() + + if ( + yield* this.emitBoundaryInterrupts( + 'afterModel', + this.finishedEvent ?? this.createSyntheticFinishedEvent(), + ) + ) { + this.setToolPhase('wait') + return + } + if (this.shouldExecuteToolPhase()) { + this.deferredToolCallRunFinishedChunks.push( + ...this.deferredModelRunFinishedChunks, + ) + this.deferredModelRunFinishedChunks = [] + } else { + yield* this.flushDeferredModelRunFinishedChunks() + } } else { yield* this.processToolCalls() } @@ -1429,10 +1665,17 @@ class TextEngine< ) { continue } + if (outputChunk.type === EventType.RUN_FINISHED) { + this.deferredModelRunFinishedChunks.push(outputChunk) + continue + } if (this.shouldDeferToolCallRunFinished(outputChunk)) { this.deferredToolCallRunFinishedChunks.push(outputChunk) continue } + if (outputChunk.type === EventType.RUN_STARTED) { + this.hasPublicRunStarted = true + } this.logger.output(`type=${outputChunk.type}`, { chunk: outputChunk }) yield outputChunk this.middlewareCtx.chunkIndex++ @@ -1664,6 +1907,18 @@ class TextEngine< return 'continue' } + this.middlewareCtx.phase = 'beforeTools' + if ( + yield* this.emitBoundaryInterrupts( + 'beforeTools', + finishEvent, + executablePendingCalls, + ) + ) { + this.setToolPhase('wait') + return 'wait' + } + const { approvals, clientToolResults } = this.collectClientState() const generator = executeToolCalls( @@ -1825,6 +2080,17 @@ class TextEngine< } this.middlewareCtx.phase = 'beforeTools' + if ( + yield* this.emitBoundaryInterrupts( + 'beforeTools', + finishEvent, + executableToolCalls, + ) + ) { + this.setToolPhase('wait') + return + } + const { approvals, clientToolResults } = this.collectClientState() const generator = executeToolCalls( @@ -1892,15 +2158,36 @@ class TextEngine< needsClientExecution: executionResult.needsClientExecution, }) + const afterToolBoundaryChunks = this.buildToolResultChunks( + allResults, + finishEvent, + ) + const afterToolRequests = + await this.middlewareRunner.runOnInterruptBoundary( + this.middlewareCtx as ChatMiddlewareContext & { + phase: 'afterTools' + }, + ) + if (afterToolRequests.length > 0) { + for (const chunk of afterToolBoundaryChunks) { + yield* this.pipeThroughMiddleware(chunk) + } + yield* this.emitBoundaryInterrupts( + 'afterTools', + finishEvent, + toolCalls, + afterToolRequests, + ) + this.setToolPhase('wait') + return + } + if ( executionResult.needsApproval.length > 0 || executionResult.needsClientExecution.length > 0 ) { if (allResults.length > 0) { - for (const chunk of this.buildToolResultChunks( - allResults, - finishEvent, - )) { + for (const chunk of afterToolBoundaryChunks) { yield* this.pipeThroughMiddleware(chunk) } } @@ -1916,7 +2203,7 @@ class TextEngine< yield* this.flushDeferredToolCallRunFinishedChunks() - const toolResultChunks = this.buildToolResultChunks(allResults, finishEvent) + const toolResultChunks = afterToolBoundaryChunks for (const chunk of toolResultChunks) { yield* this.pipeThroughMiddleware(chunk) @@ -1956,6 +2243,48 @@ class TextEngine< this.deferredToolCallRunFinishedChunks = [] } + private *flushDeferredModelRunFinishedChunks(): Generator { + for (const chunk of this.deferredModelRunFinishedChunks) { + this.logger.output(`type=${chunk.type}`, { chunk }) + yield chunk + this.middlewareCtx.chunkIndex++ + } + this.deferredModelRunFinishedChunks = [] + } + + private async *emitSyntheticRunStarted( + finishEvent: RunFinishedEvent, + ): AsyncGenerator { + if (this.hasPublicRunStarted) return + yield* this.pipeThroughMiddleware({ + type: EventType.RUN_STARTED, + runId: finishEvent.runId, + threadId: finishEvent.threadId, + model: finishEvent.model, + timestamp: Date.now(), + }) + } + + private async *emitSuccessfulEarlyTermination(): AsyncGenerator< + StreamChunk, + void, + void + > { + // `stop` is a finished run, not another tool cycle. `tool_calls` here + // makes the client auto-send after afterTools, so reject looks stuck. + this.lastFinishReason = 'stop' + const finishEvent = { + ...this.createSyntheticFinishedEvent(), + finishReason: 'stop' as const, + } + yield* this.emitSyntheticRunStarted(finishEvent) + yield* this.pipeThroughMiddleware({ + ...finishEvent, + timestamp: Date.now(), + outcome: { type: 'success' }, + }) + } + private discardDeferredToolCallRunFinishedChunks(): void { this.deferredToolCallRunFinishedChunks = [] } @@ -2077,9 +2406,17 @@ class TextEngine< return { approvals, clientToolResults } } + private genericInterruptId(): string { + return this.createId('interrupt') + } + private buildActionableInterrupts( approvals: Array, clientRequests: Array, + genericRequests: ReadonlyArray< + GenericInterruptRequest> + > = [], + genericInterruptIds: ReadonlyArray = [], ): Array { const interrupts: Array = [] @@ -2149,6 +2486,64 @@ class TextEngine< }) } + for (const [index, request] of genericRequests.entries()) { + const batchIndex = interrupts.length + const id = genericInterruptIds[index] + if (!id) throw new Error('Generic interrupt id is unavailable.') + const preEmission = createInterruptBinding(request, { batchIndex }) + interrupts.push({ + id, + reason: request.reason, + message: request.message, + ...(preEmission.descriptor.responseSchemaCanonicalJson !== undefined + ? { + responseSchema: JSON.parse( + preEmission.descriptor.responseSchemaCanonicalJson, + ), + } + : {}), + ...(request.expiresAt !== undefined + ? { expiresAt: request.expiresAt } + : {}), + metadata: { + [interruptBindingMetadataKey]: { + v: INTERRUPT_BINDING_VERSION, + kind: 'generic', + interruptId: id, + definitionId: preEmission.descriptor.definitionId, + key: preEmission.descriptor.key, + batchIndex, + ...(request.expiresAt !== undefined + ? { expiresAt: request.expiresAt } + : {}), + ...(preEmission.descriptor.payloadSchemaHash + ? { + payloadSchemaHash: preEmission.descriptor.payloadSchemaHash, + } + : {}), + ...(preEmission.descriptor.responseSchemaHash !== undefined + ? { + responseSchemaHash: preEmission.descriptor.responseSchemaHash, + } + : {}), + }, + ...(preEmission.payload !== undefined + ? { [INTERRUPT_PAYLOAD_METADATA_KEY]: preEmission.payload } + : {}), + }, + }) + } + + const ids = new Set() + for (const interrupt of interrupts) { + if (ids.has(interrupt.id)) { + throw new Error( + `Duplicate interrupt id in final batch: ${interrupt.id}`, + ) + } + ids.add(interrupt.id) + } + return interrupts } @@ -2156,13 +2551,22 @@ class TextEngine< finishEvent: RunFinishedEvent, approvals: Array, clientRequests: Array, + genericRequests: ReadonlyArray< + GenericInterruptRequest> + > = [], + genericInterruptIds?: ReadonlyArray, ): StreamChunk { return { ...finishEvent, timestamp: Date.now(), outcome: { type: 'interrupt', - interrupts: this.buildActionableInterrupts(approvals, clientRequests), + interrupts: this.buildActionableInterrupts( + approvals, + clientRequests, + genericRequests, + genericInterruptIds, + ), }, } } @@ -2303,9 +2707,22 @@ class TextEngine< finishEvent: RunFinishedEvent, approvals: Array, clientRequests: Array, + genericRequests: ReadonlyArray< + GenericInterruptRequest> + > = [], ): AsyncGenerator { + yield* this.emitSyntheticRunStarted(finishEvent) + const genericInterruptIds = genericRequests.map(() => + this.genericInterruptId(), + ) const terminal = this.completeEphemeralInterruptBindings( - this.buildInterruptFinishedChunk(finishEvent, approvals, clientRequests), + this.buildInterruptFinishedChunk( + finishEvent, + approvals, + clientRequests, + genericRequests, + genericInterruptIds, + ), ) let terminalOutputs: Array try { @@ -2334,6 +2751,103 @@ class TextEngine< return true } + private async *emitBoundaryInterrupts( + phase: 'beforeModel' | 'afterModel' | 'beforeTools' | 'afterTools', + finishEvent: RunFinishedEvent, + toolCalls: ReadonlyArray = [], + requests?: ReadonlyArray< + GenericInterruptRequest> + >, + ): AsyncGenerator { + this.middlewareCtx.phase = phase + const boundaryRequests = + requests ?? + (await this.middlewareRunner.runOnInterruptBoundary( + this.middlewareCtx as ChatMiddlewareContext & { + phase: typeof phase + }, + )) + if (boundaryRequests.length === 0) return false + for (const request of boundaryRequests) { + if ( + this.interruptDefinitions.get(request.definition.id) !== + request.definition + ) { + throw new Error( + `Generic interrupt definition ${request.definition.id} is not registered on this chat.`, + ) + } + } + if (phase === 'afterModel') { + if (this.toolCallManager.hasToolCalls()) { + this.addAssistantToolCallMessage(this.toolCallManager.getToolCalls()) + } else { + this.addAssistantTextMessageForInterrupt() + } + } + const actionable = this.getBoundaryActionableToolRequests(toolCalls) + yield* this.emitActionableInterruptBoundary( + finishEvent, + actionable.approvals, + actionable.clientRequests, + boundaryRequests, + ) + return true + } + + private addAssistantTextMessageForInterrupt(): void { + if (this.accumulatedContent.length === 0) return + this.messages = [ + ...this.messages, + { role: 'assistant', content: this.accumulatedContent }, + ] + this.middlewareCtx.messages = this.messages + } + + private getBoundaryActionableToolRequests( + toolCalls: ReadonlyArray, + ): { + approvals: Array + clientRequests: Array + } { + const { approvals, clientToolResults } = this.collectClientState() + const approvalRequests: Array = [] + const clientRequests: Array = [] + for (const toolCall of toolCalls) { + const tool = this.resolveExecutableTools([toolCall]).find( + (candidate) => candidate.name === toolCall.function.name, + ) as RuntimeToolWithApproval | undefined + if (!tool) continue + let input: unknown = {} + try { + const parsed = JSON.parse(toolCall.function.arguments.trim() || '{}') + input = parsed && typeof parsed === 'object' ? parsed : {} + } catch { + input = {} + } + const approvalId = `approval_${toolCall.id}` + if (tool.needsApproval && !approvals.has(approvalId)) { + approvalRequests.push({ + toolCallId: toolCall.id, + toolName: toolCall.function.name, + input, + approvalId, + }) + } else if ( + !tool.execute && + !clientToolResults.has(toolCall.id) && + !this.resumeCancelledToolCallIds.has(toolCall.id) + ) { + clientRequests.push({ + toolCallId: toolCall.id, + toolName: toolCall.function.name, + input, + }) + } + } + return { approvals: approvalRequests, clientRequests } + } + private completeEphemeralInterruptBindings(chunk: StreamChunk): StreamChunk { if ( chunk.type !== EventType.RUN_FINISHED || @@ -2581,7 +3095,7 @@ class TextEngine< private createSyntheticFinishedEvent(): RunFinishedEvent { return { type: 'RUN_FINISHED', - runId: this.createId('pending'), + runId: this.runIdOverride ?? this.requestId, threadId: this.threadId, model: this.params.model, timestamp: Date.now(), @@ -3401,7 +3915,15 @@ class TextEngine< } } - const pending = this.buildActionableInterrupts( + const genericPending = this.getGenericContinuationPending(interruptedRunId) + const pending: Array<{ + interruptId: string + payload: unknown + binding: InterruptBinding + genericRequest?: GenericInterruptRequest< + InterruptDefinition + > + }> = this.buildActionableInterrupts( approvalRequests, clientRequests, ).flatMap((descriptor) => { @@ -3420,6 +3942,7 @@ class TextEngine< ] : [] }) + pending.push(...genericPending) const validated = await validateInterruptResumeBatch({ threadId: this.threadId, interruptedRunId, @@ -3447,6 +3970,187 @@ class TextEngine< ...validated.resumeToolState, approvals, }) + + const genericResolutions = validated.resumeToolState.genericInterrupts + if (genericPending.length > 0 && genericResolutions) { + const resolutions = genericPending + .sort((left, right) => { + const leftIndex = + left.binding.kind === 'generic' ? (left.binding.batchIndex ?? 0) : 0 + const rightIndex = + right.binding.kind === 'generic' + ? (right.binding.batchIndex ?? 0) + : 0 + return leftIndex - rightIndex + }) + .flatMap((record) => { + const resolution = genericResolutions.get(record.interruptId) + if (!resolution || !record.genericRequest) return [] + return [ + resolution.status === 'resolved' + ? { + request: record.genericRequest, + status: 'resolved' as const, + response: resolution.payload, + } + : { + request: record.genericRequest, + status: 'cancelled' as const, + }, + ] + }) + const collection: InterruptResolutionCollection = { + for: (definition) => + resolutions.filter( + (resolution) => resolution.request.definition === definition, + ) as never, + all: ( + ...definitions: Array> + ) => + definitions.length === 0 + ? resolutions + : resolutions.filter((resolution) => + definitions.includes(resolution.request.definition), + ), + } + const policy = await this.middlewareRunner.runOnInterruptResolution( + this.middlewareCtx, + collection, + ) + if (policy.toolResume === 'stop') { + this.earlyTermination = true + } else if (policy.toolResume === 'cancel') { + for (const request of pendingToolCalls) { + this.resumeCancelledToolCallIds.add(request.id) + } + } + } + } + + private getGenericContinuationPending(interruptedRunId: string): Array<{ + interruptId: string + payload: unknown + binding: InterruptBinding + genericRequest: GenericInterruptRequest< + InterruptDefinition + > + }> { + const fail = (message: string): never => { + throw new InterruptResumeValidationError([ + { + scope: 'batch', + threadId: this.threadId, + interruptedRunId, + generation: 0, + interruptIds: [], + code: 'stale', + message, + source: 'server', + retryable: false, + }, + ]) + } + const pending: Array<{ + interruptId: string + payload: unknown + binding: InterruptBinding + genericRequest: GenericInterruptRequest< + InterruptDefinition + > + }> = [] + const ids = new Set() + const batchIndexes = new Set() + for (const resumeItem of this.params.resume ?? []) { + const parsed = readGenericInterruptContinuation(resumeItem.metadata) + if (parsed.status === 'absent') continue + if (parsed.status === 'invalid') { + return fail(parsed.message) + } + const entry = parsed.value + const id = resumeItem.interruptId + const definition = this.interruptDefinitions.get(entry.definitionId) + if (!definition) { + return fail( + `Generic interrupt definition ${entry.definitionId} is unavailable.`, + ) + } + if (ids.has(id) || batchIndexes.has(entry.batchIndex)) { + return fail( + 'Generic interrupt continuation contains duplicate entries.', + ) + } + ids.add(id) + batchIndexes.add(entry.batchIndex) + let request: GenericInterruptRequest< + InterruptDefinition + > + try { + request = Reflect.apply(definition.interrupt, definition, [ + { + key: entry.key, + reason: entry.reason, + message: entry.message, + ...(typeof entry.expiresAt === 'string' + ? { expiresAt: entry.expiresAt } + : {}), + ...(Object.prototype.hasOwnProperty.call(entry, 'payload') + ? { payload: entry.payload } + : {}), + }, + ]) + } catch (error) { + return fail( + `Generic interrupt continuation ${id} is invalid: ${ + error instanceof Error ? error.message : String(error) + }`, + ) + } + const emitted = createInterruptBinding(request, { + batchIndex: entry.batchIndex, + }) + if ( + entry.responseSchemaHash !== emitted.descriptor.responseSchemaHash || + entry.payloadSchemaHash !== emitted.descriptor.payloadSchemaHash + ) { + return fail( + `Generic interrupt continuation ${id} does not match its definition.`, + ) + } + pending.push({ + interruptId: id, + payload: { + id, + ...(emitted.descriptor.responseSchemaCanonicalJson !== undefined + ? { + responseSchema: JSON.parse( + emitted.descriptor.responseSchemaCanonicalJson, + ), + } + : {}), + }, + binding: { + v: INTERRUPT_BINDING_VERSION, + kind: 'generic', + interruptId: id, + interruptedRunId, + generation: 0, + definitionId: entry.definitionId, + key: entry.key, + batchIndex: entry.batchIndex, + ...(typeof entry.expiresAt === 'string' + ? { expiresAt: entry.expiresAt } + : {}), + ...(emitted.descriptor.payloadSchemaHash + ? { payloadSchemaHash: emitted.descriptor.payloadSchemaHash } + : {}), + ...(entry.responseSchemaHash !== undefined + ? { responseSchemaHash: entry.responseSchemaHash } + : {}), + }, + genericRequest: request, + }) + } + return pending } private applyResumeToolState(state: ChatResumeToolState | undefined): void { @@ -3470,6 +4174,58 @@ class TextEngine< this.resumeCancelledToolCallIds.add(toolCallId) } } + if (state?.genericInterrupts) { + for (const [interruptId, resolution] of state.genericInterrupts) { + this.resumeGenericInterrupts.set(interruptId, resolution) + } + } + if (state?.genericInterruptRequests) { + for (const [interruptId, request] of state.genericInterruptRequests) { + this.resumeGenericInterruptRequests.set(interruptId, request) + } + } + } + + private async applyDurableGenericInterruptResolution(): Promise { + if (this.resumeGenericInterruptRequests.size === 0) return + const resolutions = [ + ...this.resumeGenericInterruptRequests.entries(), + ].flatMap(([interruptId, request]) => { + const resolution = this.resumeGenericInterrupts.get(interruptId) + if (!resolution) return [] + return [ + resolution.status === 'resolved' + ? { + request, + status: 'resolved' as const, + response: resolution.payload, + } + : { request, status: 'cancelled' as const }, + ] + }) + const collection: InterruptResolutionCollection = { + for: (definition) => + resolutions.filter( + (resolution) => resolution.request.definition === definition, + ) as never, + all: (...definitions: Array>) => + definitions.length === 0 + ? resolutions + : resolutions.filter((resolution) => + definitions.includes(resolution.request.definition), + ), + } + const policy = await this.middlewareRunner.runOnInterruptResolution( + this.middlewareCtx, + collection, + ) + if (policy.toolResume === 'stop') { + this.earlyTermination = true + } else if (policy.toolResume === 'cancel') { + for (const toolCall of this.getPendingToolCallsFromMessages()) { + this.resumeCancelledToolCallIds.add(toolCall.id) + } + } } private applyMiddlewareConfig(config: ChatMiddlewareConfig): void { @@ -3507,6 +4263,9 @@ class TextEngine< chunk, ) for (const outputChunk of outputChunks) { + if (outputChunk.type === EventType.RUN_STARTED) { + this.hasPublicRunStarted = true + } yield outputChunk this.middlewareCtx.chunkIndex++ } @@ -3647,64 +4406,133 @@ export function chat< TStream, any >['tools'] = TextActivityOptions['tools'], - const TMiddleware extends TextActivityOptions< - TAdapter, - TSchema, - TStream, - any - >['middleware'] = TextActivityOptions< - TAdapter, - TSchema, - TStream, - any - >['middleware'], + const TInterrupts extends ReadonlyArray< + InterruptDefinition + > = [], + TContext = unknown, + const TMiddleware extends Array | undefined = undefined, >( options: TextActivityOptionsWithContext< TAdapter, TSchema, TStream, TTools, + TInterrupts, + TContext, TMiddleware >, ): TextActivityResult { - validateCapabilities(options.middleware ?? [], options.adapter) + validateInterruptDefinitions(options.interrupts) + validateCapabilities( + readRuntimeMiddleware(options.middleware) ?? [], + options.adapter, + ) const { outputSchema, stream } = options - // outputSchema + stream:true is the only branch that streams structured - // output. Without an explicit `stream: true`, schema-bearing calls run the - // agent loop and resolve to a typed Promise>. if (outputSchema && stream === true) { - return runStreamingStructuredOutput({ - ...options, - outputSchema, - stream, - }) as TextActivityResult + return runStreamingStructuredOutput( + toRuntimeTextActivityOptions(options, { + outputSchema, + stream: true, + }), + ) as TextActivityResult } - // If outputSchema is provided, run agentic structured output (Promise) if (outputSchema) { - return runAgenticStructuredOutput({ - ...options, - outputSchema, - }) as TextActivityResult + return runAgenticStructuredOutput( + toRuntimeTextActivityOptions(options, { + outputSchema, + stream: false, + }), + ) as TextActivityResult } - // If stream is explicitly false, run non-streaming text if (stream === false) { - return runNonStreamingText({ - ...options, + return runNonStreamingText( + toRuntimeTextActivityOptions(options, { + outputSchema: undefined, + stream: false, + }), + ) as TextActivityResult + } + + return runStreamingText( + toRuntimeTextActivityOptions(options, { outputSchema: undefined, - stream, - }) as TextActivityResult + stream: true, + }), + ) as TextActivityResult +} + +type RuntimeTextActivityOptions< + TAdapter extends AnyTextAdapter, + TSchema extends SchemaInput | undefined, + TStream extends boolean, +> = Omit, 'middleware'> & { + middleware?: Array +} + +function readRuntimeMiddleware( + middleware: unknown, +): Array | undefined { + if (middleware === undefined) return undefined + if (!Array.isArray(middleware)) { + throw new TypeError('Chat middleware must be an array.') } + return middleware +} - // Otherwise, run streaming text (default) - return runStreamingText({ - ...options, - outputSchema: undefined, - stream, - }) as TextActivityResult +function toRuntimeTextActivityOptions< + TAdapter extends AnyTextAdapter, + TInputSchema extends SchemaInput | undefined, + TInputStream extends boolean, + TOutputSchema extends SchemaInput | undefined, + TOutputStream extends boolean, + TTools extends TextActivityOptions< + TAdapter, + TInputSchema, + TInputStream, + any + >['tools'], + TInterrupts extends ReadonlyArray>, + TContext, + TMiddleware extends Array | undefined, +>( + options: TextActivityOptionsWithContext< + TAdapter, + TInputSchema, + TInputStream, + TTools, + TInterrupts, + TContext, + TMiddleware + >, + overrides: { outputSchema: TOutputSchema; stream: TOutputStream }, +): RuntimeTextActivityOptions { + const { middleware, ...rest } = options + return { + ...rest, + ...overrides, + ...(middleware === undefined + ? {} + : { middleware: readRuntimeMiddleware(middleware) }), + } +} + +function validateInterruptDefinitions( + definitions: + | ReadonlyArray> + | undefined, +): void { + if (!definitions) return + const seen = new Set() + for (const definition of definitions) { + if (seen.has(definition.id)) { + throw new Error(`Duplicate interrupt definition id: ${definition.id}`) + } + seen.add(definition.id) + } } /** @@ -3756,8 +4584,8 @@ function publishDeliverySeams( * returns, so the identity has to be minted out here and the engine reached back * through `engineRef`, which the body fills as soon as its engine exists. */ -function runStreamingText( - options: TextActivityOptions, +function runStreamingText( + options: RuntimeTextActivityOptions, ): AsyncIterable { const engineRef: DeliveryEngineRef = {} const stream = streamTextChunks(options, engineRef) @@ -3765,8 +4593,8 @@ function runStreamingText( return stream } -async function* streamTextChunks( - options: TextActivityOptions, +async function* streamTextChunks( + options: RuntimeTextActivityOptions, engineRef: DeliveryEngineRef, ): AsyncIterable { const { adapter, middleware, context, debug, mcp, ...textOptions } = options @@ -3785,7 +4613,7 @@ async function* streamTextChunks( params: { ...textOptions, model, logger } as TextOptions< Record, Record, - TContext + any >, middleware, context, @@ -3807,19 +4635,13 @@ async function* streamTextChunks( * Run non-streaming text - collects all content and returns as a string. * Runs the full agentic loop (if tools are provided) but returns collected text. */ -function runNonStreamingText( - options: TextActivityOptions, +function runNonStreamingText( + options: RuntimeTextActivityOptions, ): Promise { - // Run the streaming text and collect all text using streamToText. - const stream = runStreamingText( - // oxlint-disable-next-line eslint-js/no-restricted-syntax -- generic-stream remap: caller is non-streaming (false), but runStreamingText is invoked internally to collect text; concrete `false`→`true` literals don't structurally overlap. - options as unknown as TextActivityOptions< - AnyTextAdapter, - undefined, - true, - TContext - >, - ) + const stream = runStreamingText({ + ...options, + stream: true, + }) return streamToText(stream) } @@ -3830,11 +4652,8 @@ function runNonStreamingText( * 2. Once complete, call adapter.structuredOutput with the conversation context * 3. Validate and return the structured result */ -async function runAgenticStructuredOutput< - TSchema extends SchemaInput, - TContext = unknown, ->( - options: TextActivityOptions, +async function runAgenticStructuredOutput( + options: RuntimeTextActivityOptions, ): Promise> { const { adapter, @@ -3903,7 +4722,7 @@ async function runAgenticStructuredOutput< params: { ...textOptions, model, logger } as TextOptions< Record, Record, - TContext + any >, middleware, context, @@ -4110,11 +4929,8 @@ async function* fallbackStructuredOutputStream( * synchronously at call time rather than as a yielded RUN_ERROR mid-stream — * those are programmer errors, not runtime conditions. */ -function runStreamingStructuredOutput< - TSchema extends SchemaInput, - TContext = unknown, ->( - options: TextActivityOptions, +function runStreamingStructuredOutput( + options: RuntimeTextActivityOptions, ): StructuredOutputStream> { const { outputSchema } = options @@ -4175,11 +4991,8 @@ type StructuredOutputStreamInternal = AsyncIterable< StreamChunk | StructuredOutputCompleteEvent > -async function* runStreamingStructuredOutputImpl< - TSchema extends SchemaInput, - TContext = unknown, ->( - options: TextActivityOptions, +async function* runStreamingStructuredOutputImpl( + options: RuntimeTextActivityOptions, jsonSchema: NonNullable>, normalize: (data: unknown) => unknown, engineRef: DeliveryEngineRef, @@ -4220,7 +5033,7 @@ async function* runStreamingStructuredOutputImpl< params: { ...textOptions, model, logger } as TextOptions< Record, Record, - TContext + any >, middleware, context, diff --git a/packages/ai/src/activities/chat/middleware/builder.ts b/packages/ai/src/activities/chat/middleware/builder.ts index 66e237b340..5d889ba3df 100644 --- a/packages/ai/src/activities/chat/middleware/builder.ts +++ b/packages/ai/src/activities/chat/middleware/builder.ts @@ -1,6 +1,9 @@ import type { CapabilityHandle } from './capabilities' import type { AnyChatMiddleware, ChatMiddleware } from './types' import type { DefinedChatMiddleware } from './define' +import type { InterruptDefinition } from '../../../interrupt-definition' + +type AnyInterruptDefinition = InterruptDefinition /** Union of capability NAME literals from a tuple of handles. */ export type NamesOf> = @@ -67,19 +70,41 @@ export type CheckCoverage> = [ export interface ChatMiddlewareBuilder< TList extends ReadonlyArray, TProvided extends string, + TInterruptDefinitions extends AnyInterruptDefinition = never, > { use: < TRequires extends ReadonlyArray, TProvides extends ReadonlyArray, TContext = unknown, + TMiddlewareInterruptDefinitions extends AnyInterruptDefinition = + TInterruptDefinitions, >( middleware: [NamesOf] extends [TProvided] - ? DefinedChatMiddleware - : DefinedChatMiddleware & + ? DefinedChatMiddleware< + TContext, + TRequires, + TProvides, + TMiddlewareInterruptDefinitions + > + : DefinedChatMiddleware< + TContext, + TRequires, + TProvides, + TMiddlewareInterruptDefinitions + > & MissingCapabilities, TProvided>>, ) => ChatMiddlewareBuilder< - readonly [...TList, DefinedChatMiddleware], - TProvided | NamesOf + readonly [ + ...TList, + DefinedChatMiddleware< + TContext, + TRequires, + TProvides, + TMiddlewareInterruptDefinitions + >, + ], + TProvided | NamesOf, + TInterruptDefinitions | TMiddlewareInterruptDefinitions > build: () => [...TList] diff --git a/packages/ai/src/activities/chat/middleware/compose.ts b/packages/ai/src/activities/chat/middleware/compose.ts index dd3cf548bd..3e811cfbaa 100644 --- a/packages/ai/src/activities/chat/middleware/compose.ts +++ b/packages/ai/src/activities/chat/middleware/compose.ts @@ -10,6 +10,9 @@ import type { ChatMiddlewareContext, ErrorInfo, FinishInfo, + InterruptBoundaryPhase, + InterruptResolutionCollection, + InterruptToolResume, IterationInfo, SandboxFileHookEvent, StructuredOutputMiddlewareConfig, @@ -17,6 +20,10 @@ import type { ToolPhaseCompleteInfo, UsageInfo, } from './types' +import type { + GenericInterruptRequest, + InterruptDefinition, +} from '../../../interrupt-definition' /** One middleware's terminal-hook throw, captured instead of propagated. */ interface HookFailure { @@ -25,7 +32,7 @@ interface HookFailure { } /** Check if a middleware should be skipped for instrumentation events. */ -function shouldSkipInstrumentation(mw: ChatMiddleware): boolean { +function shouldSkipInstrumentation(mw: ChatMiddleware): boolean { return mw.name === 'devtools' || mw.name === 'strip-to-spec' } @@ -43,12 +50,18 @@ function instrumentCtx(ctx: ChatMiddlewareContext) { * Internal middleware runner that manages composed execution of middleware hooks. * Created once per chat() invocation. */ -export class MiddlewareRunner { - private readonly middlewares: ReadonlyArray> +export class MiddlewareRunner< + TContext = unknown, + TInterruptDefinitions extends InterruptDefinition = + InterruptDefinition, +> { + private readonly middlewares: ReadonlyArray< + ChatMiddleware + > private readonly logger: InternalLogger constructor( - middlewares: ReadonlyArray>, + middlewares: ReadonlyArray>, logger: InternalLogger, ) { this.middlewares = middlewares @@ -59,6 +72,83 @@ export class MiddlewareRunner { return this.middlewares.length > 0 } + async runOnInterruptBoundary( + ctx: ChatMiddlewareContext & { phase: InterruptBoundaryPhase }, + ): Promise>> { + const requests: Array> = [] + for (const mw of this.middlewares) { + if (mw.onInterruptBoundary) { + const skip = shouldSkipInstrumentation(mw) + const start = Date.now() + const result = await mw.onInterruptBoundary(ctx) + if (result?.interrupts) requests.push(...result.interrupts) + if (!skip) { + this.logger.middleware( + `hook=onInterruptBoundary middleware=${mw.name ?? 'unnamed'}`, + { + middleware: mw.name ?? 'unnamed', + hook: 'onInterruptBoundary', + }, + ) + aiEventClient.emit('middleware:hook:executed', { + ...instrumentCtx(ctx), + middlewareName: mw.name || 'unnamed', + hookName: 'onInterruptBoundary', + iteration: ctx.iteration, + duration: Date.now() - start, + hasTransform: result?.interrupts !== undefined, + }) + } + } + } + return requests + } + + async runOnInterruptResolution( + ctx: ChatMiddlewareContext, + resolutions: InterruptResolutionCollection, + ): Promise<{ toolResume?: InterruptToolResume }> { + let toolResume: InterruptToolResume | undefined + for (const mw of this.middlewares) { + if (mw.onInterruptResolution) { + const skip = shouldSkipInstrumentation(mw) + const start = Date.now() + const next = await mw.onInterruptResolution(ctx, resolutions) + if (next?.toolResume !== undefined) { + const priority: Record = { + continue: 0, + cancel: 1, + stop: 2, + } + if ( + toolResume === undefined || + priority[next.toolResume] > priority[toolResume] + ) { + toolResume = next.toolResume + } + } + if (!skip) { + this.logger.middleware( + `hook=onInterruptResolution middleware=${mw.name ?? 'unnamed'}`, + { + middleware: mw.name ?? 'unnamed', + hook: 'onInterruptResolution', + }, + ) + aiEventClient.emit('middleware:hook:executed', { + ...instrumentCtx(ctx), + middlewareName: mw.name || 'unnamed', + hookName: 'onInterruptResolution', + iteration: ctx.iteration, + duration: Date.now() - start, + hasTransform: next !== undefined, + }) + } + } + } + return toolResume === undefined ? {} : { toolResume } + } + /** * Pipe config through all middleware onConfig hooks in order. * Each middleware receives the merged config from previous middleware. @@ -492,7 +582,7 @@ export class MiddlewareRunner { * {@link runOnError}. */ private async captureTerminalHook( - mw: ChatMiddleware, + mw: ChatMiddleware, hookName: 'onFinish' | 'onAbort' | 'onError', invoke: () => void | Promise, ): Promise { diff --git a/packages/ai/src/activities/chat/middleware/define.ts b/packages/ai/src/activities/chat/middleware/define.ts index 92234f1a2f..18ff5d2e36 100644 --- a/packages/ai/src/activities/chat/middleware/define.ts +++ b/packages/ai/src/activities/chat/middleware/define.ts @@ -1,5 +1,8 @@ import type { CapabilityHandle } from './capabilities' import type { ChatMiddleware } from './types' +import type { InterruptDefinition } from '../../../interrupt-definition' + +type AnyInterruptDefinition = InterruptDefinition /** * A middleware whose `requires`/`provides` tuple types are captured precisely @@ -9,7 +12,8 @@ export interface DefinedChatMiddleware< TContext, TRequires extends ReadonlyArray, TProvides extends ReadonlyArray, -> extends ChatMiddleware { + TInterruptDefinitions extends AnyInterruptDefinition = never, +> extends ChatMiddleware { requires?: TRequires provides?: TProvides } @@ -24,11 +28,17 @@ export function defineChatMiddleware< TContext = unknown, const TRequires extends ReadonlyArray = readonly [], const TProvides extends ReadonlyArray = readonly [], + TInterruptDefinitions extends AnyInterruptDefinition = never, >( - middleware: ChatMiddleware & { + middleware: ChatMiddleware & { requires?: TRequires provides?: TProvides }, -): DefinedChatMiddleware { +): DefinedChatMiddleware< + TContext, + TRequires, + TProvides, + TInterruptDefinitions +> { return middleware } diff --git a/packages/ai/src/activities/chat/middleware/generic-interrupts.ts b/packages/ai/src/activities/chat/middleware/generic-interrupts.ts new file mode 100644 index 0000000000..01f2c2f535 --- /dev/null +++ b/packages/ai/src/activities/chat/middleware/generic-interrupts.ts @@ -0,0 +1,26 @@ +import { createCapability } from './capabilities' +import type { InterruptDefinition } from '../../../interrupt-definition' + +/** + * Internal per-run registry of first-party generic interrupt definitions. + * + * Persistence uses this bridge while restoring a durable interrupt. It never + * receives a definition from the stored record; it can only look up one that + * the current chat call registered. + */ +export interface GenericInterruptDefinitionRegistry { + readonly definitions: ReadonlyMap< + string, + InterruptDefinition + > +} + +export const GenericInterruptDefinitionRegistryCapability = + createCapability()( + 'generic-interrupt-definition-registry', + ) + +export const [ + getGenericInterruptDefinitionRegistry, + provideGenericInterruptDefinitionRegistry, +] = GenericInterruptDefinitionRegistryCapability diff --git a/packages/ai/src/activities/chat/middleware/index.ts b/packages/ai/src/activities/chat/middleware/index.ts index 2611cd6ae4..1f913cfa77 100644 --- a/packages/ai/src/activities/chat/middleware/index.ts +++ b/packages/ai/src/activities/chat/middleware/index.ts @@ -18,8 +18,23 @@ export type { SandboxFileEvent, SandboxFileHookEvent, ChatSandboxHooks, + InterruptBoundaryPhase, + InterruptToolResume, + InterruptResolutionCollection, + GenericInterruptResolution, + InterruptBoundaryResult, + InterruptResolutionResult, } from './types' +export { INTERRUPT_BOUNDARY_PHASES, INTERRUPT_TOOL_RESUMES } from './types' + +export { + GenericInterruptDefinitionRegistryCapability, + getGenericInterruptDefinitionRegistry, + provideGenericInterruptDefinitionRegistry, +} from './generic-interrupts' +export type { GenericInterruptDefinitionRegistry } from './generic-interrupts' + export { MiddlewareRunner } from './compose' export { createCapability, CapabilityRegistry } from './capabilities' diff --git a/packages/ai/src/activities/chat/middleware/types.ts b/packages/ai/src/activities/chat/middleware/types.ts index 28096517a2..6cd5de10fb 100644 --- a/packages/ai/src/activities/chat/middleware/types.ts +++ b/packages/ai/src/activities/chat/middleware/types.ts @@ -1,3 +1,7 @@ +import type { + StandardJSONSchemaV1, + StandardSchemaV1, +} from '@standard-schema/spec' import type { AgentLoopState, JSONSchema, @@ -10,6 +14,10 @@ import type { } from '../../../types' import type { SystemPrompt } from '../../../system-prompts' import type { ToolApprovalResolution } from '../../../interrupts' +import type { + GenericInterruptRequest, + InterruptDefinition, +} from '../../../interrupt-definition' import type { Capability, CapabilityHandle, @@ -68,6 +76,7 @@ export interface ChatSandboxHooks { * Phase of the chat middleware lifecycle. * - 'init': Initial config transform before the chat engine starts * - 'beforeModel': Before each adapter chatStream call (per agent iteration) + * - 'afterModel': After each adapter chatStream call (per agent iteration) * - 'modelStream': During model streaming * - 'beforeTools': Before tool execution phase * - 'afterTools': After tool execution phase @@ -77,11 +86,97 @@ export interface ChatSandboxHooks { export type ChatMiddlewarePhase = | 'init' | 'beforeModel' + | 'afterModel' | 'modelStream' | 'beforeTools' | 'afterTools' | 'structuredOutput' +export const INTERRUPT_BOUNDARY_PHASES = [ + 'beforeModel', + 'afterModel', + 'beforeTools', + 'afterTools', +] as const + +export type InterruptBoundaryPhase = (typeof INTERRUPT_BOUNDARY_PHASES)[number] + +export const INTERRUPT_TOOL_RESUMES = ['continue', 'cancel', 'stop'] as const + +export type InterruptToolResume = (typeof INTERRUPT_TOOL_RESUMES)[number] + +type AnyInterruptDefinition = InterruptDefinition + +type InterruptResponse = + TDefinition extends InterruptDefinition + ? TResponseSchema extends StandardSchemaV1 + ? TResponse + : TResponseSchema extends StandardJSONSchemaV1 + ? TResponse + : unknown + : unknown + +export type GenericInterruptResolution< + TDefinition extends AnyInterruptDefinition, +> = TDefinition extends AnyInterruptDefinition + ? + | { + readonly request: GenericInterruptRequest + readonly status: 'resolved' + readonly response: InterruptResponse + } + | { + readonly request: GenericInterruptRequest + readonly status: 'cancelled' + readonly response?: never + } + : never + +export interface InterruptResolutionCollection< + TDefinitions extends AnyInterruptDefinition = AnyInterruptDefinition, +> { + for: < + TDefinition extends ([TDefinitions] extends [never] + ? AnyInterruptDefinition + : TDefinitions), + >( + definition: TDefinition, + ) => ReadonlyArray> + all: { + (): ReadonlyArray> + >( + ...definitions: TSelected + ): ReadonlyArray> + } +} + +type BivariantInterruptResolutionHook< + TContext, + TDefinitions extends AnyInterruptDefinition, +> = InterruptResolutionHookSignature['call'] + +declare abstract class InterruptResolutionHookSignature< + TContext, + TDefinitions extends AnyInterruptDefinition, +> { + abstract call( + ctx: ChatMiddlewareContext, + resolutions: InterruptResolutionCollection, + ): InterruptResolutionResult | Promise +} + +export type InterruptBoundaryResult< + TDefinitions extends AnyInterruptDefinition = AnyInterruptDefinition, +> = + | undefined + | { + readonly interrupts: ReadonlyArray> + } + +export type InterruptResolutionResult = void | { + readonly toolResume: InterruptToolResume +} + /** * Stable context object passed to all middleware hooks. * Created once per chat() invocation and shared across all hooks. @@ -230,6 +325,13 @@ export interface ChatResumeToolState { genericInterrupts?: | ReadonlyMap | undefined + /** Durable generic requests reconstructed by server middleware. */ + genericInterruptRequests?: + | ReadonlyMap< + string, + GenericInterruptRequest> + > + | undefined deniedToolResults?: ReadonlyMap | undefined cancelledToolCallIds?: ReadonlySet | undefined } @@ -460,10 +562,32 @@ export interface ErrorInfo { * } * ``` */ -export interface ChatMiddleware { +export interface ChatMiddleware< + TContext = unknown, + TInterruptDefinitions extends AnyInterruptDefinition = never, +> { /** Optional name for debugging and identification */ name?: string + /** + * Called at a lifecycle boundary. Return interrupt requests to pause the run. + * Requests from every middleware in the same boundary form one batch. + */ + onInterruptBoundary?: ( + ctx: ChatMiddlewareContext & { phase: InterruptBoundaryPhase }, + ) => + | InterruptBoundaryResult + | Promise> + + /** + * Called on a continuation run after the client answers registered interrupts. + * Return `toolResume` to decide whether pending tools continue, cancel, or stop. + */ + onInterruptResolution?: BivariantInterruptResolutionHook< + TContext, + TInterruptDefinitions + > + /** * Capabilities this middleware requires. `chat()` validates that some * middleware (or the adapter) provides each one; unsatisfied requirements are @@ -653,4 +777,5 @@ export interface ChatMiddleware { } /** A `ChatMiddleware` with a permissive context — for use as a constraint. */ -export type AnyChatMiddleware = ChatMiddleware +/** A permissive middleware constraint that retains the definition parameter. */ +export type AnyChatMiddleware = ChatMiddleware diff --git a/packages/ai/src/activities/chat/stream/processor.ts b/packages/ai/src/activities/chat/stream/processor.ts index 1994a5b37c..90991e1b95 100644 --- a/packages/ai/src/activities/chat/stream/processor.ts +++ b/packages/ai/src/activities/chat/stream/processor.ts @@ -38,6 +38,7 @@ import { updateToolResultPart, } from './message-updaters' import { ImmediateStrategy } from './strategies' +import { INTERRUPT_BINDING_METADATA_KEY } from '../../../interrupt-resume' import type { ChunkRecording, ChunkStrategy, @@ -138,6 +139,22 @@ export interface StreamProcessorOptions { const STRUCTURED_OUTPUT_UPDATE_BATCH_SIZE = 12 +function interruptBatchHasGeneric(interrupts: Array): boolean { + return interrupts.some((interrupt) => { + const metadata = interrupt.metadata + if (!metadata || typeof metadata !== 'object' || Array.isArray(metadata)) { + return false + } + const binding = metadata[INTERRUPT_BINDING_METADATA_KEY] + return ( + binding !== null && + typeof binding === 'object' && + !Array.isArray(binding) && + binding.kind === 'generic' + ) + }) +} + /** * StreamProcessor - State machine for processing AI response streams * @@ -1548,6 +1565,7 @@ export class StreamProcessor { } private handleInterrupts(interrupts: Array): void { + const hasGeneric = interruptBatchHasGeneric(interrupts) for (const interrupt of interrupts) { const metadata = interrupt.metadata && typeof interrupt.metadata === 'object' @@ -1598,6 +1616,9 @@ export class StreamProcessor { } if (kind === 'client_tool' || interrupt.reason === 'client_tool_input') { + // Generic interrupts in the same batch decide `toolResume`. Do not + // run client tools until that policy is `continue`. + if (hasGeneric) continue this.events.onToolCall?.({ toolCallId, toolName, diff --git a/packages/ai/src/adapter-internals.ts b/packages/ai/src/adapter-internals.ts index b0c2612bb6..4e7305dba8 100644 --- a/packages/ai/src/adapter-internals.ts +++ b/packages/ai/src/adapter-internals.ts @@ -25,3 +25,23 @@ export { PendingTurnCapability, providePendingTurn, } from './activities/chat/middleware/pending-turn' +export { + getGenericInterruptDefinitionRegistry, + GenericInterruptDefinitionRegistryCapability, + provideGenericInterruptDefinitionRegistry, +} from './activities/chat/middleware/generic-interrupts' +export type { GenericInterruptDefinitionRegistry } from './activities/chat/middleware/generic-interrupts' +export { + createInterruptBinding, + getInterruptRequestInput, + rehydrateInterruptRequest, +} from './interrupt-definition' +export type { + GenericInterruptRequest, + InterruptDefinition, +} from './interrupt-definition' +export { + readInterruptBinding, + validateInterruptResumeBatch, +} from './interrupt-resume' +export type { PendingInterruptResumeRecord } from './interrupt-resume' diff --git a/packages/ai/src/client.ts b/packages/ai/src/client.ts index a677e92ea4..d5f048cf75 100644 --- a/packages/ai/src/client.ts +++ b/packages/ai/src/client.ts @@ -343,6 +343,26 @@ export { INTERRUPT_BINDING_VERSION, canonicalizeInterruptResolutions, } from './interrupts' +export { + defineInterrupt, + hashInterruptDefinitionSchema, + INTERRUPT_PAYLOAD_METADATA_KEY, +} from './interrupt-definition' +export { + INTERRUPT_CONTINUATION_METADATA_KEY, + INTERRUPT_CONTINUATION_VERSION, + genericInterruptContinuationFromDescriptor, + readGenericInterruptContinuation, + wrapGenericInterruptContinuation, +} from './generic-interrupt-continuation' +export type { + GenericInterruptContinuation, + GenericInterruptContinuationReadResult, +} from './generic-interrupt-continuation' +export type { + GenericInterruptRequest, + InterruptDefinition, +} from './interrupt-definition' export type { BatchInterruptError, BatchInterruptErrorCode, diff --git a/packages/ai/src/generic-interrupt-continuation.ts b/packages/ai/src/generic-interrupt-continuation.ts new file mode 100644 index 0000000000..cc84892f88 --- /dev/null +++ b/packages/ai/src/generic-interrupt-continuation.ts @@ -0,0 +1,162 @@ +import { INTERRUPT_PAYLOAD_METADATA_KEY } from './interrupt-definition' +import { readUnopenedInterruptBinding } from './interrupt-resume' +import type { Interrupt } from './types' + +/** + * `ResumeEntry.metadata` key for a first-party generic request. + * + * AG-UI `resume` only carries the answer (`interruptId`, `status`, `payload`). + * The original request rides here so an ephemeral server can rebuild it. + */ +export const INTERRUPT_CONTINUATION_METADATA_KEY = + 'tanstack:interruptContinuation' as const + +export const INTERRUPT_CONTINUATION_VERSION = 1 as const + +export interface GenericInterruptContinuation { + v: typeof INTERRUPT_CONTINUATION_VERSION + definitionId: string + key: string + batchIndex: number + reason: string + message: string + expiresAt?: string + responseSchemaHash?: string + payloadSchemaHash?: string + payload?: unknown +} + +export type GenericInterruptContinuationReadResult = + | { status: 'absent' } + | { status: 'invalid'; message: string } + | { status: 'ok'; value: GenericInterruptContinuation } + +function isRecord(value: unknown): value is Record { + return value !== null && typeof value === 'object' && !Array.isArray(value) +} + +function invalid( + message: string, +): Extract { + return { status: 'invalid', message } +} + +/** + * Read one generic request from `resume[].metadata`. + * + * Missing key means this resume item is not a first-party generic continuation. + * A present key that fails the shape is a protocol error. + */ +export function readGenericInterruptContinuation( + metadata: unknown, +): GenericInterruptContinuationReadResult { + if (metadata === undefined) return { status: 'absent' } + if (!isRecord(metadata)) { + return invalid('Generic interrupt resume metadata must be an object.') + } + if ( + !Object.prototype.hasOwnProperty.call( + metadata, + INTERRUPT_CONTINUATION_METADATA_KEY, + ) + ) { + return { status: 'absent' } + } + const raw = metadata[INTERRUPT_CONTINUATION_METADATA_KEY] + if (!isRecord(raw)) { + return invalid('Generic interrupt continuation is invalid.') + } + if ( + raw.v !== INTERRUPT_CONTINUATION_VERSION || + typeof raw.definitionId !== 'string' || + typeof raw.key !== 'string' || + typeof raw.reason !== 'string' || + typeof raw.message !== 'string' || + typeof raw.batchIndex !== 'number' || + !Number.isInteger(raw.batchIndex) || + raw.batchIndex < 0 || + (raw.responseSchemaHash !== undefined && + typeof raw.responseSchemaHash !== 'string') || + (raw.expiresAt !== undefined && typeof raw.expiresAt !== 'string') || + (raw.payloadSchemaHash !== undefined && + typeof raw.payloadSchemaHash !== 'string') + ) { + return invalid('Generic interrupt continuation contains invalid fields.') + } + return { + status: 'ok', + value: { + v: INTERRUPT_CONTINUATION_VERSION, + definitionId: raw.definitionId, + key: raw.key, + batchIndex: raw.batchIndex, + reason: raw.reason, + message: raw.message, + ...(typeof raw.expiresAt === 'string' + ? { expiresAt: raw.expiresAt } + : {}), + ...(typeof raw.responseSchemaHash === 'string' + ? { responseSchemaHash: raw.responseSchemaHash } + : {}), + ...(typeof raw.payloadSchemaHash === 'string' + ? { payloadSchemaHash: raw.payloadSchemaHash } + : {}), + ...(Object.prototype.hasOwnProperty.call(raw, 'payload') + ? { payload: raw.payload } + : {}), + }, + } +} + +/** Put a parsed continuation on `ResumeEntry.metadata`. */ +export function wrapGenericInterruptContinuation( + continuation: GenericInterruptContinuation, +): Record { + return { [INTERRUPT_CONTINUATION_METADATA_KEY]: continuation } +} + +/** + * Build the resume-metadata continuation from an outbound AG-UI interrupt. + * + * Returns `undefined` when the descriptor is not a first-party generic item. + */ +export function genericInterruptContinuationFromDescriptor( + interrupt: Interrupt, +): GenericInterruptContinuation | undefined { + const binding = readUnopenedInterruptBinding(interrupt) + if ( + binding?.kind !== 'generic' || + binding.definitionId === undefined || + binding.key === undefined || + binding.batchIndex === undefined + ) { + return undefined + } + const metadata = isRecord(interrupt.metadata) ? interrupt.metadata : undefined + const hasPayload = + metadata !== undefined && + Object.prototype.hasOwnProperty.call( + metadata, + INTERRUPT_PAYLOAD_METADATA_KEY, + ) + return { + v: INTERRUPT_CONTINUATION_VERSION, + definitionId: binding.definitionId, + key: binding.key, + batchIndex: binding.batchIndex, + reason: interrupt.reason, + message: interrupt.message ?? '', + ...(interrupt.expiresAt !== undefined + ? { expiresAt: interrupt.expiresAt } + : {}), + ...(binding.responseSchemaHash !== undefined + ? { responseSchemaHash: binding.responseSchemaHash } + : {}), + ...(binding.payloadSchemaHash !== undefined + ? { payloadSchemaHash: binding.payloadSchemaHash } + : {}), + ...(hasPayload + ? { payload: metadata[INTERRUPT_PAYLOAD_METADATA_KEY] } + : {}), + } +} diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index cbc12c7cbb..6cc8a1cdc0 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -184,6 +184,17 @@ export type { SandboxFileEvent, SandboxFileHookEvent, ChatSandboxHooks, + InterruptBoundaryPhase, + InterruptToolResume, + InterruptResolutionCollection, + GenericInterruptResolution, + InterruptBoundaryResult, + InterruptResolutionResult, +} from './activities/chat/middleware/index' + +export { + INTERRUPT_BOUNDARY_PHASES, + INTERRUPT_TOOL_RESUMES, } from './activities/chat/middleware/index' // Interrupt protocol surface. Deliberately enumerated rather than @@ -192,6 +203,29 @@ export type { // a commitment. Only the ephemeral contract this release actually implements // is exported — no durable-recovery or persisted-state types, which would // pre-decide a question the orchestration RFC still owns. +export { + defineInterrupt, + createInterruptBinding, + INTERRUPT_PAYLOAD_METADATA_KEY, +} from './interrupt-definition' +export { + INTERRUPT_CONTINUATION_METADATA_KEY, + INTERRUPT_CONTINUATION_VERSION, + genericInterruptContinuationFromDescriptor, + readGenericInterruptContinuation, + wrapGenericInterruptContinuation, +} from './generic-interrupt-continuation' +export type { + GenericInterruptContinuation, + GenericInterruptContinuationReadResult, +} from './generic-interrupt-continuation' +export type { + InterruptDefinition, + GenericInterruptRequest, + InterruptDefinitionOptions, + InterruptBindingDescriptor, +} from './interrupt-definition' + export { INTERRUPT_BINDING_VERSION, canonicalizeInterruptResolutions, diff --git a/packages/ai/src/interrupt-definition.ts b/packages/ai/src/interrupt-definition.ts new file mode 100644 index 0000000000..48c23ef386 --- /dev/null +++ b/packages/ai/src/interrupt-definition.ts @@ -0,0 +1,581 @@ +import type { + StandardJSONSchemaV1, + StandardSchemaV1, +} from '@standard-schema/spec' +import { + canonicalInterruptJson, + cloneAndDeepFreezeJson, + digestInterruptJson, +} from './interrupt-serialization' +import { + isStandardSchema, + isStandardJSONSchema, +} from './activities/chat/tools/schema-converter' + +export const INTERRUPT_PAYLOAD_METADATA_KEY = + 'tanstack:interruptPayload' as const +export const INTERRUPT_BINDING_KIND = 'generic' as const +export const INTERRUPT_BINDING_VERSION = 1 as const + +type PortableSchema = + | StandardJSONSchemaV1 + | StandardSchemaV1 + +type InferSchemaOutput = + TSchema extends StandardSchemaV1 + ? TOutput + : TSchema extends StandardJSONSchemaV1 + ? TOutput + : never +type InferSchemaInput = + TSchema extends StandardSchemaV1 + ? TInput + : TSchema extends StandardJSONSchemaV1 + ? TInput + : never +type DefinitionSchemaState = { + responseSchemaCanonicalJson?: string + responseSchemaHash?: string + payloadSchemaCanonicalJson?: string + payloadSchemaHash?: string +} +const definitionSchemaState = new WeakMap() + +export interface InterruptDefinitionOptions< + TId extends string, + TPayloadSchema extends PortableSchema | undefined, + TResponseSchema extends PortableSchema | undefined, +> { + id: TId + payloadSchema?: TPayloadSchema + responseSchema?: TResponseSchema +} + +export interface InterruptBindingDescriptor { + v: typeof INTERRUPT_BINDING_VERSION + kind: typeof INTERRUPT_BINDING_KIND + definitionId: string + key: string + threadId?: string + interruptedRunId?: string + generation?: number + batchIndex?: number + responseSchemaCanonicalJson?: string + payloadSchemaCanonicalJson?: string + payloadSchemaHash?: string + responseSchemaHash?: string +} + +export interface InterruptPreEmissionData { + descriptor: InterruptBindingDescriptor + payload?: unknown +} + +type InterruptInput< + TPayloadSchema extends PortableSchema | undefined, + TPayload = unknown, +> = { + key: string + reason: string + message: string + expiresAt?: string +} & ([TPayloadSchema] extends [undefined] ? {} : { payload?: TPayload }) + +type GenericInterruptRequestBase< + TDefinition extends InterruptDefinition, +> = { + readonly definition: TDefinition + readonly key: string + readonly reason: string + readonly message: string + readonly expiresAt?: string +} + +type GenericInterruptRequestFor< + TDefinition extends InterruptDefinition, + TPayloadSchema extends PortableSchema | undefined, + TPayload, +> = GenericInterruptRequestBase & + ([TPayloadSchema] extends [undefined] + ? {} + : { readonly payload: TPayload | undefined }) + +export type GenericInterruptRequest< + TDefinition extends InterruptDefinition, +> = [TDefinition] extends [never] + ? never + : TDefinition extends InterruptDefinition< + any, + infer TPayloadSchema, + any, + infer TPayload + > + ? GenericInterruptRequestFor + : GenericInterruptRequestBase + +type InterruptInputKey = 'key' | 'reason' | 'message' | 'expiresAt' | 'payload' +type RejectUnexpectedInputKeys = + Exclude extends never + ? unknown + : { [K in Exclude]: never } +type RejectUnexpectedPayload = 'payload' extends keyof TInput + ? { payload: never } + : unknown +type ValidInterruptInput< + TInput, + TPayloadSchema extends PortableSchema | undefined, + TPayload = unknown, +> = + TInput extends InterruptInput + ? RejectUnexpectedInputKeys & + ([TPayloadSchema] extends [undefined] + ? RejectUnexpectedPayload + : unknown) + : never + +/** + * Extracting a class method preserves the intentional bivariant assignment + * behavior of the public `interrupt` callback without exposing a method + * signature in an interface. + */ +declare abstract class InterruptRequestMethodSignature< + TId extends string, + TPayloadSchema extends PortableSchema | undefined, + TResponseSchema extends PortableSchema | undefined, + TPayload, + TPayloadInput, +> { + abstract call( + input: TInput & ValidInterruptInput, + ): GenericInterruptRequestFor< + InterruptDefinition< + TId, + TPayloadSchema, + TResponseSchema, + TPayload, + TPayloadInput + >, + TPayloadSchema, + TPayload + > +} + +type InterruptRequestMethod< + TId extends string, + TPayloadSchema extends PortableSchema | undefined, + TResponseSchema extends PortableSchema | undefined, + TPayload, + TPayloadInput, +> = InterruptRequestMethodSignature< + TId, + TPayloadSchema, + TResponseSchema, + TPayload, + TPayloadInput +>['call'] + +type DefinedInterruptDefinition< + TId extends string, + TPayloadSchema extends PortableSchema | undefined, + TResponseSchema extends PortableSchema | undefined, + TPayload = unknown, + TPayloadInput = TPayload, +> = Omit< + InterruptDefinition< + TId, + TPayloadSchema, + TResponseSchema, + TPayload, + TPayloadInput + >, + 'interrupt' +> & { + interrupt: InterruptRequestMethod< + TId, + TPayloadSchema, + TResponseSchema, + TPayload, + TPayloadInput + > +} + +export interface InterruptDefinition< + TId extends string, + TPayloadSchema extends PortableSchema | undefined, + TResponseSchema extends PortableSchema | undefined, + TPayload = unknown, + TPayloadInput = TPayload, +> { + readonly id: TId + readonly payloadSchema: TPayloadSchema + readonly responseSchema: TResponseSchema + interrupt: InterruptRequestMethod< + TId, + TPayloadSchema, + TResponseSchema, + TPayload, + TPayloadInput + > +} + +export function createInterruptBinding( + request: GenericInterruptRequest>, + fields: Pick< + InterruptBindingDescriptor, + 'threadId' | 'interruptedRunId' | 'generation' | 'batchIndex' + > = {}, +): InterruptPreEmissionData { + const schemaState = definitionSchemaState.get(request.definition) + if (!schemaState) { + throw new TypeError('Interrupt definition schema state is unavailable.') + } + const { threadId, interruptedRunId, generation, batchIndex } = fields + return { + descriptor: { + v: INTERRUPT_BINDING_VERSION, + kind: INTERRUPT_BINDING_KIND, + definitionId: request.definition.id, + key: request.key, + ...(threadId !== undefined ? { threadId } : {}), + ...(interruptedRunId !== undefined ? { interruptedRunId } : {}), + ...(generation !== undefined ? { generation } : {}), + ...(batchIndex !== undefined ? { batchIndex } : {}), + ...(schemaState.responseSchemaCanonicalJson + ? { + responseSchemaCanonicalJson: + schemaState.responseSchemaCanonicalJson, + } + : {}), + ...(schemaState.payloadSchemaCanonicalJson + ? { payloadSchemaCanonicalJson: schemaState.payloadSchemaCanonicalJson } + : {}), + ...(schemaState.payloadSchemaHash + ? { payloadSchemaHash: schemaState.payloadSchemaHash } + : {}), + ...(schemaState.responseSchemaHash + ? { responseSchemaHash: schemaState.responseSchemaHash } + : {}), + }, + ...('payload' in request && request.payload !== undefined + ? { payload: request.payload } + : {}), + } +} + +type ParsedInterruptInput = { + key: string + reason: string + message: string + expiresAt?: string + payload?: unknown +} + +type InterruptRequestFactory = ( + input: ParsedInterruptInput, + payloadIsParsed: boolean, +) => GenericInterruptRequest> + +const interruptRequestFactories = new WeakMap() +const interruptRequestInputs = new WeakMap< + object, + Readonly +>() + +/** + * Returns the schema input captured for a newly emitted request. This is + * internal because continuation state can cross a client boundary and must be + * parsed again when it returns to the server. + */ +export function getInterruptRequestInput( + request: GenericInterruptRequest>, +): Readonly { + const input = interruptRequestInputs.get(request) + if (!input) { + throw new TypeError('Interrupt request input is unavailable.') + } + return input +} + +/** + * Rebuild a request from a persisted display payload that has already passed + * the definition's payload schema. This is internal because callers must not + * bypass public input validation for new requests. + */ +export function rehydrateInterruptRequest( + definition: InterruptDefinition, + input: ParsedInterruptInput, +): GenericInterruptRequest> { + const factory = interruptRequestFactories.get(definition) + if (!factory) { + throw new TypeError('Interrupt definition request factory is unavailable.') + } + return factory(input, true) +} + +interface CanonicalSchemaJson { + json: Record + canonicalJson: string +} + +function schemaJson(schema: unknown, name: string): CanonicalSchemaJson { + if (!isStandardJSONSchema(schema)) { + throw new TypeError( + `${name} must be a Standard Schema with a JSON Schema converter.`, + ) + } + try { + const exported = schema['~standard'].jsonSchema.input({ + target: 'draft-07', + }) + if (exported === undefined) { + throw new TypeError('The exported schema is undefined.') + } + if (typeof exported === 'function') { + throw new TypeError('The exported schema must not be a function.') + } + if (Array.isArray(exported)) { + throw new TypeError('The exported schema must be a plain JSON object.') + } + if ( + !exported || + typeof exported !== 'object' || + ![Object.prototype, null].includes(Object.getPrototypeOf(exported)) + ) { + throw new TypeError('The exported schema must be a plain JSON object.') + } + const converted: Record = {} + for (const [key, value] of Object.entries(exported)) { + if (key !== '$schema') converted[key] = value + } + const canonicalJson = canonicalInterruptJson(converted) + return { json: converted, canonicalJson } + } catch (error) { + throw new TypeError( + `${name} could not export compatible JSON Schema: ${error instanceof Error ? error.message : String(error)}`, + ) + } +} + +/** Same hash the producer stamps on a first-party generic binding. */ +export function hashInterruptDefinitionSchema(schema: unknown): string { + return digestInterruptJson( + schemaJson(schema, 'Interrupt schema').canonicalJson, + ) +} + +function validateJson(value: unknown, label: string): void { + try { + canonicalInterruptJson(value) + } catch (error) { + throw new TypeError( + `${label} must be JSON-compatible: ${error instanceof Error ? error.message : String(error)}`, + ) + } +} + +function validateNonEmptyString(value: unknown, label: string): string { + if (typeof value !== 'string' || value.trim() === '') { + throw new TypeError(`${label} must be a non-empty string.`) + } + return value +} + +function validateExpiresAt(value: unknown): string { + if (typeof value !== 'string') { + throw new TypeError('Interrupt expiresAt must be a string.') + } + return value +} + +function isPromiseLike(value: unknown): value is PromiseLike { + return ( + value !== null && + (typeof value === 'object' || typeof value === 'function') && + 'then' in value && + typeof value.then === 'function' + ) +} + +function parseInterruptPayload( + schema: PortableSchema, + value: unknown, +): unknown { + if (!isStandardSchema(schema)) return value + const result = schema['~standard'].validate(value) + if (isPromiseLike(result)) { + throw new TypeError( + 'Interrupt payloadSchema validation must be synchronous.', + ) + } + if (result.issues !== undefined) { + throw new TypeError( + `Interrupt payload is invalid: ${result.issues.map((issue) => issue.message).join(' ')}`, + ) + } + return result.value +} + +export function defineInterrupt< + const TId extends string, + const TPayloadSchema extends PortableSchema, + const TResponseSchema extends PortableSchema, +>(options: { + id: TId + payloadSchema: TPayloadSchema + responseSchema: TResponseSchema +}): DefinedInterruptDefinition< + TId, + TPayloadSchema, + TResponseSchema, + InferSchemaOutput, + InferSchemaInput +> +export function defineInterrupt< + const TId extends string, + const TPayloadSchema extends PortableSchema, +>(options: { + id: TId + payloadSchema: TPayloadSchema + responseSchema?: never +}): DefinedInterruptDefinition< + TId, + TPayloadSchema, + undefined, + InferSchemaOutput, + InferSchemaInput +> +export function defineInterrupt< + const TId extends string, + const TResponseSchema extends PortableSchema, +>(options: { + id: TId + responseSchema: TResponseSchema + payloadSchema?: never +}): DefinedInterruptDefinition< + TId, + undefined, + TResponseSchema, + undefined, + undefined +> +export function defineInterrupt< + const TId extends string, + const TPayloadSchema extends PortableSchema | undefined, + const TResponseSchema extends PortableSchema | undefined, +>( + options: InterruptDefinitionOptions, +): InterruptDefinition< + TId, + TPayloadSchema, + TResponseSchema, + InferSchemaOutput, + InferSchemaInput +> { + validateNonEmptyString(options.id, 'Interrupt definition id') + const hasResponseSchema = options.responseSchema !== undefined + const responseJson = hasResponseSchema + ? schemaJson(options.responseSchema, 'responseSchema') + : undefined + const hasPayloadSchema = Object.prototype.hasOwnProperty.call( + options, + 'payloadSchema', + ) + const payloadJson = hasPayloadSchema + ? schemaJson(options.payloadSchema, 'payloadSchema') + : undefined + const schemaState: DefinitionSchemaState = { + ...(responseJson + ? { + responseSchemaCanonicalJson: responseJson.canonicalJson, + responseSchemaHash: digestInterruptJson(responseJson.canonicalJson), + } + : {}), + ...(payloadJson + ? { + payloadSchemaCanonicalJson: payloadJson.canonicalJson, + payloadSchemaHash: digestInterruptJson(payloadJson.canonicalJson), + } + : {}), + } + const definition = { + id: options.id, + payloadSchema: options.payloadSchema, + responseSchema: options.responseSchema, + interrupt(input: InterruptInput) { + return createRequest(input, false) + }, + } as InterruptDefinition< + TId, + TPayloadSchema, + TResponseSchema, + InferSchemaOutput, + InferSchemaInput + > + const parsePayload = (payload: unknown): unknown => { + const payloadSchema = options.payloadSchema + if (payloadSchema === undefined) { + throw new TypeError( + 'This interrupt definition does not accept a payload.', + ) + } + return parseInterruptPayload(payloadSchema, payload) + } + const createRequest: InterruptRequestFactory = (input, payloadIsParsed) => { + for (const key of Object.keys(input)) { + if (!['key', 'payload', 'reason', 'message', 'expiresAt'].includes(key)) { + throw new TypeError(`Interrupt input field ${key} is not allowed.`) + } + } + const key = validateNonEmptyString(input.key, 'Interrupt key') + const reason = validateNonEmptyString(input.reason, 'Interrupt reason') + const message = validateNonEmptyString(input.message, 'Interrupt message') + if ('payload' in input) { + if (!hasPayloadSchema) { + throw new TypeError( + 'This interrupt definition does not accept a payload.', + ) + } + if (input.payload !== undefined) { + validateJson(input.payload, 'Interrupt payload') + } + } + const parsedPayload = + 'payload' in input + ? payloadIsParsed + ? input.payload + : parsePayload(input.payload) + : undefined + const payload = + parsedPayload === undefined + ? undefined + : cloneAndDeepFreezeJson(parsedPayload) + const expiresAt = + input.expiresAt === undefined + ? undefined + : validateExpiresAt(input.expiresAt) + const request = Object.freeze({ + definition, + key, + ...(hasPayloadSchema && payload !== undefined ? { payload } : {}), + reason, + message, + ...(expiresAt !== undefined ? { expiresAt } : {}), + }) + if (!payloadIsParsed) { + interruptRequestInputs.set( + request, + cloneAndDeepFreezeJson({ + key, + reason, + message, + ...(expiresAt !== undefined ? { expiresAt } : {}), + ...(parsedPayload !== undefined ? { payload: parsedPayload } : {}), + }), + ) + } + return request + } + definitionSchemaState.set(definition, schemaState) + interruptRequestFactories.set(definition, createRequest) + return Object.freeze(definition) +} diff --git a/packages/ai/src/interrupt-resume.ts b/packages/ai/src/interrupt-resume.ts index ba7360c3fc..6694e452ff 100644 --- a/packages/ai/src/interrupt-resume.ts +++ b/packages/ai/src/interrupt-resume.ts @@ -25,6 +25,10 @@ import type { ChatMiddlewareConfig, ChatResumeToolState, } from './activities/chat/middleware/types' +import type { + GenericInterruptRequest, + InterruptDefinition, +} from './interrupt-definition' import type { Interrupt, RunAgentResumeItem } from './types' /** @@ -46,6 +50,10 @@ export interface PendingInterruptResumeRecord { interruptId: string payload: unknown binding: InterruptBinding + /** Present for a first-party generic interrupt. */ + genericRequest?: GenericInterruptRequest< + InterruptDefinition + > } export interface ValidateInterruptResumeBatchInput { @@ -164,6 +172,17 @@ function runtimeTool( return tools.find((tool) => tool.name === name) as RuntimeTool | undefined } +async function parseSchemaValue( + schema: unknown, + value: unknown, +): Promise<{ success: true; data: unknown } | { success: false }> { + if (!isStandardSchema(schema)) return { success: true, data: value } + const result = await validateWithStandardSchema(schema, value) + return result.success + ? { success: true, data: result.data } + : { success: false } +} + function descriptorResponseSchema( record: PendingInterruptResumeRecord, ): unknown { @@ -219,9 +238,14 @@ function validateDescriptorSchema( errors: Array, ): unknown { const schema = descriptorResponseSchema(record) + const responseSchemaHash = binding.responseSchemaHash + if (schema === undefined && responseSchemaHash === undefined) { + return undefined + } if ( schema === undefined || - schemaHash(schema) !== binding.responseSchemaHash + responseSchemaHash === undefined || + schemaHash(schema) !== responseSchemaHash ) { errors.push( interruptItemError( @@ -276,12 +300,24 @@ export async function validateInterruptResumeBatch( } } + const pendingGenerics = input.pending.filter( + (record) => record.binding.kind === 'generic', + ) + const genericBatchSatisfied = + pendingGenerics.length > 0 && + pendingGenerics.every((record) => resumeById.has(record.interruptId)) + 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) { + // Client tools that share a generic interrupt batch wait for + // `toolResume`. `continue` re-emits them; `cancel` / `stop` skip them. + if (genericBatchSatisfied && binding.kind === 'client-tool-execution') { + continue + } incomplete = true errors.push( interruptItemError( @@ -342,6 +378,28 @@ export async function validateInterruptResumeBatch( continue } if (binding.kind === 'generic') { + const genericRequest = record.genericRequest + if (genericRequest !== undefined) { + const batchIndex = binding.batchIndex + if ( + binding.definitionId !== genericRequest.definition.id || + binding.key !== genericRequest.key || + binding.interruptId !== record.interruptId || + batchIndex === undefined || + !Number.isInteger(batchIndex) || + batchIndex < 0 + ) { + errors.push( + interruptItemError( + input, + record.interruptId, + 'stale', + `Generic interrupt ${record.interruptId} has stale definition metadata.`, + { source: 'server' }, + ), + ) + } + } if (entry.status === 'cancelled') { if (entry.payload !== undefined) { errors.push( @@ -353,6 +411,16 @@ export async function validateInterruptResumeBatch( ), ) } + } else if (genericRequest !== undefined) { + await pushSchemaIssues({ + request: input, + errors, + interruptId: record.interruptId, + schema: genericRequest.definition.responseSchema, + value: entry.payload, + code: 'invalid-payload', + label: `Interrupt ${record.interruptId} payload is invalid`, + }) } else if (responseSchema !== undefined) { await pushSchemaIssues({ request: input, @@ -616,13 +684,20 @@ export async function validateInterruptResumeBatch( if (!entry) continue const binding = record.binding if (binding.kind === 'generic') { + const parsed = + entry.status === 'resolved' && record.genericRequest !== undefined + ? await parseSchemaValue( + record.genericRequest.definition.responseSchema, + entry.payload, + ) + : undefined genericInterrupts.set( record.interruptId, entry.status === 'resolved' ? { interruptId: record.interruptId, status: 'resolved', - payload: entry.payload, + payload: parsed?.success ? parsed.data : entry.payload, } : { interruptId: record.interruptId, status: 'cancelled' }, ) @@ -711,17 +786,41 @@ export function readUnopenedInterruptBinding( const interruptId = stringField(raw, 'interruptId') const responseSchemaHash = stringField(raw, 'responseSchemaHash') const expiresAt = stringField(raw, 'expiresAt') - if (!interruptId || !responseSchemaHash) return undefined + if (!interruptId || responseSchemaHash === '') return undefined const v = INTERRUPT_BINDING_VERSION if (kind === 'generic') { + const definitionId = stringField(raw, 'definitionId') + const key = stringField(raw, 'key') + const batchIndex = raw['batchIndex'] + const payloadSchemaHash = stringField(raw, 'payloadSchemaHash') + const hasFirstPartyFields = + definitionId !== undefined || + key !== undefined || + batchIndex !== undefined || + payloadSchemaHash !== undefined + if ( + hasFirstPartyFields && + (!definitionId || + !key || + typeof batchIndex !== 'number' || + !Number.isInteger(batchIndex) || + batchIndex < 0) + ) { + return undefined + } return { v, kind, interruptId, - responseSchemaHash, + ...(responseSchemaHash ? { responseSchemaHash } : {}), ...(expiresAt ? { expiresAt } : {}), + ...(definitionId ? { definitionId } : {}), + ...(key ? { key } : {}), + ...(typeof batchIndex === 'number' ? { batchIndex } : {}), + ...(payloadSchemaHash ? { payloadSchemaHash } : {}), } } + if (!responseSchemaHash) return undefined const toolName = stringField(raw, 'toolName') const toolCallId = stringField(raw, 'toolCallId') if (!toolName || !toolCallId) return undefined diff --git a/packages/ai/src/interrupts.ts b/packages/ai/src/interrupts.ts index e11246e0ec..59f30057c8 100644 --- a/packages/ai/src/interrupts.ts +++ b/packages/ai/src/interrupts.ts @@ -84,12 +84,15 @@ interface InterruptBindingBase { interruptId: string interruptedRunId: string generation: number - responseSchemaHash: string expiresAt?: string } +interface ResponseSchemaInterruptBindingBase extends InterruptBindingBase { + responseSchemaHash: string +} + export type InterruptBinding = - | (InterruptBindingBase & { + | (ResponseSchemaInterruptBindingBase & { kind: 'tool-approval' toolName: string toolCallId: string @@ -97,7 +100,7 @@ export type InterruptBinding = inputSchemaHash: string approvalSchemaHash: string }) - | (InterruptBindingBase & { + | (ResponseSchemaInterruptBindingBase & { kind: 'client-tool-execution' toolName: string toolCallId: string @@ -105,6 +108,13 @@ export type InterruptBinding = }) | (InterruptBindingBase & { kind: 'generic' + /** Omitted when the generic interrupt accepts an unvalidated response. */ + responseSchemaHash?: string + /** Present only for a first-party generic interrupt. */ + definitionId?: string + key?: string + batchIndex?: number + payloadSchemaHash?: string }) export type UnopenedInterruptBinding = InterruptBinding extends infer TBinding diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index e80d02175a..b9777606f1 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -1023,8 +1023,7 @@ export interface TextOptions< /** * 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 - * upstream-native interrupts. + * A first-party generic item carries the original request in `metadata`. */ resume?: Array @@ -1121,7 +1120,10 @@ export type Interrupt = AGUIInterrupt export type RunFinishedOutcome = AGUIRunFinishedOutcome -export type RunAgentResumeItem = AGUIResumeEntry +export type RunAgentResumeItem = AGUIResumeEntry & { + /** AG-UI resume metadata. First-party generic requests ride here. */ + metadata?: Record +} /** * Emitted when a run completes successfully. diff --git a/packages/ai/src/utilities/chat-params.ts b/packages/ai/src/utilities/chat-params.ts index 36a37e67da..de0f928f89 100644 --- a/packages/ai/src/utilities/chat-params.ts +++ b/packages/ai/src/utilities/chat-params.ts @@ -2,7 +2,6 @@ import { AGUIError } from '@ag-ui/core' import type { Context as AGUIContext, Message as AGUIMessage, - ResumeEntry as AGUIResumeEntry, Role as AGUIRole, } from '@ag-ui/core' import type { @@ -173,20 +172,29 @@ function validateContext(value: unknown, index: number): AGUIContext { } } -function validateResumeEntry(value: unknown, index: number): AGUIResumeEntry { +function validateResumeEntry( + value: unknown, + index: number, +): RunAgentResumeItem { const at = `resume[${index}]` if (!isRecord(value)) invalidBody(`${at} must be an object`) const status = value.status if (status !== 'resolved' && status !== 'cancelled') { invalidBody(`${at}.status must be "resolved" or "cancelled"`) } - const entry: AGUIResumeEntry = { + const entry: RunAgentResumeItem = { interruptId: requireString(value.interruptId, `${at}.interruptId`), status, } // Omit the key entirely when absent, matching the optional-field shape the // schema produced. if (value.payload !== undefined) entry.payload = value.payload + if (value.metadata !== undefined) { + if (!isRecord(value.metadata)) { + invalidBody(`${at}.metadata must be an object`) + } + entry.metadata = value.metadata + } return entry } diff --git a/packages/ai/tests/chat-params.test.ts b/packages/ai/tests/chat-params.test.ts index 0c46ea8b87..b2e3cb21d9 100644 --- a/packages/ai/tests/chat-params.test.ts +++ b/packages/ai/tests/chat-params.test.ts @@ -269,6 +269,49 @@ describe('chatParamsFromRequestBody — RunAgentInput validation', () => { expect('payload' in result.resume![0]!).toBe(false) }) + it('keeps resume metadata for generic continuation', async () => { + const metadata = { + 'tanstack:interruptContinuation': { + v: 1, + definitionId: 'review', + key: 'one', + batchIndex: 0, + reason: 'review', + message: 'Review', + }, + } + const result = await chatParamsFromRequestBody({ + ...base, + messages: [], + resume: [ + { + interruptId: 'i1', + status: 'resolved', + payload: { approved: true }, + metadata, + }, + ], + }) + expect(result.resume).toEqual([ + { + interruptId: 'i1', + status: 'resolved', + payload: { approved: true }, + metadata, + }, + ]) + }) + + it('rejects non-object resume metadata', async () => { + await expect( + chatParamsFromRequestBody({ + ...base, + messages: [], + resume: [{ interruptId: 'i1', status: 'cancelled', metadata: 'nope' }], + }), + ).rejects.toThrow(/resume\[0\]\.metadata/) + }) + it('defaults forwardedProps to {} and rejects a non-object one', async () => { const result = await chatParamsFromRequestBody(withMessages([])) expect(result.forwardedProps).toEqual({}) diff --git a/packages/ai/tests/chat.test.ts b/packages/ai/tests/chat.test.ts index 0e362875c9..340a780087 100644 --- a/packages/ai/tests/chat.test.ts +++ b/packages/ai/tests/chat.test.ts @@ -1,6 +1,11 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { chat, createChatOptions } from '../src/activities/chat/index' +import { defineInterrupt } from '../src/interrupt-definition' +import { + genericInterruptContinuationFromDescriptor, + wrapGenericInterruptContinuation, +} from '../src/generic-interrupt-continuation' import { defineChatMiddleware } from '../src/activities/chat/middleware/define' import { DISCOVERY_TOOL_NAME } from '../src/activities/chat/tools/lazy-tool-manager' import { EventType } from '../src/types' @@ -3791,4 +3796,504 @@ describe('chat()', () => { expect((resultChunks[0] as any).toolCallId).toBeDefined() }) }) + + describe('generic interrupts', () => { + it('emits one afterModel interrupt terminal and stores continuation state', async () => { + const review = defineInterrupt({ + id: 'review-plan', + responseSchema: z.object({ approved: z.boolean() }), + }) + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.textStart(), + ev.textContent('Plan'), + ev.textEnd(), + ev.runFinished('stop'), + ], + ], + }) + + const chunks = await collectChunks( + chat({ + adapter, + interrupts: [review], + middleware: [ + defineChatMiddleware({ + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterModel') return + return { + interrupts: [ + review.interrupt({ + key: 'turn-1', + reason: 'review', + message: 'Review the plan', + }), + ], + } + }, + }), + ], + messages: [{ role: 'user', content: 'Make a plan' }], + }) as AsyncIterable, + ) + + const terminal = expectSingleRunFinished(chunks) + expect(terminal.outcome).toMatchObject({ + type: 'interrupt', + interrupts: [ + { + reason: 'review', + message: 'Review the plan', + }, + ], + }) + expect( + chunks.findIndex((chunk) => chunk.type === EventType.MESSAGES_SNAPSHOT), + ).toBeLessThan( + chunks.findIndex((chunk) => chunk.type === EventType.RUN_FINISHED), + ) + expect( + chunks.some((chunk) => chunk.type === EventType.STATE_SNAPSHOT), + ).toBe(false) + }) + + it('starts a synthetic run before a beforeModel interrupt', async () => { + const review = defineInterrupt({ + id: 'before-model-lifecycle', + responseSchema: z.object({ approved: z.boolean() }), + }) + const { adapter, calls } = createMockAdapter({ iterations: [] }) + + const chunks = await collectChunks( + chat({ + adapter, + interrupts: [review], + runId: 'before-model-run', + middleware: [ + defineChatMiddleware({ + onInterruptBoundary(ctx) { + if (ctx.phase !== 'beforeModel') return + return { + interrupts: [ + review.interrupt({ + key: 'one', + reason: 'review', + message: 'Review', + }), + ], + } + }, + }), + ], + messages: [{ role: 'user', content: 'Start' }], + }) as AsyncIterable, + ) + + expect(calls).toHaveLength(0) + expect(chunks.map((chunk) => chunk.type)).toEqual([ + EventType.RUN_STARTED, + EventType.MESSAGES_SNAPSHOT, + EventType.RUN_FINISHED, + ]) + expect(expectSingleRunFinished(chunks).outcome?.type).toBe('interrupt') + }) + + it('pauses before tools without executing them', async () => { + const decision = defineInterrupt({ + id: 'before-tools', + responseSchema: z.object({ proceed: z.boolean() }), + }) + const execute = vi.fn(() => ({ ok: true })) + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.textStart(), + ev.toolStart('tool-1', 'write'), + ev.toolArgs('tool-1', '{}'), + ev.toolEnd('tool-1', 'write', { input: {} }), + ev.runFinished('tool_calls'), + ], + ], + }) + + const chunks = await collectChunks( + chat({ + adapter, + interrupts: [decision], + middleware: [ + defineChatMiddleware({ + onInterruptBoundary(ctx) { + if (ctx.phase !== 'beforeTools') return + return { + interrupts: [ + decision.interrupt({ + key: 'tools', + reason: 'review', + message: 'Run tools?', + }), + ], + } + }, + }), + ], + messages: [{ role: 'user', content: 'Write it' }], + tools: [serverTool('write', execute)], + }) as AsyncIterable, + ) + + expect(execute).not.toHaveBeenCalled() + expect(expectSingleRunFinished(chunks).outcome).toMatchObject({ + type: 'interrupt', + }) + }) + + it('stops after afterTools when resolution returns toolResume stop', async () => { + const review = defineInterrupt({ + id: 'after-tools-stop', + responseSchema: z.object({ approved: z.boolean() }), + }) + const execute = vi.fn(() => ({ ok: true })) + const { adapter, calls } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('tool-1', 'write'), + ev.toolArgs('tool-1', '{}'), + ev.toolEnd('tool-1', 'write', { input: {} }), + ev.runFinished('tool_calls'), + ], + ], + }) + const middleware = defineChatMiddleware({ + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterTools') return + return { + interrupts: [ + review.interrupt({ + key: 'after-tools', + reason: 'review', + message: 'Review the tool result', + }), + ], + } + }, + onInterruptResolution(_ctx, resolutions) { + for (const resolution of resolutions.for(review)) { + if ( + resolution.status === 'resolved' && + !resolution.response.approved + ) { + return { toolResume: 'stop' } + } + } + return { toolResume: 'continue' } + }, + }) + + const first = await collectChunks( + chat({ + adapter, + interrupts: [review], + middleware: [middleware], + messages: [{ role: 'user', content: 'Write it' }], + tools: [serverTool('write', execute)], + threadId: 'thread-after-tools', + runId: 'run-after-tools', + }) as AsyncIterable, + ) + const interrupt = expectSingleRunFinished(first).outcome + if (interrupt?.type !== 'interrupt') { + throw new Error('Expected afterTools interrupt') + } + const paused = interrupt.interrupts[0] + if (!paused) throw new Error('Expected interrupt id') + const continuation = genericInterruptContinuationFromDescriptor(paused) + if (!continuation) { + throw new Error('Expected generic continuation metadata') + } + + const resume = await collectChunks( + chat({ + adapter, + interrupts: [review], + middleware: [middleware], + messages: [ + { role: 'user', content: 'Write it' }, + { + role: 'assistant', + content: null, + toolCalls: [ + { + id: 'tool-1', + type: 'function', + function: { name: 'write', arguments: '{}' }, + }, + ], + }, + { role: 'tool', content: '{"ok":true}', toolCallId: 'tool-1' }, + ], + tools: [serverTool('write', execute)], + threadId: 'thread-after-tools', + runId: 'run-after-tools-resume', + parentRunId: 'run-after-tools', + resume: [ + { + interruptId: paused.id, + status: 'resolved', + payload: { approved: false }, + metadata: wrapGenericInterruptContinuation(continuation), + }, + ], + }) as AsyncIterable, + ) + + expect(execute).toHaveBeenCalledTimes(1) + expect(calls).toHaveLength(1) + const terminal = expectSingleRunFinished(resume) + expect(terminal.outcome).toEqual({ type: 'success' }) + expect(terminal.finishReason).toBe('stop') + }) + + it('runs onFinish when resolution returns toolResume stop', async () => { + const review = defineInterrupt({ + id: 'after-tools-stop-finish', + responseSchema: z.object({ approved: z.boolean() }), + }) + const onFinish = vi.fn() + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('tool-1', 'write'), + ev.toolArgs('tool-1', '{}'), + ev.toolEnd('tool-1', 'write', { input: {} }), + ev.runFinished('tool_calls'), + ], + ], + }) + const reviewMiddleware = defineChatMiddleware({ + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterTools') return + return { + interrupts: [ + review.interrupt({ + key: 'after-tools', + reason: 'review', + message: 'Review the tool result', + }), + ], + } + }, + onInterruptResolution() { + return { toolResume: 'stop' } + }, + }) + const first = await collectChunks( + chat({ + adapter, + interrupts: [review], + middleware: [reviewMiddleware], + messages: [{ role: 'user', content: 'Write it' }], + tools: [serverTool('write', () => ({ ok: true }))], + threadId: 'thread-stop-finish', + runId: 'run-stop-finish', + }) as AsyncIterable, + ) + const paused = expectSingleRunFinished(first).outcome + if (paused?.type !== 'interrupt') { + throw new Error('Expected afterTools interrupt') + } + const interrupt = paused.interrupts[0] + if (!interrupt) throw new Error('Expected interrupt id') + const continuation = genericInterruptContinuationFromDescriptor(interrupt) + if (!continuation) { + throw new Error('Expected generic continuation metadata') + } + + await collectChunks( + chat({ + adapter, + interrupts: [review], + middleware: [ + reviewMiddleware, + defineChatMiddleware({ name: 'observe-finish', onFinish }), + ], + messages: [ + { role: 'user', content: 'Write it' }, + { + role: 'assistant', + content: null, + toolCalls: [ + { + id: 'tool-1', + type: 'function', + function: { name: 'write', arguments: '{}' }, + }, + ], + }, + { role: 'tool', content: '{"ok":true}', toolCallId: 'tool-1' }, + ], + tools: [serverTool('write', () => ({ ok: true }))], + threadId: 'thread-stop-finish', + runId: 'run-stop-finish-resume', + parentRunId: 'run-stop-finish', + resume: [ + { + interruptId: interrupt.id, + status: 'resolved', + payload: { approved: false }, + metadata: wrapGenericInterruptContinuation(continuation), + }, + ], + }) as AsyncIterable, + ) + + expect(onFinish).toHaveBeenCalledOnce() + expect(onFinish.mock.calls[0]?.[1]).toMatchObject({ + finishReason: 'stop', + }) + }) + + it('keeps pending tool calls in the afterModel snapshot', async () => { + const review = defineInterrupt({ + id: 'after-model-tools', + responseSchema: z.object({ approved: z.boolean() }), + }) + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('tool-1', 'write'), + ev.toolArgs('tool-1', '{}'), + ev.toolEnd('tool-1', 'write', { input: {} }), + ev.runFinished('tool_calls'), + ], + ], + }) + + const chunks = await collectChunks( + chat({ + adapter, + interrupts: [review], + middleware: [ + defineChatMiddleware({ + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterModel') return + return { + interrupts: [ + review.interrupt({ + key: 'after-model', + reason: 'review', + message: 'Review before tools', + }), + ], + } + }, + }), + ], + messages: [{ role: 'user', content: 'Write it' }], + tools: [serverTool('write', () => ({ ok: true }))], + }) as AsyncIterable, + ) + + const snapshot = chunks.find( + (chunk) => chunk.type === EventType.MESSAGES_SNAPSHOT, + ) + if (!snapshot || snapshot.type !== EventType.MESSAGES_SNAPSHOT) { + throw new Error('Expected messages snapshot') + } + const assistant = snapshot.messages.at(-1) + expect(assistant).toMatchObject({ + role: 'assistant', + toolCalls: [ + { + id: 'tool-1', + function: { name: 'write' }, + }, + ], + }) + }) + + it('keeps tool approvals in the afterTools generic batch', async () => { + const review = defineInterrupt({ + id: 'after-tools-with-approval', + responseSchema: z.object({ approved: z.boolean() }), + }) + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('tool-1', 'delete'), + ev.toolArgs('tool-1', '{}'), + ev.toolEnd('tool-1', 'delete', { input: {} }), + ev.runFinished('tool_calls'), + ], + ], + }) + + const chunks = await collectChunks( + chat({ + adapter, + interrupts: [review], + middleware: [ + defineChatMiddleware({ + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterTools') return + return { + interrupts: [ + review.interrupt({ + key: 'after-tools', + reason: 'review', + message: 'Review after tools', + }), + ], + } + }, + }), + ], + messages: [{ role: 'user', content: 'Delete it' }], + tools: [ + { + ...serverTool('delete', () => ({ ok: true })), + needsApproval: true, + }, + ], + }) as AsyncIterable, + ) + + const terminal = expectSingleRunFinished(chunks) + expect(terminal.outcome?.type).toBe('interrupt') + if (terminal.outcome?.type !== 'interrupt') return + const reasons = terminal.outcome.interrupts.map( + (interrupt) => interrupt.reason, + ) + expect(reasons).toContain('tool_call') + expect(reasons).toContain('review') + }) + + it('rejects duplicate interrupt definition ids before adapter work', () => { + const first = defineInterrupt({ + id: 'duplicate-chat-id', + responseSchema: z.object({ ok: z.boolean() }), + }) + const second = defineInterrupt({ + id: 'duplicate-chat-id', + responseSchema: z.object({ ok: z.boolean() }), + }) + const { adapter } = createMockAdapter({ iterations: [] }) + + expect(() => + Reflect.apply(chat, undefined, [ + { + adapter, + interrupts: [first, second], + messages: [{ role: 'user', content: 'Hello' }], + }, + ]), + ).toThrow('Duplicate interrupt definition id: duplicate-chat-id') + }) + }) }) diff --git a/packages/ai/tests/generic-interrupt-continuation.test.ts b/packages/ai/tests/generic-interrupt-continuation.test.ts new file mode 100644 index 0000000000..c81396d104 --- /dev/null +++ b/packages/ai/tests/generic-interrupt-continuation.test.ts @@ -0,0 +1,86 @@ +import { describe, expect, it } from 'vitest' +import { z } from 'zod' +import { + INTERRUPT_CONTINUATION_METADATA_KEY, + genericInterruptContinuationFromDescriptor, + readGenericInterruptContinuation, + wrapGenericInterruptContinuation, +} from '../src/generic-interrupt-continuation' +import { + INTERRUPT_PAYLOAD_METADATA_KEY, + createInterruptBinding, + defineInterrupt, +} from '../src/interrupt-definition' +import { INTERRUPT_BINDING_METADATA_KEY } from '../src/interrupt-resume' +import { INTERRUPT_BINDING_VERSION } from '../src/interrupts' + +describe('generic interrupt continuation', () => { + it('treats missing metadata as absent', () => { + expect(readGenericInterruptContinuation(undefined)).toEqual({ + status: 'absent', + }) + expect(readGenericInterruptContinuation({})).toEqual({ status: 'absent' }) + }) + + it('rejects a present key with a bad shape', () => { + expect( + readGenericInterruptContinuation({ + [INTERRUPT_CONTINUATION_METADATA_KEY]: 'nope', + }), + ).toMatchObject({ status: 'invalid' }) + expect( + readGenericInterruptContinuation({ + [INTERRUPT_CONTINUATION_METADATA_KEY]: { v: 2 }, + }), + ).toMatchObject({ status: 'invalid' }) + }) + + it('round-trips a first-party interrupt descriptor', () => { + const review = defineInterrupt({ + id: 'review-plan', + payloadSchema: z.object({ title: z.string() }), + responseSchema: z.object({ approved: z.boolean() }), + }) + const request = review.interrupt({ + key: 'turn-1', + reason: 'review', + message: 'Review the plan', + payload: { title: 'Ship it' }, + }) + const emission = createInterruptBinding(request, { batchIndex: 0 }) + const interrupt = { + id: 'generic-1', + reason: request.reason, + message: request.message, + metadata: { + [INTERRUPT_BINDING_METADATA_KEY]: { + v: INTERRUPT_BINDING_VERSION, + kind: 'generic', + interruptId: 'generic-1', + definitionId: emission.descriptor.definitionId, + key: emission.descriptor.key, + batchIndex: 0, + responseSchemaHash: emission.descriptor.responseSchemaHash, + payloadSchemaHash: emission.descriptor.payloadSchemaHash, + }, + [INTERRUPT_PAYLOAD_METADATA_KEY]: { title: 'Ship it' }, + }, + } + const continuation = genericInterruptContinuationFromDescriptor(interrupt) + expect(continuation).toMatchObject({ + v: 1, + definitionId: 'review-plan', + key: 'turn-1', + batchIndex: 0, + reason: 'review', + message: 'Review the plan', + payload: { title: 'Ship it' }, + }) + if (!continuation) throw new Error('Expected continuation') + expect( + readGenericInterruptContinuation( + wrapGenericInterruptContinuation(continuation), + ), + ).toEqual({ status: 'ok', value: continuation }) + }) +}) diff --git a/packages/ai/tests/interrupt-resume.test.ts b/packages/ai/tests/interrupt-resume.test.ts index 0a8195354a..7b7cf91bdf 100644 --- a/packages/ai/tests/interrupt-resume.test.ts +++ b/packages/ai/tests/interrupt-resume.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from 'vitest' import { z } from 'zod' import { + defineInterrupt, hashSchemaInput, normalizeApprovalSchema, toolDefinition, @@ -267,4 +268,131 @@ describe('validateInterruptResumeBatch', () => { ), ).toBe(true) }) + + it('allows a missing client-tool resume when every generic in the batch is answered', async () => { + const review = defineInterrupt({ + id: 'review-plan', + payloadSchema: z.object({ title: z.string() }), + responseSchema: z.object({ + approved: z.boolean(), + note: z.string(), + }), + }) + const request = review.interrupt({ + key: 'one', + reason: 'review', + message: 'Review', + payload: { title: 'Plan' }, + }) + const renderDef = toolDefinition({ + name: 'render_review', + description: 'Render a review', + inputSchema: z.object({ reviewId: z.string() }), + outputSchema: z.object({ rendered: z.boolean() }), + }) + const genericBinding: Extract = { + v: INTERRUPT_BINDING_VERSION, + kind: 'generic', + interruptId: 'generic-1', + interruptedRunId: 'run-1', + generation: 0, + definitionId: 'review-plan', + key: 'one', + batchIndex: 0, + } + const clientBinding: Extract< + InterruptBinding, + { kind: 'client-tool-execution' } + > = { + v: INTERRUPT_BINDING_VERSION, + kind: 'client-tool-execution', + interruptId: 'client_tool_call-2', + interruptedRunId: 'run-1', + generation: 0, + toolName: 'render_review', + toolCallId: 'call-2', + outputSchemaHash: hashSchemaInput(renderDef.outputSchema), + responseSchemaHash: 'sha256:client-tool', + } + const result = await validateInterruptResumeBatch({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 0, + pending: [ + { + interruptId: genericBinding.interruptId, + payload: request, + binding: genericBinding, + genericRequest: request, + }, + { + interruptId: clientBinding.interruptId, + payload: {}, + binding: clientBinding, + }, + ], + resume: [ + { + interruptId: genericBinding.interruptId, + status: 'resolved', + payload: { approved: true, note: 'ok' }, + }, + ], + tools: [transfer, renderDef.client(async () => ({ rendered: true }))], + }) + expect(result.errors).toEqual([]) + expect(result.resumeToolState).toBeDefined() + expect(result.resumeToolState?.genericInterrupts?.get('generic-1')).toEqual( + { + interruptId: 'generic-1', + status: 'resolved', + payload: { approved: true, note: 'ok' }, + }, + ) + expect(result.resumeToolState?.clientToolResults?.size).toBe(0) + }) + + it('still requires a client-tool resume when the batch has no generic interrupt', async () => { + const renderDef = toolDefinition({ + name: 'render_review', + description: 'Render a review', + inputSchema: z.object({ reviewId: z.string() }), + outputSchema: z.object({ rendered: z.boolean() }), + }) + const clientBinding: Extract< + InterruptBinding, + { kind: 'client-tool-execution' } + > = { + v: INTERRUPT_BINDING_VERSION, + kind: 'client-tool-execution', + interruptId: 'client_tool_call-2', + interruptedRunId: 'run-1', + generation: 0, + toolName: 'render_review', + toolCallId: 'call-2', + outputSchemaHash: hashSchemaInput(renderDef.outputSchema), + responseSchemaHash: 'sha256:client-tool', + } + const result = await validateInterruptResumeBatch({ + threadId: 'thread-1', + interruptedRunId: 'run-1', + generation: 0, + pending: [ + { + interruptId: clientBinding.interruptId, + payload: {}, + binding: clientBinding, + }, + ], + resume: [], + tools: [renderDef.client(async () => ({ rendered: true }))], + }) + expect( + result.errors.some( + (error) => + error.code === 'incomplete-batch' || + error.code === 'unknown-interrupt', + ), + ).toBe(true) + }) }) diff --git a/packages/ai/tests/interrupts-types.test-d.ts b/packages/ai/tests/interrupts-types.test-d.ts index 3b61016d16..ce37a6d3f6 100644 --- a/packages/ai/tests/interrupts-types.test-d.ts +++ b/packages/ai/tests/interrupts-types.test-d.ts @@ -1,6 +1,9 @@ import { expectTypeOf } from 'vitest' import { z } from 'zod' -import { toolDefinition } from '../src' +import { defineInterrupt, toolDefinition } from '../src' +import { defineInterrupt as defineClientInterrupt } from '../src/client' +import type { GenericInterruptRequest, InferSchemaType } from '../src' +import type { GenericInterruptRequest as ClientGenericInterruptRequest } from '../src/client' import type { ApprovalCapabilityOf, ApprovalSchemaOf, @@ -72,3 +75,144 @@ expectTypeOf().toEqualTypeOf< expectTypeOf().toEqualTypeOf< ReadonlySet | undefined >() + +const interruptWithPayload = defineInterrupt({ + id: 'with-payload', + payloadSchema: z.object({ label: z.string() }), + responseSchema: z.object({ accepted: z.boolean() }), +}) +const interruptWithoutPayload = defineInterrupt({ + id: 'without-payload', + responseSchema: z.object({ accepted: z.boolean() }), +}) +const transformingInterrupt = defineInterrupt({ + id: 'transforming-payload', + payloadSchema: z.string().transform((value) => value.length), + responseSchema: z.object({ accepted: z.boolean() }), +}) +const payloadRequest = interruptWithPayload.interrupt({ + key: 'one', + payload: { label: 'hello' }, + reason: 'test', + message: 'Test', +}) +expectTypeOf(payloadRequest.payload).toEqualTypeOf< + { label: string } | undefined +>() +const transformingRequest = transformingInterrupt.interrupt({ + key: 'transforming', + payload: 'five', + reason: 'test', + message: 'Test', +}) +expectTypeOf(transformingRequest.payload).toEqualTypeOf() +// @ts-expect-error Payload must use the schema input type. +transformingInterrupt.interrupt({ + key: 'invalid-transforming', + payload: 5, + reason: 'test', + message: 'Test', +}) +const nestedTransformingRequest = transformingRequest.definition.interrupt({ + key: 'nested-transforming', + payload: 'six', + reason: 'test', + message: 'Test', +}) +expectTypeOf(nestedTransformingRequest.payload).toEqualTypeOf< + number | undefined +>() +// @ts-expect-error Nested definitions keep the schema input payload type. +transformingRequest.definition.interrupt({ + key: 'invalid-nested-transforming', + payload: 6, + reason: 'test', + message: 'Test', +}) +// @ts-expect-error Request definition is readonly. +payloadRequest.definition = interruptWithoutPayload +// @ts-expect-error Request key is readonly. +payloadRequest.key = 'changed' +// @ts-expect-error Request reason is readonly. +payloadRequest.reason = 'changed' +// @ts-expect-error Request message is readonly. +payloadRequest.message = 'changed' +// @ts-expect-error Request payload is readonly. +payloadRequest.payload = { label: 'changed' } +const noPayloadRequest = interruptWithoutPayload.interrupt({ + key: 'two', + reason: 'test', + message: 'Test', +}) +const clientInterrupt = defineClientInterrupt({ + id: 'client-interrupt', + responseSchema: z.object({ ok: z.boolean() }), +}) +expectTypeOf< + InferSchemaType +>().toEqualTypeOf<{ label: string }>() +expectTypeOf().toEqualTypeOf< + typeof interruptWithPayload.payloadSchema +>() +expectTypeOf().toEqualTypeOf<'with-payload'>() +expectTypeOf< + keyof GenericInterruptRequest +>().toEqualTypeOf< + 'definition' | 'key' | 'payload' | 'reason' | 'message' | 'expiresAt' +>() +expectTypeOf< + keyof GenericInterruptRequest +>().toEqualTypeOf<'definition' | 'key' | 'reason' | 'message' | 'expiresAt'>() +const clientRequest = clientInterrupt.interrupt({ + key: 'client', + reason: 'test', + message: 'Test', +}) +expectTypeOf< + keyof ClientGenericInterruptRequest +>().toEqualTypeOf<'definition' | 'key' | 'reason' | 'message' | 'expiresAt'>() +const extraIdInput = { + key: 'extra-id', + reason: 'test', + message: 'Test', + id: 'ag-ui-id', +} +// @ts-expect-error Unknown id input field is not allowed. +interruptWithPayload.interrupt(extraIdInput) +const extraRunInput = { + key: 'extra-run', + reason: 'test', + message: 'Test', + runId: 'run-id', +} +// @ts-expect-error Unknown runId input field is not allowed. +interruptWithPayload.interrupt(extraRunInput) +const extraResponseInput = { + key: 'extra-response', + reason: 'test', + message: 'Test', + response: { accepted: true }, +} +// @ts-expect-error Unknown response input field is not allowed. +interruptWithPayload.interrupt(extraResponseInput) +const missingResponseOptions = { id: 'missing-response' } +// @ts-expect-error responseSchema is required. +defineInterrupt(missingResponseOptions) +// @ts-expect-error A definition without a payload schema has no payload property. +noPayloadRequest.payload +interruptWithoutPayload.interrupt({ + key: 'three', + // @ts-expect-error A payload key is forbidden without a payload schema. + payload: undefined, + reason: 'test', + message: 'Test', +}) + +// @ts-expect-error Internal helpers are not exported from the root barrel. +import { createInterruptBinding } from '../src' +// @ts-expect-error Internal constants are not exported from the root barrel. +import { INTERRUPT_BINDING_KIND } from '../src' +// @ts-expect-error Internal descriptor types are not exported from the root barrel. +import type { InterruptBindingDescriptor } from '../src' +// @ts-expect-error Internal helpers are not exported from the client barrel. +import { createInterruptBinding as createClientInterruptBinding } from '../src/client' diff --git a/packages/ai/tests/interrupts.test.ts b/packages/ai/tests/interrupts.test.ts index f9ff696ccc..0811288568 100644 --- a/packages/ai/tests/interrupts.test.ts +++ b/packages/ai/tests/interrupts.test.ts @@ -1,5 +1,10 @@ import { describe, expect, expectTypeOf, it } from 'vitest' +import { z } from 'zod' import { EventType } from '@ag-ui/core' +import { + createInterruptBinding, + defineInterrupt, +} from '../src/interrupt-definition' import { hashSchemaInput, normalizeApprovalSchema, @@ -28,6 +33,76 @@ import type { TextOptions, } from '../src/types' +describe('first-party interrupt definitions', () => { + const approval = defineInterrupt({ + id: 'approval', + payloadSchema: z.object({ amount: z.number() }), + responseSchema: z.object({ approved: z.boolean() }), + }) + + it('creates a portable request and deterministic pre-emission binding data', () => { + const request = approval.interrupt({ + key: 'payment-1', + payload: { amount: 10 }, + reason: 'tool_call', + message: 'Approve payment?', + }) + expect(request).toMatchObject({ definition: approval, key: 'payment-1' }) + expect(request.payload).toEqual({ amount: 10 }) + const binding = createInterruptBinding(request) + expect(binding.descriptor.definitionId).toBe('approval') + expect(binding.descriptor.key).toBe('payment-1') + expect(binding.payload).toEqual({ amount: 10 }) + }) + + it('accepts undefined for an optional payload schema', () => { + const definition = defineInterrupt({ + id: 'optional-note', + payloadSchema: z.string().optional(), + responseSchema: z.object({ ok: z.boolean() }), + }) + expect( + definition.interrupt({ + key: 'one', + reason: 'test', + message: 'Test', + payload: undefined, + }).payload, + ).toBeUndefined() + }) + + it('rejects a payload when the definition has no payload schema', () => { + const definition = defineInterrupt({ + id: 'simple', + responseSchema: z.object({ ok: z.boolean() }), + }) + expect(() => + Reflect.apply(definition.interrupt, definition, [ + { + key: 'simple-2', + payload: undefined, + reason: 'test', + message: 'Test', + }, + ]), + ).toThrow() + }) + + it('rejects duplicate-looking extra request fields at runtime', () => { + expect(() => + Reflect.apply(approval.interrupt, approval, [ + { + key: 'filtered', + payload: { amount: 1 }, + reason: 'test', + message: 'Test', + id: 'ag-ui-id', + }, + ]), + ).toThrow(/Interrupt input field id is not allowed/) + }) +}) + describe('AG-UI interrupt protocol types', () => { it('allows RUN_FINISHED success, interrupt, and legacy outcomes', () => { const success = { diff --git a/packages/ai/tests/middleware-interrupt-types.test-d.ts b/packages/ai/tests/middleware-interrupt-types.test-d.ts new file mode 100644 index 0000000000..d73f0a03c5 --- /dev/null +++ b/packages/ai/tests/middleware-interrupt-types.test-d.ts @@ -0,0 +1,438 @@ +import { expectTypeOf } from 'vitest' +import { z } from 'zod' +import type { StandardJSONSchemaV1 } from '@standard-schema/spec' +import { defineInterrupt } from '../src' +import { + chat, + createChatOptions, + defineChatMiddleware, + INTERRUPT_BOUNDARY_PHASES, + INTERRUPT_TOOL_RESUMES, + toolDefinition, +} from '../src' +import type { AnyTextAdapter } from '../src' +import type { + ChatMiddleware, + DefinedChatMiddleware, + GenericInterruptRequest, + GenericInterruptResolution, + InterruptBoundaryPhase, + InterruptResolutionCollection, + InterruptResolutionResult, + InterruptToolResume, +} from '../src' +import { createChatMiddleware } from '../src' + +expectTypeOf(INTERRUPT_BOUNDARY_PHASES).toEqualTypeOf< + readonly ['beforeModel', 'afterModel', 'beforeTools', 'afterTools'] +>() +expectTypeOf().toEqualTypeOf< + (typeof INTERRUPT_BOUNDARY_PHASES)[number] +>() +expectTypeOf(INTERRUPT_TOOL_RESUMES).toEqualTypeOf< + readonly ['continue', 'cancel', 'stop'] +>() +expectTypeOf().toEqualTypeOf< + (typeof INTERRUPT_TOOL_RESUMES)[number] +>() + +const standardDefinition = defineInterrupt({ + id: 'standard', + responseSchema: z.object({ approved: z.boolean() }), +}) +const duplicateIdDefinition = defineInterrupt({ + id: 'standard', + responseSchema: z.object({ approved: z.boolean() }), +}) +declare const runtimeInterruptId: string +const runtimeIdDefinition = defineInterrupt({ + id: runtimeInterruptId, + responseSchema: z.object({ approved: z.boolean() }), +}) + +declare const jsonResponseSchema: StandardJSONSchemaV1< + unknown, + { accepted: boolean } +> +const jsonDefinition = defineInterrupt({ + id: 'json', + responseSchema: jsonResponseSchema, +}) +declare const unrelatedDefinition: ReturnType + +type Definitions = typeof standardDefinition | typeof jsonDefinition +declare const collection: InterruptResolutionCollection + +const standardRequest = standardDefinition.interrupt({ + key: 'standard', + reason: 'test', + message: 'Standard', +}) +const jsonRequest = jsonDefinition.interrupt({ + key: 'json', + reason: 'test', + message: 'JSON', +}) + +const contextToolDefinition = toolDefinition({ + name: 'context-tool', + description: 'Context tool', + inputSchema: z.object({ value: z.string() }), +}) +const contextTool = contextToolDefinition.server( + async (_args, context: { context: { toolFlag: boolean } }) => + context.context.toolFlag, +) +const unregisteredRequest = unrelatedDefinition.interrupt({ + key: 'x', + reason: 'x', + message: 'x', +}) + +const validStandardResolution: GenericInterruptResolution = { + request: standardRequest, + status: 'resolved', + response: { approved: true }, +} +const validJsonResolution: GenericInterruptResolution = { + request: jsonRequest, + status: 'resolved', + response: { accepted: true }, +} +void validStandardResolution +void validJsonResolution + +const validCancelledResolution: GenericInterruptResolution< + typeof standardDefinition +> = { + request: standardRequest, + status: 'cancelled', +} +void validCancelledResolution + +// @ts-expect-error Resolved resolutions must include the exact response. +const resolvedWithoutResponse: GenericInterruptResolution< + typeof standardDefinition +> = { + request: standardRequest, + status: 'resolved', +} +void resolvedWithoutResponse + +// @ts-expect-error Cancelled resolutions cannot include a response. +const cancelledWithResponse: GenericInterruptResolution< + typeof standardDefinition +> = { + request: standardRequest, + status: 'cancelled', + response: { approved: true }, +} +void cancelledWithResponse + +// @ts-expect-error A standard request cannot carry the JSON definition response. +const invalidResolutionPair: GenericInterruptResolution = { + request: standardRequest, + status: 'resolved', + response: { accepted: true }, +} +void invalidResolutionPair + +// @ts-expect-error A collection rejects definitions outside its registered union. +collection.for(unrelatedDefinition) + +const standardResolutions = collection.for(standardDefinition) +expectTypeOf<(typeof standardResolutions)[number]['response']>().toEqualTypeOf< + { approved: boolean } | undefined +>() +const jsonResolutions = collection.for(jsonDefinition) +expectTypeOf<(typeof jsonResolutions)[number]['response']>().toEqualTypeOf< + { accepted: boolean } | undefined +>() +const allResolutions = collection.all() +expectTypeOf(allResolutions).toEqualTypeOf< + ReadonlyArray> +>() +expectTypeOf(collection.all(standardDefinition)).toEqualTypeOf< + ReadonlyArray> +>() +expectTypeOf(collection.all(standardDefinition, jsonDefinition)).toEqualTypeOf< + ReadonlyArray> +>() + +const builderStandard: DefinedChatMiddleware< + unknown, + readonly [], + readonly [], + typeof standardDefinition +> = { + onInterruptResolution() { + return { toolResume: 'continue' } + }, +} +const builderJson: DefinedChatMiddleware< + unknown, + readonly [], + readonly [], + typeof jsonDefinition +> = { + onInterruptResolution() { + return { toolResume: 'continue' } + }, +} +const interruptFreeBuilt = createChatMiddleware() + .use({ + name: 'logging-middleware', + onConfig(_ctx, config) { + return config + }, + }) + .use({ + name: 'second-logging-middleware', + onConfig(_ctx, config) { + return config + }, + }) + .build() + +const built = createChatMiddleware() + .use(builderStandard) + .use(builderJson) + .build() +expectTypeOf(built).toEqualTypeOf< + [typeof builderStandard, typeof builderJson] +>() +const unionAwareBuilder = createChatMiddleware() + .use(builderStandard) + .use(builderJson) + .use({ + onInterruptResolution( + _ctx, + resolutions: InterruptResolutionCollection, + ) { + expectTypeOf(resolutions.all()).toEqualTypeOf< + ReadonlyArray> + >() + return { toolResume: 'continue' } + }, + }) +void unionAwareBuilder + +const validPhase: InterruptBoundaryPhase = 'beforeModel' +expectTypeOf(validPhase).toEqualTypeOf<'beforeModel'>() +const validMiddleware: ChatMiddleware = { + onInterruptBoundary(ctx) { + const phase: InterruptBoundaryPhase = ctx.phase + void phase + return + }, + onInterruptResolution() { + return { toolResume: 'continue' } + }, +} +expectTypeOf(validMiddleware).toMatchTypeOf< + ChatMiddleware +>() + +// @ts-expect-error The boundary phase is restricted to the four engine phases. +const invalidPhase: InterruptBoundaryPhase = 'init' +void invalidPhase + +const invalidReturn: ChatMiddleware = { + // @ts-expect-error Boundary hooks may return only an interrupts collection. + onInterruptBoundary: () => ({ requests: [] }), +} +void invalidReturn + +// @ts-expect-error Empty resolution objects are not valid decisions. +const invalidResolution: InterruptResolutionResult = {} +void invalidResolution + +const invalidToolResume: ChatMiddleware = { + // @ts-expect-error Only continue, cancel, or stop are valid. + onInterruptResolution() { + return { toolResume: 'pause' } + }, +} +void invalidToolResume + +declare const adapter: AnyTextAdapter +chat({ + adapter, + middleware: interruptFreeBuilt, +}) +chat({ + adapter, + // @ts-expect-error Literal interrupt definition ids must be unique in one chat registry. + interrupts: [standardDefinition, duplicateIdDefinition] as const, +}) +// Runtime IDs remain compatible. The runtime validation rejects duplicates +// only when the values are known at execution time. +chat({ + adapter, + interrupts: [runtimeIdDefinition, runtimeIdDefinition] as const, +}) +chat({ + adapter, + interrupts: [standardDefinition, jsonDefinition] as const, + middleware: [ + { + onInterruptBoundary: () => ({ + interrupts: [standardRequest, jsonRequest], + }), + onInterruptResolution(_ctx, resolutions) { + resolutions.for(standardDefinition) + resolutions.for(jsonDefinition) + return { toolResume: 'continue' } + }, + }, + ], +}) + +// Public calls must infer the merged tool and user context without a +// predeclared ChatMiddleware annotation. +chat({ + adapter, + interrupts: [standardDefinition, jsonDefinition] as const, + tools: [contextTool] as const, + context: { tenantId: 'tenant', toolFlag: true }, + middleware: [ + { + setup(ctx) { + expectTypeOf(ctx.context.tenantId).toEqualTypeOf() + expectTypeOf(ctx.context.toolFlag).toEqualTypeOf() + // @ts-expect-error Unknown context fields are rejected. + ctx.context.unknownField + }, + onInterruptBoundary(ctx) { + expectTypeOf(ctx.context.tenantId).toEqualTypeOf() + expectTypeOf(ctx.context.toolFlag).toEqualTypeOf() + // @ts-expect-error Unknown context fields are rejected. + ctx.context.unknownField + return { interrupts: [standardRequest] } + }, + onInterruptResolution(ctx, resolutions) { + expectTypeOf(ctx.context.tenantId).toEqualTypeOf() + expectTypeOf(ctx.context.toolFlag).toEqualTypeOf() + resolutions.for(standardDefinition) + resolutions.for(jsonDefinition) + return { toolResume: 'continue' } + }, + }, + ], +}) + +const publicOptions = createChatOptions({ + adapter, + interrupts: [standardDefinition, jsonDefinition] as const, + tools: [contextTool] as const, + context: { tenantId: 'tenant', toolFlag: true }, + middleware: [ + { + setup(ctx) { + expectTypeOf(ctx.context.tenantId).toEqualTypeOf() + expectTypeOf(ctx.context.toolFlag).toEqualTypeOf() + // @ts-expect-error Unknown context fields are rejected. + ctx.context.unknownField + }, + onInterruptBoundary(ctx) { + expectTypeOf(ctx.context.tenantId).toEqualTypeOf() + expectTypeOf(ctx.context.toolFlag).toEqualTypeOf() + return { interrupts: [jsonRequest] } + }, + onInterruptResolution(_ctx, resolutions) { + resolutions.for(standardDefinition) + resolutions.for(jsonDefinition) + return { toolResume: 'continue' } + }, + }, + ], +}) +chat({ + ...publicOptions, + context: { tenantId: 'tenant', toolFlag: true }, +}) +publicOptions.middleware?.push({ + // @ts-expect-error Foreign definitions are rejected by options.middleware.push. + onInterruptBoundary: () => ({ + interrupts: [unregisteredRequest], + }), +}) + +chat({ + adapter, + interrupts: [standardDefinition] as const, + middleware: [ + { + // @ts-expect-error Inline middleware may request only registered definitions. + onInterruptBoundary: () => ({ interrupts: [unregisteredRequest] }), + onInterruptResolution: (_ctx, resolutions) => { + // @ts-expect-error Inline middleware resolutions may inspect only registered definitions. + resolutions.for(unrelatedDefinition) + return { toolResume: 'continue' } + }, + }, + ], +}) + +const typedMiddleware: ChatMiddleware<{ tenantId: string }> = { + setup(ctx) { + expectTypeOf(ctx.context).toEqualTypeOf<{ tenantId: string }>() + }, +} +chat({ + adapter, + interrupts: [] as const, + tools: [] as const, + context: { tenantId: 'tenant' }, + middleware: [typedMiddleware], +}) +chat({ + adapter, + tools: [] as const, + context: { tenantId: 'tenant' }, + middleware: [typedMiddleware], +}) + +chat({ + adapter, + // @ts-expect-error Middleware cannot emit a first-party interrupt without a registry. + middleware: [ + { + onInterruptBoundary: () => ({ interrupts: [standardRequest] }), + }, + ], +}) + +const noRegistryMiddleware: ChatMiddleware = { + // @ts-expect-error A never registry cannot emit any first-party request. + onInterruptBoundary: () => ({ interrupts: [standardRequest] }), +} +void noRegistryMiddleware +// @ts-expect-error A never definition cannot produce a request. +const noRegistryRequest: GenericInterruptRequest = standardRequest +void noRegistryRequest + +const standaloneRegisteredEmitter = defineChatMiddleware({ + onInterruptBoundary: () => ({ interrupts: [standardRequest] }), +}) +chat({ + adapter, + interrupts: [standardDefinition] as const, + middleware: [standaloneRegisteredEmitter], +}) +chat({ + adapter, + // @ts-expect-error A reusable emitter also requires its definition in the chat registry. + middleware: [standaloneRegisteredEmitter], +}) + +chat({ + adapter, + interrupts: [] as const, + // @ts-expect-error Middleware cannot emit a first-party interrupt with an empty registry. + middleware: [ + { + onInterruptBoundary: () => ({ interrupts: [standardRequest] }), + }, + ], +}) diff --git a/packages/ai/tests/middleware-interrupt.test.ts b/packages/ai/tests/middleware-interrupt.test.ts new file mode 100644 index 0000000000..acd7946e2a --- /dev/null +++ b/packages/ai/tests/middleware-interrupt.test.ts @@ -0,0 +1,199 @@ +import { describe, expect, it } from 'vitest' +import { z } from 'zod' +import { aiEventClient } from '@tanstack/ai-event-client' +import { defineInterrupt } from '../src/interrupt-definition' +import { InternalLogger } from '../src/logger/internal-logger' +import type { ResolvedCategories } from '../src/logger/internal-logger' +import { MiddlewareRunner } from '../src/activities/chat/middleware/compose' +import { CapabilityRegistry } from '../src/activities/chat/middleware/capabilities' +import type { + ChatMiddleware, + ChatMiddlewareContext, + InterruptBoundaryPhase, +} from '../src/activities/chat/middleware/types' + +const categories: ResolvedCategories = { + request: false, + provider: false, + output: false, + middleware: false, + tools: false, + agentLoop: false, + config: false, + errors: false, + sandbox: false, +} + +const logger = new InternalLogger( + { debug() {}, info() {}, warn() {}, error() {} }, + categories, +) + +const context: ChatMiddlewareContext = { + requestId: 'request', + streamId: 'stream', + runId: 'run', + threadId: 'thread', + phase: 'beforeModel', + iteration: 0, + chunkIndex: 0, + abort() {}, + context: undefined, + defer() {}, + activity: 'chat', + provider: 'provider', + model: 'model', + source: 'server', + streaming: true, + systemPrompts: [], + messageCount: 0, + hasTools: false, + currentMessageId: null, + accumulatedContent: '', + messages: [], + createId: (prefix) => prefix, + capabilities: new CapabilityRegistry(), + get: () => { + throw new Error('unused in middleware interrupt tests') + }, + getOptional: () => undefined, + provide() {}, +} + +const approval = defineInterrupt({ + id: 'middleware-approval', + responseSchema: z.object({ approved: z.boolean() }), +}) + +describe('generic interrupt middleware composition', () => { + it('runs every boundary hook in order and aggregates requests', async () => { + const phases: InterruptBoundaryPhase[] = [ + 'beforeModel', + 'afterModel', + 'beforeTools', + 'afterTools', + ] + const calls: string[] = [] + const firstRequest = approval.interrupt({ + key: 'first', + reason: 'test', + message: 'First', + }) + const secondRequest = approval.interrupt({ + key: 'second', + reason: 'test', + message: 'Second', + }) + const first: ChatMiddleware = { + onInterruptBoundary(ctx) { + calls.push(`first:${ctx.phase}`) + return { interrupts: [firstRequest] } + }, + } + const second: ChatMiddleware = { + onInterruptBoundary(ctx) { + calls.push(`second:${ctx.phase}`) + return { interrupts: [secondRequest] } + }, + } + const runner = new MiddlewareRunner( + [first, second], + logger, + ) + + const requestsByPhase = [] + for (const phase of phases) { + requestsByPhase.push( + await runner.runOnInterruptBoundary({ ...context, phase }), + ) + } + + expect(requestsByPhase).toEqual( + phases.map(() => [firstRequest, secondRequest]), + ) + expect(calls).toEqual( + phases.flatMap((phase) => [`first:${phase}`, `second:${phase}`]), + ) + expect(phases).toHaveLength(4) + }) + + it('handles void boundary results and composes the last resolution decision', async () => { + const first: ChatMiddleware = { + onInterruptBoundary() {}, + onInterruptResolution() { + return { toolResume: 'continue' } + }, + } + const second: ChatMiddleware = { + onInterruptBoundary() {}, + onInterruptResolution() { + return { toolResume: 'stop' } + }, + } + const runner = new MiddlewareRunner( + [first, second], + logger, + ) + + await expect( + runner.runOnInterruptBoundary({ ...context, phase: 'afterTools' }), + ).resolves.toEqual([]) + await expect( + runner.runOnInterruptResolution(context, { + for: () => [], + all: () => [], + }), + ).resolves.toEqual({ toolResume: 'stop' }) + }) + + it('instruments lifecycle hooks and skips internal middleware', async () => { + const events: Array<{ + middlewareName: string + hookName: string + duration: number + }> = [] + const unsubscribe = aiEventClient.on( + 'middleware:hook:executed', + (event) => { + events.push(event.payload) + }, + { withEventTarget: true }, + ) + const visible: ChatMiddleware = { + name: 'visible', + onInterruptBoundary() {}, + onInterruptResolution() {}, + } + const internal: ChatMiddleware = { + name: 'devtools', + onInterruptBoundary() {}, + onInterruptResolution() {}, + } + const runner = new MiddlewareRunner( + [visible, internal], + logger, + ) + + await runner.runOnInterruptBoundary({ ...context, phase: 'beforeModel' }) + await runner.runOnInterruptResolution(context, { + for: () => [], + all: () => [], + }) + unsubscribe() + + expect(events).toHaveLength(2) + expect(events).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + middlewareName: 'visible', + hookName: 'onInterruptBoundary', + }), + expect.objectContaining({ + middlewareName: 'visible', + hookName: 'onInterruptResolution', + }), + ]), + ) + expect(events.every((event) => event.duration >= 0)).toBe(true) + }) +}) diff --git a/packages/ai/tests/stream-processor.test.ts b/packages/ai/tests/stream-processor.test.ts index 72b825771a..ea2dbb972e 100644 --- a/packages/ai/tests/stream-processor.test.ts +++ b/packages/ai/tests/stream-processor.test.ts @@ -1271,6 +1271,51 @@ describe('StreamProcessor', () => { }) }) + it('does not auto-run client tools when a generic interrupt shares the batch', () => { + const events = spyEvents() + const processor = new StreamProcessor({ events }) + processor.prepareAssistantMessage() + + processor.processChunk(ev.runStarted()) + processor.processChunk(ev.toolStart('tc-1', 'clientSearch')) + processor.processChunk(ev.toolArgs('tc-1', '{"query":"test"}')) + processor.processChunk({ + ...ev.runFinished('tool_calls'), + outcome: { + type: 'interrupt', + interrupts: [ + { + id: 'generic-1', + reason: 'review_required', + message: 'Review the plan', + metadata: { + 'tanstack:interruptBinding': { + v: 1, + kind: 'generic', + interruptId: 'generic-1', + definitionId: 'review-plan', + key: 'one', + batchIndex: 0, + }, + }, + }, + { + id: 'client_tool_tc-1', + reason: 'tanstack:client_tool_execution', + toolCallId: 'tc-1', + metadata: { + kind: 'client_tool', + toolName: 'clientSearch', + input: { query: 'test' }, + }, + }, + ], + }, + }) + + expect(events.onToolCall).not.toHaveBeenCalled() + }) + it('should deliver client tool interrupt before stream end for mixed tool streams', async () => { const order: Array = [] const events = spyEvents() diff --git a/testing/e2e/fixtures/middleware-test/generic-after-model.json b/testing/e2e/fixtures/middleware-test/generic-after-model.json new file mode 100644 index 0000000000..649b21315e --- /dev/null +++ b/testing/e2e/fixtures/middleware-test/generic-after-model.json @@ -0,0 +1,11 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[generic-after-model] run test", + "sequenceIndex": 0 + }, + "response": { "content": "AFTER_MODEL_CONTENT" } + } + ] +} diff --git a/testing/e2e/fixtures/middleware-test/generic-after-tools.json b/testing/e2e/fixtures/middleware-test/generic-after-tools.json new file mode 100644 index 0000000000..57226df8fb --- /dev/null +++ b/testing/e2e/fixtures/middleware-test/generic-after-tools.json @@ -0,0 +1,25 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[generic-after-tools] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "inspect_review", + "arguments": "{\"reviewId\":\"review-4\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[generic-after-tools] run test", + "sequenceIndex": 1 + }, + "response": { "content": "AFTER_TOOLS_RESOLVED" } + } + ] +} diff --git a/testing/e2e/fixtures/middleware-test/generic-before-model.json b/testing/e2e/fixtures/middleware-test/generic-before-model.json new file mode 100644 index 0000000000..004635a2c1 --- /dev/null +++ b/testing/e2e/fixtures/middleware-test/generic-before-model.json @@ -0,0 +1,11 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[generic-before-model] run test", + "sequenceIndex": 0 + }, + "response": { "content": "BEFORE_MODEL_RESOLVED" } + } + ] +} diff --git a/testing/e2e/fixtures/middleware-test/generic-before-tools-cancel.json b/testing/e2e/fixtures/middleware-test/generic-before-tools-cancel.json new file mode 100644 index 0000000000..17c5dce062 --- /dev/null +++ b/testing/e2e/fixtures/middleware-test/generic-before-tools-cancel.json @@ -0,0 +1,29 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[generic-before-tools-cancel] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "delete_review", + "arguments": "{\"reviewId\":\"review-2\"}" + }, + { + "name": "render_review", + "arguments": "{\"reviewId\":\"review-2\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[generic-before-tools-cancel] run test", + "sequenceIndex": 1 + }, + "response": { "content": "TOOLS_CANCELLED" } + } + ] +} diff --git a/testing/e2e/fixtures/middleware-test/generic-before-tools-continue.json b/testing/e2e/fixtures/middleware-test/generic-before-tools-continue.json new file mode 100644 index 0000000000..e1ee4aa6e7 --- /dev/null +++ b/testing/e2e/fixtures/middleware-test/generic-before-tools-continue.json @@ -0,0 +1,29 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[generic-before-tools-continue] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "delete_review", + "arguments": "{\"reviewId\":\"review-1\"}" + }, + { + "name": "render_review", + "arguments": "{\"reviewId\":\"review-1\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[generic-before-tools-continue] run test", + "sequenceIndex": 1 + }, + "response": { "content": "TOOLS_CONTINUED" } + } + ] +} diff --git a/testing/e2e/fixtures/middleware-test/generic-before-tools-stop.json b/testing/e2e/fixtures/middleware-test/generic-before-tools-stop.json new file mode 100644 index 0000000000..32ac87da64 --- /dev/null +++ b/testing/e2e/fixtures/middleware-test/generic-before-tools-stop.json @@ -0,0 +1,22 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[generic-before-tools-stop] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "delete_review", + "arguments": "{\"reviewId\":\"review-3\"}" + }, + { + "name": "render_review", + "arguments": "{\"reviewId\":\"review-3\"}" + } + ] + } + } + ] +} diff --git a/testing/e2e/src/lib/devtools-test.ts b/testing/e2e/src/lib/devtools-test.ts index ccda6f2c13..3f3ff40910 100644 --- a/testing/e2e/src/lib/devtools-test.ts +++ b/testing/e2e/src/lib/devtools-test.ts @@ -3,19 +3,24 @@ export interface DevtoolsRouteSearch { aimockPort?: number } +export function parseAimockPort(value: unknown): number | undefined { + const port = + typeof value === 'number' + ? value + : typeof value === 'string' + ? Number.parseInt(value, 10) + : undefined + return port != null && !Number.isNaN(port) ? port : undefined +} + export function parseDevtoolsRouteSearch( search: Record, ): DevtoolsRouteSearch { - const aimockPort = - typeof search.aimockPort === 'string' - ? Number.parseInt(search.aimockPort, 10) - : undefined + const aimockPort = parseAimockPort(search.aimockPort) return { ...(typeof search.testId === 'string' ? { testId: search.testId } : {}), - ...(aimockPort !== undefined && !Number.isNaN(aimockPort) - ? { aimockPort } - : {}), + ...(aimockPort !== undefined ? { aimockPort } : {}), } } diff --git a/testing/e2e/src/lib/generic-middleware-interrupts.ts b/testing/e2e/src/lib/generic-middleware-interrupts.ts new file mode 100644 index 0000000000..407381b221 --- /dev/null +++ b/testing/e2e/src/lib/generic-middleware-interrupts.ts @@ -0,0 +1,69 @@ +import { defineInterrupt, toolDefinition } from '@tanstack/ai' +import { z } from 'zod' + +export const reviewPlan = defineInterrupt({ + id: 'review-plan', + payloadSchema: z.object({ + title: z.string(), + boundary: z.enum([ + 'beforeModel', + 'afterModel', + 'beforeTools', + 'afterTools', + ]), + }), + responseSchema: z.object({ + approved: z.boolean(), + note: z.string(), + }), +}) + +export const deleteReviewTool = toolDefinition({ + name: 'delete_review', + description: 'Delete a review after approval', + inputSchema: z.object({ reviewId: z.string() }), + outputSchema: z.object({ deleted: z.boolean(), reviewId: z.string() }), + needsApproval: true, +}) + +export const renderReviewTool = toolDefinition({ + name: 'render_review', + description: 'Render a review in the browser', + inputSchema: z.object({ reviewId: z.string() }), + outputSchema: z.object({ rendered: z.boolean(), reviewId: z.string() }), +}) + +export const inspectReviewTool = toolDefinition({ + name: 'inspect_review', + description: 'Inspect a review on the server', + inputSchema: z.object({ reviewId: z.string() }), + outputSchema: z.object({ inspected: z.boolean(), reviewId: z.string() }), +}) + +export const genericScenarios = [ + 'generic-before-model', + 'generic-after-model', + 'generic-before-tools-continue', + 'generic-before-tools-cancel', + 'generic-before-tools-stop', + 'generic-after-tools', +] as const + +export type GenericScenario = (typeof genericScenarios)[number] + +export function isGenericScenario(value: string): value is GenericScenario { + return genericScenarios.some((scenario) => scenario === value) +} + +export function boundaryForScenario(scenario: GenericScenario) { + if (scenario === 'generic-before-model') return 'beforeModel' as const + if (scenario === 'generic-after-model') return 'afterModel' as const + if (scenario === 'generic-after-tools') return 'afterTools' as const + return 'beforeTools' as const +} + +export function toolResumeForScenario(scenario: GenericScenario) { + if (scenario === 'generic-before-tools-cancel') return 'cancel' as const + if (scenario === 'generic-before-tools-stop') return 'stop' as const + return 'continue' as const +} diff --git a/testing/e2e/src/lib/phase-capture.ts b/testing/e2e/src/lib/phase-capture.ts index e9ece87ff0..9fc78d432c 100644 --- a/testing/e2e/src/lib/phase-capture.ts +++ b/testing/e2e/src/lib/phase-capture.ts @@ -10,6 +10,26 @@ export interface YieldedChunkSummary { /** The chunk's discriminant (e.g. RUN_STARTED, TEXT_MESSAGE_CONTENT). */ type: string + /** The run ID emitted on this specific stream chunk, when the event has one. */ + runId?: string + outcomeType?: string + interruptCount?: number +} + +export interface GenericBoundaryCapture { + phase: string + runId: string +} + +export interface GenericResolutionCapture { + definitionId: string + status: 'resolved' | 'cancelled' + response?: unknown +} + +export interface GenericToolExecutionCapture { + name: string + side: 'server' } export interface PhaseCapture { @@ -27,6 +47,10 @@ export interface PhaseCapture { * chain has applied its transformations. */ yieldedChunks: Array + boundaries: Array + resolutions: Array + policies: Array<'continue' | 'cancel' | 'stop'> + toolExecutions: Array } const captures: Map = new Map() @@ -34,7 +58,15 @@ const captures: Map = new Map() function bucketFor(captureId: string): PhaseCapture { let bucket = captures.get(captureId) if (!bucket) { - bucket = { phases: [], onFinishCount: 0, yieldedChunks: [] } + bucket = { + phases: [], + onFinishCount: 0, + yieldedChunks: [], + boundaries: [], + resolutions: [], + policies: [], + toolExecutions: [], + } captures.set(captureId, bucket) } return bucket @@ -45,6 +77,10 @@ export function resetPhaseCapture(captureId: string): void { phases: [], onFinishCount: 0, yieldedChunks: [], + boundaries: [], + resolutions: [], + policies: [], + toolExecutions: [], }) } @@ -66,3 +102,31 @@ export function recordYieldedChunk( ): void { bucketFor(captureId).yieldedChunks.push(chunk) } + +export function recordGenericBoundary( + captureId: string, + boundary: GenericBoundaryCapture, +): void { + bucketFor(captureId).boundaries.push(boundary) +} + +export function recordGenericResolution( + captureId: string, + resolution: GenericResolutionCapture, +): void { + bucketFor(captureId).resolutions.push(resolution) +} + +export function recordGenericPolicy( + captureId: string, + policy: 'continue' | 'cancel' | 'stop', +): void { + bucketFor(captureId).policies.push(policy) +} + +export function recordGenericToolExecution( + captureId: string, + execution: GenericToolExecutionCapture, +): void { + bucketFor(captureId).toolExecutions.push(execution) +} diff --git a/testing/e2e/src/routes/$provider/$feature.tsx b/testing/e2e/src/routes/$provider/$feature.tsx index 5cf36c3a26..bd91cfc470 100644 --- a/testing/e2e/src/routes/$provider/$feature.tsx +++ b/testing/e2e/src/routes/$provider/$feature.tsx @@ -10,6 +10,7 @@ import type { } from '@tanstack/ai-client' import type { GeminiInteractionsCustomEventValue } from '@tanstack/ai-gemini/experimental' import type { Feature, Mode, Provider } from '@/lib/types' +import { parseAimockPort } from '@/lib/devtools-test' import { ALL_FEATURES, ALL_PROVIDERS } from '@/lib/types' import { isSupported } from '@/lib/feature-support' import { addToCartToolDef } from '@/lib/tools' @@ -27,16 +28,10 @@ const VALID_MODES = new Set(['sse', 'http-stream', 'fetcher']) export const Route = createFileRoute('/$provider/$feature')({ component: FeaturePage, validateSearch: (search: Record) => { - const port = - typeof search.aimockPort === 'number' - ? search.aimockPort - : typeof search.aimockPort === 'string' - ? parseInt(search.aimockPort, 10) - : undefined const rawMode = typeof search.mode === 'string' ? search.mode : undefined return { testId: typeof search.testId === 'string' ? search.testId : undefined, - aimockPort: port != null && !isNaN(port) ? port : undefined, + aimockPort: parseAimockPort(search.aimockPort), mode: rawMode && VALID_MODES.has(rawMode as Mode) ? (rawMode as Mode) diff --git a/testing/e2e/src/routes/api.foreign-interrupt.ts b/testing/e2e/src/routes/api.foreign-interrupt.ts index a4b7c127d8..a30670047d 100644 --- a/testing/e2e/src/routes/api.foreign-interrupt.ts +++ b/testing/e2e/src/routes/api.foreign-interrupt.ts @@ -148,15 +148,16 @@ export const Route = createFileRoute('/api/foreign-interrupt')({ POST: async ({ request }) => { const body: unknown = await request.json() const threadId = stringField(body, 'threadId') ?? 'thread-1' + // Bindings correlate on the request runId. A fresh server id would + // make `ours` generic but not resolvable. + const runId = stringField(body, 'runId') ?? `run-${threadId}` const resumed = resumedInterruptIds(body) if (resumed.length > 0) { return toServerSentEventsResponse( - resumedRun(threadId, `run-${threadId}-continuation`, resumed), + resumedRun(threadId, `${runId}-continuation`, resumed), ) } - return toServerSentEventsResponse( - foreignRun(threadId, `run-${threadId}`), - ) + return toServerSentEventsResponse(foreignRun(threadId, runId)) }, }, }, diff --git a/testing/e2e/src/routes/api.middleware-test.ts b/testing/e2e/src/routes/api.middleware-test.ts index 33d588783a..031f38946d 100644 --- a/testing/e2e/src/routes/api.middleware-test.ts +++ b/testing/e2e/src/routes/api.middleware-test.ts @@ -34,11 +34,25 @@ import type { import { guitarRecommendationSchema } from '@/lib/schemas' import { getPhaseCapture, + recordGenericBoundary, + recordGenericPolicy, + recordGenericResolution, + recordGenericToolExecution, recordOnFinish, recordPhase, recordYieldedChunk, resetPhaseCapture, } from '@/lib/phase-capture' +import { + boundaryForScenario, + deleteReviewTool, + inspectReviewTool, + isGenericScenario, + renderReviewTool, + reviewPlan, + toolResumeForScenario, +} from '@/lib/generic-middleware-interrupts' +import type { GenericScenario } from '@/lib/generic-middleware-interrupts' import { createTextAdapter } from '@/lib/providers' import { getOtelCapture, @@ -167,12 +181,99 @@ async function* teeForPhaseCapture( source: AsyncIterable, captureId: string, ): AsyncIterable { + let currentRunId: string | undefined for await (const chunk of source) { - recordYieldedChunk(captureId, { type: chunk.type }) + if (chunk.type === 'RUN_STARTED' && typeof chunk.runId === 'string') { + currentRunId = chunk.runId + } + const runId = + 'runId' in chunk && typeof chunk.runId === 'string' + ? chunk.runId + : currentRunId + recordYieldedChunk(captureId, { + type: chunk.type, + ...(runId !== undefined ? { runId } : {}), + ...(chunk.type === 'RUN_FINISHED' && chunk.outcome + ? { outcomeType: chunk.outcome.type } + : {}), + ...(chunk.type === 'RUN_FINISHED' && chunk.outcome?.type === 'interrupt' + ? { interruptCount: chunk.outcome.interrupts.length } + : {}), + }) yield chunk } } +function lifecycleMiddlewareStack( + existing: Array, + extra: ChatMiddleware, +): Array> { + return [...existing, extra] +} + +function createGenericLifecycleMiddleware( + captureId: string, + scenario: GenericScenario, +): ChatMiddleware { + const boundary = boundaryForScenario(scenario) + return { + name: 'generic-lifecycle', + onInterruptBoundary(ctx) { + if (ctx.phase !== boundary || ctx.parentRunId) return + recordGenericBoundary(captureId, { phase: ctx.phase, runId: ctx.runId }) + return { + interrupts: [ + reviewPlan.interrupt({ + key: `${scenario}-review`, + reason: 'review_required', + message: `Review the plan at ${ctx.phase}`, + payload: { title: 'Middleware review plan', boundary: ctx.phase }, + }), + ], + } + }, + onInterruptResolution(_ctx, resolutions) { + for (const resolution of resolutions.for(reviewPlan)) { + recordGenericResolution(captureId, { + definitionId: resolution.request.definition.id, + status: resolution.status, + ...(resolution.status === 'resolved' + ? { response: resolution.response } + : {}), + }) + } + const policy = toolResumeForScenario(scenario) + recordGenericPolicy(captureId, policy) + return { toolResume: policy } + }, + } +} + +function genericTools(captureId: string, scenario: GenericScenario) { + if (scenario === 'generic-after-tools') { + return [ + inspectReviewTool.server(async ({ reviewId }) => { + recordGenericToolExecution(captureId, { + name: 'inspect_review', + side: 'server', + }) + return { inspected: true, reviewId } + }), + ] + } + if (boundaryForScenario(scenario) !== 'beforeTools') return [] + return [ + deleteReviewTool.server(async ({ reviewId }) => { + recordGenericToolExecution(captureId, { + name: 'delete_review', + side: 'server', + }) + return { deleted: true, reviewId } + }), + renderReviewTool.client(), + ] +} + /** * Fake memory adapter for `memory` mode. `recall` unconditionally returns a * known system-prompt block plus a memory tool (so the spec can assert both the @@ -385,6 +486,12 @@ export const Route = createFileRoute('/api/middleware-test')({ ) const middleware: Array = [] + let genericLifecycleMiddleware: + | ChatMiddleware + | undefined + const genericScenario = isGenericScenario(scenario) + ? scenario + : undefined if (middlewareMode === 'chunk-transform') middleware.push(chunkTransformMiddleware) @@ -410,6 +517,25 @@ export const Route = createFileRoute('/api/middleware-test')({ resetPhaseCapture(testId) middleware.push(createPhaseRecorderMiddleware(testId)) } + if (middlewareMode === 'generic-lifecycle') { + if (!testId || !genericScenario) { + return new Response( + JSON.stringify({ + error: + 'generic-lifecycle mode requires testId and a generic scenario', + }), + { + status: 400, + headers: { 'Content-Type': 'application/json' }, + }, + ) + } + if (!params.parentRunId) resetPhaseCapture(testId) + genericLifecycleMiddleware = createGenericLifecycleMiddleware( + testId, + genericScenario, + ) + } if (middlewareMode === 'memory') { if (!testId) { return new Response( @@ -452,7 +578,12 @@ export const Route = createFileRoute('/api/middleware-test')({ ) } - const tools = scenario === 'with-tool' ? [weatherTool] : [] + const tools = + genericScenario && testId + ? genericTools(testId, genericScenario) + : scenario === 'with-tool' + ? [weatherTool] + : [] // The two `structured-output*` scenarios both bind the same // guitar schema; they differ only in what the spec asserts (phases @@ -473,22 +604,44 @@ export const Route = createFileRoute('/api/middleware-test')({ stream: true, abortController, }) - : chat({ - ...adapterOptions, - messages: params.messages, - tools, - middleware, - threadId: params.threadId, - runId: params.runId, - agentLoopStrategy: maxIterations(10), - abortController, - }) + : genericLifecycleMiddleware + ? chat({ + ...adapterOptions, + messages: params.messages, + tools, + middleware: lifecycleMiddlewareStack( + middleware, + genericLifecycleMiddleware, + ), + threadId: params.threadId, + runId: params.runId, + parentRunId: params.parentRunId, + resume: params.resume, + interrupts: [reviewPlan] as const, + agentLoopStrategy: maxIterations(10), + abortController, + }) + : chat({ + ...adapterOptions, + messages: params.messages, + tools, + middleware, + threadId: params.threadId, + runId: params.runId, + parentRunId: params.parentRunId, + resume: params.resume, + interrupts: genericScenario ? [reviewPlan] : undefined, + agentLoopStrategy: maxIterations(10), + abortController, + }) // Tee the post-middleware stream when `phase-recorder` is active // so the spec can assert on what the consumer ultimately sees // (e.g. exactly one RUN_STARTED/RUN_FINISHED pair). const stream = - middlewareMode === 'phase-recorder' && testId + (middlewareMode === 'phase-recorder' || + middlewareMode === 'generic-lifecycle') && + testId ? teeForPhaseCapture(rawStream, testId) : rawStream diff --git a/testing/e2e/src/routes/interrupts-test.tsx b/testing/e2e/src/routes/interrupts-test.tsx index 183435cfaf..2a594d5645 100644 --- a/testing/e2e/src/routes/interrupts-test.tsx +++ b/testing/e2e/src/routes/interrupts-test.tsx @@ -15,6 +15,7 @@ import { shareAdoptionStory, } from '@/lib/interrupt-scenario-tools' import type { ResolutionConfig } from '@/lib/interrupt-scenario-tools' +import { parseAimockPort } from '@/lib/devtools-test' /** * Interrupt playground — deterministic e2e host for every wildlife interrupt @@ -533,13 +534,9 @@ function sectionStyle(accent: string) { export const Route = createFileRoute('/interrupts-test')({ component: InterruptsTestPage, validateSearch: (search: Record) => { - const port = - typeof search.aimockPort === 'string' - ? parseInt(search.aimockPort, 10) - : undefined return { testId: typeof search.testId === 'string' ? search.testId : undefined, - aimockPort: port != null && !isNaN(port) ? port : undefined, + aimockPort: parseAimockPort(search.aimockPort), scenario: typeof search.scenario === 'string' ? search.scenario : undefined, } diff --git a/testing/e2e/src/routes/middleware-test.tsx b/testing/e2e/src/routes/middleware-test.tsx index 71ddf027bd..f9ae3899fe 100644 --- a/testing/e2e/src/routes/middleware-test.tsx +++ b/testing/e2e/src/routes/middleware-test.tsx @@ -1,6 +1,17 @@ -import { useState } from 'react' +import { useMemo, useState } from 'react' import { createFileRoute } from '@tanstack/react-router' -import { useChat, fetchServerSentEvents } from '@tanstack/ai-react' +import { + useChat, + fetchServerSentEvents, + localStoragePersistence, +} from '@tanstack/ai-react' +import { clientTools } from '@tanstack/ai-client' +import { parseAimockPort } from '@/lib/devtools-test' +import { + deleteReviewTool, + renderReviewTool, + reviewPlan, +} from '@/lib/generic-middleware-interrupts' const MIDDLEWARE_MODES = [ { id: 'none', label: 'No Middleware' }, @@ -10,8 +21,11 @@ const MIDDLEWARE_MODES = [ { id: 'phase-recorder', label: 'Phase Recorder (capture phase + chunks)' }, { id: 'otel', label: 'OpenTelemetry (capture spans/metrics)' }, { id: 'memory', label: 'Memory (recall/save)' }, + { id: 'generic-lifecycle', label: 'Generic Interrupt Lifecycle' }, ] as const +const genericPersistence = localStoragePersistence() + interface PhaseCaptureSnapshot { phases: Array onFinishCount: number @@ -60,13 +74,9 @@ function toPhaseCapture(raw: unknown): PhaseCaptureSnapshot { export const Route = createFileRoute('/middleware-test')({ component: MiddlewareTestPage, validateSearch: (search: Record) => { - const port = - typeof search.aimockPort === 'string' - ? parseInt(search.aimockPort, 10) - : undefined return { testId: typeof search.testId === 'string' ? search.testId : undefined, - aimockPort: port != null && !isNaN(port) ? port : undefined, + aimockPort: parseAimockPort(search.aimockPort), // `provider` / `model` are forwarded to the server route so the // structured-output × middleware spec can exercise both the // native-combined-mode path (modern openai / claude 4.5+) and the @@ -74,14 +84,25 @@ export const Route = createFileRoute('/middleware-test')({ provider: typeof search.provider === 'string' ? search.provider : undefined, model: typeof search.model === 'string' ? search.model : undefined, + scenario: + typeof search.scenario === 'string' ? search.scenario : undefined, + middlewareMode: + typeof search.middlewareMode === 'string' + ? search.middlewareMode + : undefined, } }, }) function MiddlewareTestPage() { - const { testId, aimockPort, provider, model } = Route.useSearch() - const [scenario, setScenario] = useState('basic-text') - const [middlewareMode, setMiddlewareMode] = useState('none') + const { testId, aimockPort, provider, model, ...searchSelection } = + Route.useSearch() + const [scenario, setScenario] = useState( + searchSelection.scenario ?? 'basic-text', + ) + const [middlewareMode, setMiddlewareMode] = useState( + searchSelection.middlewareMode ?? 'none', + ) const [testComplete, setTestComplete] = useState(false) const [phaseCapture, setPhaseCapture] = useState(EMPTY_PHASE_CAPTURE) @@ -89,11 +110,33 @@ function MiddlewareTestPage() { configs: Array<{ systemPrompts: Array; toolNames: Array }> saveCount: number }>({ configs: [], saveCount: 0 }) + const [clientToolExecutions, setClientToolExecutions] = useState(0) - const { messages, sendMessage, isLoading } = useChat({ - threadId: `mw-test-${scenario}-${middlewareMode}-${provider ?? 'openai'}-${model ?? 'default'}`, + const clientToolList = useMemo( + () => + clientTools( + deleteReviewTool.client(), + renderReviewTool.client(async ({ reviewId }) => { + setClientToolExecutions((count) => count + 1) + return { rendered: true, reviewId } + }), + ), + [], + ) + + const { messages, sendMessage, isLoading, interrupts } = useChat< + typeof clientToolList, + undefined, + unknown, + readonly [typeof reviewPlan] + >({ + threadId: `mw-test-${testId ?? 'manual'}-${scenario}-${middlewareMode}-${provider ?? 'openai'}-${model ?? 'default'}`, connection: fetchServerSentEvents('/api/middleware-test'), body: { scenario, middlewareMode, testId, aimockPort, provider, model }, + tools: clientToolList, + interrupts: [reviewPlan], + persistence: + middlewareMode === 'generic-lifecycle' ? genericPersistence : undefined, onFinish: () => { // For phase-recorder mode the spec reads `#mw-phases-json` / // `#mw-onfinish-count` / `#mw-yielded-chunks-json` AFTER @@ -137,9 +180,26 @@ function MiddlewareTestPage() { const handleRun = () => { setTestComplete(false) setPhaseCapture(EMPTY_PHASE_CAPTURE) + setClientToolExecutions(0) sendMessage(`[${scenario}] run test`) } + type ActiveInterrupt = (typeof interrupts)[number] + const reviewInterrupts = interrupts.filter( + ( + interrupt, + ): interrupt is Extract => + interrupt.kind === 'generic' && + 'definitionId' in interrupt && + interrupt.definitionId === reviewPlan.id, + ) + const approvalInterrupts = interrupts.filter( + ( + interrupt, + ): interrupt is Extract => + interrupt.kind === 'tool-approval', + ) + return (
Structured Output (Stream) + + + + + +
@@ -214,6 +286,51 @@ function MiddlewareTestPage() { Run Test + {reviewInterrupts.map((interrupt) => ( +
+ {interrupt.message} + + +
+ ))} + + {approvalInterrupts.map((interrupt) => ( +
+ + +
+ ))} +
     
   )
diff --git a/testing/e2e/src/routes/tools-test.tsx b/testing/e2e/src/routes/tools-test.tsx
index 6fc18ffc83..99c6a4970b 100644
--- a/testing/e2e/src/routes/tools-test.tsx
+++ b/testing/e2e/src/routes/tools-test.tsx
@@ -5,6 +5,7 @@ import { modelMessagesToUIMessages, toolDefinition } from '@tanstack/ai'
 import { z } from 'zod'
 import type { ModelMessage, ToolCallPart } from '@tanstack/ai'
 import type { UIMessage } from '@tanstack/ai-react'
+import { parseAimockPort } from '@/lib/devtools-test'
 import { SCENARIO_LIST } from '@/lib/tools-test-tools'
 
 /**
@@ -782,13 +783,9 @@ function ToolsTestPage() {
 export const Route = createFileRoute('/tools-test')({
   component: ToolsTestPage,
   validateSearch: (search: Record) => {
-    const port =
-      typeof search.aimockPort === 'string'
-        ? parseInt(search.aimockPort, 10)
-        : undefined
     return {
       testId: typeof search.testId === 'string' ? search.testId : undefined,
-      aimockPort: port != null && !isNaN(port) ? port : undefined,
+      aimockPort: parseAimockPort(search.aimockPort),
       historyFixture:
         typeof search.historyFixture === 'string'
           ? search.historyFixture
diff --git a/testing/e2e/tests/generic-middleware-interrupts.spec.ts b/testing/e2e/tests/generic-middleware-interrupts.spec.ts
new file mode 100644
index 0000000000..083e841fd3
--- /dev/null
+++ b/testing/e2e/tests/generic-middleware-interrupts.spec.ts
@@ -0,0 +1,321 @@
+import type { Page } from '@playwright/test'
+import type { PhaseCapture } from '../src/lib/phase-capture'
+import { expect, test } from './fixtures'
+
+const genericMode = 'generic-lifecycle'
+
+function middlewareUrl(testId: string, aimockPort: number, scenario: string) {
+  const query = new URLSearchParams({
+    testId,
+    aimockPort: String(aimockPort),
+    scenario,
+    middlewareMode: genericMode,
+  })
+  return `/middleware-test?${query}`
+}
+
+async function startScenario(
+  page: Page,
+  testId: string,
+  aimockPort: number,
+  scenario: string,
+) {
+  await page.goto(middlewareUrl(testId, aimockPort, scenario))
+  await page.waitForSelector('#mw-run-button')
+  await page.waitForFunction(
+    () =>
+      document
+        .getElementById('mw-metadata')
+        ?.getAttribute('data-is-loading') === 'false',
+  )
+  // Let React attach delegated handlers before interaction.
+  await page.waitForTimeout(300)
+
+  const readCount = () =>
+    page.evaluate(() =>
+      parseInt(
+        document
+          .getElementById('mw-metadata')
+          ?.getAttribute('data-message-count') || '0',
+        10,
+      ),
+    )
+
+  for (let attempt = 0; attempt < 5; attempt++) {
+    const baseline = await readCount()
+    await page.locator('#mw-run-button').click()
+    const started = await page
+      .waitForFunction(
+        (base) => {
+          const meta = document.getElementById('mw-metadata')
+          if (meta?.getAttribute('data-is-loading') === 'true') return true
+          if (
+            parseInt(meta?.getAttribute('data-interrupt-count') || '0', 10) > 0
+          )
+            return true
+          if (
+            parseInt(meta?.getAttribute('data-message-count') || '0', 10) > base
+          )
+            return true
+          return false
+        },
+        baseline,
+        { timeout: 2000 },
+      )
+      .then(() => true)
+      .catch(() => false)
+    if (started) {
+      await expect(page.getByTestId('generic-review-plan')).toBeVisible()
+      return
+    }
+  }
+
+  throw new Error('Run test button did not start a chat run')
+}
+
+async function resolveReview(page: Page) {
+  await page.getByTestId('resolve-review-plan').click()
+}
+
+async function resolveMixedBatch(page: Page) {
+  await resolveReview(page)
+  await page.getByTestId('approve-delete-review').click()
+}
+
+async function waitForFinished(page: Page) {
+  await expect(page.locator('#mw-metadata')).toHaveAttribute(
+    'data-interrupt-count',
+    '0',
+  )
+  await expect(page.locator('#mw-metadata')).toHaveAttribute(
+    'data-is-loading',
+    'false',
+  )
+}
+
+async function fetchCapture(page: Page, testId: string): Promise {
+  const response = await page.request.get(
+    `/api/middleware-test?testId=${encodeURIComponent(testId)}&kind=phase`,
+  )
+  expect(response.ok()).toBe(true)
+  return response.json()
+}
+
+async function expectAssistantText(page: Page, text: string) {
+  await expect(page.locator('#mw-messages-json')).toContainText(text)
+}
+
+function expectOneInterruptTerminalForEachBoundary(capture: PhaseCapture) {
+  const interruptedRunIds = Array.from(
+    new Set(capture.boundaries.map((boundary) => boundary.runId)),
+  )
+
+  expect(interruptedRunIds).not.toHaveLength(0)
+
+  for (const runId of interruptedRunIds) {
+    const chunks = capture.yieldedChunks.filter(
+      (chunk) => chunk.runId === runId,
+    )
+    const startedIndexes = chunks.flatMap((chunk, index) =>
+      chunk.type === 'RUN_STARTED' ? [index] : [],
+    )
+    const terminalIndexes = chunks.flatMap((chunk, index) =>
+      chunk.type === 'RUN_FINISHED' ? [index] : [],
+    )
+
+    expect(startedIndexes).toHaveLength(1)
+    expect(terminalIndexes).toHaveLength(1)
+    const startedIndex = startedIndexes[0]
+    const terminalIndex = terminalIndexes[0]
+    expect(startedIndex).toBeLessThan(terminalIndex)
+    expect(chunks[terminalIndex]?.outcomeType).toBe('interrupt')
+    expect(terminalIndex).toBe(chunks.length - 1)
+    expect(chunks[terminalIndex - 1]?.type).toBe('MESSAGES_SNAPSHOT')
+  }
+}
+
+const boundaryCases = [
+  {
+    scenario: 'generic-before-model',
+    boundary: 'beforeModel',
+    result: 'BEFORE_MODEL_RESOLVED',
+  },
+  {
+    scenario: 'generic-after-model',
+    boundary: 'afterModel',
+    result: 'AFTER_MODEL_CONTENT',
+  },
+  {
+    scenario: 'generic-before-tools-continue',
+    boundary: 'beforeTools',
+    result: 'TOOLS_CONTINUED',
+    mixed: true,
+  },
+  {
+    scenario: 'generic-after-tools',
+    boundary: 'afterTools',
+    result: 'AFTER_TOOLS_RESOLVED',
+  },
+] as const
+
+for (const boundaryCase of boundaryCases) {
+  test(`resolves a typed generic interrupt at ${boundaryCase.boundary}`, async ({
+    page,
+    testId,
+    aimockPort,
+  }) => {
+    await startScenario(page, testId, aimockPort, boundaryCase.scenario)
+    await expect(page.getByTestId('generic-review-plan')).toHaveAttribute(
+      'data-definition-id',
+      'review-plan',
+    )
+    await expect(page.getByTestId('generic-review-plan')).toHaveAttribute(
+      'data-payload',
+      new RegExp(`"boundary":"${boundaryCase.boundary}"`),
+    )
+
+    if (boundaryCase.mixed) await resolveMixedBatch(page)
+    else await resolveReview(page)
+    await waitForFinished(page)
+    await expectAssistantText(page, boundaryCase.result)
+
+    const capture = await fetchCapture(page, testId)
+    expect(capture.boundaries).toEqual([
+      expect.objectContaining({ phase: boundaryCase.boundary }),
+    ])
+    expect(capture.resolutions).toEqual([
+      expect.objectContaining({
+        definitionId: 'review-plan',
+        status: 'resolved',
+        response: {
+          approved: true,
+          note: 'approved in middleware e2e',
+        },
+      }),
+    ])
+  })
+}
+
+test('cancels a typed generic interrupt and records the cancellation', async ({
+  page,
+  testId,
+  aimockPort,
+}) => {
+  await startScenario(page, testId, aimockPort, 'generic-before-model')
+  await page.getByTestId('cancel-review-plan').click()
+  await waitForFinished(page)
+
+  const capture = await fetchCapture(page, testId)
+  expect(capture.resolutions).toEqual([
+    expect.objectContaining({
+      definitionId: 'review-plan',
+      status: 'cancelled',
+    }),
+  ])
+})
+
+test('keeps generic, approval, and client-tool waits in one terminal batch', async ({
+  page,
+  testId,
+  aimockPort,
+}) => {
+  await startScenario(page, testId, aimockPort, 'generic-before-tools-continue')
+  await expect(page.getByTestId('generic-review-plan')).toHaveCount(1)
+  await expect(page.getByTestId('delete-review-approval')).toHaveCount(1)
+  await expect(page.locator('#mw-metadata')).toHaveAttribute(
+    'data-interrupt-count',
+    '2',
+  )
+
+  await resolveMixedBatch(page)
+  await waitForFinished(page)
+  await expect(page.locator('#mw-metadata')).toHaveAttribute(
+    'data-client-tool-executions',
+    '1',
+  )
+
+  const capture = await fetchCapture(page, testId)
+  const firstTerminal = capture.yieldedChunks.find(
+    (chunk) =>
+      chunk.type === 'RUN_FINISHED' && chunk.interruptCount !== undefined,
+  )
+  expect(firstTerminal?.interruptCount).toBe(3)
+  expect(capture.toolExecutions).toEqual([
+    { name: 'delete_review', side: 'server' },
+  ])
+  expect(capture.policies).toEqual(['continue'])
+})
+
+for (const policyCase of [
+  {
+    scenario: 'generic-before-tools-cancel',
+    policy: 'cancel',
+    result: 'TOOLS_CANCELLED',
+  },
+  {
+    scenario: 'generic-before-tools-stop',
+    policy: 'stop',
+    result: undefined,
+  },
+] as const) {
+  test(`${policyCase.policy} policy prevents pending tool execution`, async ({
+    page,
+    testId,
+    aimockPort,
+  }) => {
+    await startScenario(page, testId, aimockPort, policyCase.scenario)
+    await resolveMixedBatch(page)
+    await waitForFinished(page)
+
+    const capture = await fetchCapture(page, testId)
+    expect(capture.policies).toEqual([policyCase.policy])
+    expect(capture.toolExecutions).toEqual([])
+    await expect(page.locator('#mw-metadata')).toHaveAttribute(
+      'data-client-tool-executions',
+      '0',
+    )
+    if (policyCase.result) await expectAssistantText(page, policyCase.result)
+    else
+      await expect(page.locator('#mw-messages-json')).not.toContainText(
+        'TOOLS_',
+      )
+  })
+}
+
+for (const scenario of [
+  'generic-before-model',
+  'generic-after-model',
+  'generic-before-tools-continue',
+  'generic-before-tools-cancel',
+  'generic-before-tools-stop',
+  'generic-after-tools',
+] as const) {
+  test(`emits one final interrupt terminal for ${scenario}`, async ({
+    page,
+    testId,
+    aimockPort,
+  }) => {
+    await startScenario(page, testId, aimockPort, scenario)
+    const capture = await fetchCapture(page, testId)
+    expectOneInterruptTerminalForEachBoundary(capture)
+  })
+}
+
+test('restores a pending generic interrupt after reload and resumes it', async ({
+  page,
+  testId,
+  aimockPort,
+}) => {
+  await startScenario(page, testId, aimockPort, 'generic-before-model')
+  await page.reload()
+  await expect(page.getByTestId('generic-review-plan')).toBeVisible()
+  await resolveReview(page)
+  await waitForFinished(page)
+  await expectAssistantText(page, 'BEFORE_MODEL_RESOLVED')
+
+  const capture = await fetchCapture(page, testId)
+  expect(capture.boundaries).toHaveLength(1)
+  expect(capture.resolutions).toEqual([
+    expect.objectContaining({ status: 'resolved' }),
+  ])
+})