diff --git a/.changeset/agent-limits-keys-step.md b/.changeset/agent-limits-keys-step.md new file mode 100644 index 000000000..8af37a455 --- /dev/null +++ b/.changeset/agent-limits-keys-step.md @@ -0,0 +1,8 @@ +--- +'@tanstack/ai': minor +--- + +Subagent limits, and provider keys for agents. + +- **`subagents.limits`**: `maxDepth`, `maxConcurrent`, `maxCalls`, and `timeoutMs` for the whole tree of children the model starts through tools. The tree shares one budget (`SubagentBudget`), which a child's `ctx.chat({ subagents })` passes on, so a child cannot reset it. A refused start reaches the model as a tool error. +- **`subagents.keys`** and **`ctx.keys`** (`ProviderKeys`: `get`, `require`, and `adapter`). Pass an object by provider id, for example `{ openai: userKey }`, or a function, for example `(provider) => getByokKey(request, provider)`. An agent builds a `keyedAdapter` with `await ctx.keys.adapter(adapter)`. A missing key falls back to the provider's `env` names. When there is no key at all, the call throws an error that names the env var. A child's `ctx.chat({ subagents })` passes the keys on to its own children. Keys that the child sets itself win for those providers. diff --git a/.changeset/subagent-tool-input.md b/.changeset/subagent-tool-input.md new file mode 100644 index 000000000..a45cbc8dd --- /dev/null +++ b/.changeset/subagent-tool-input.md @@ -0,0 +1,11 @@ +--- +'@tanstack/ai': minor +'@tanstack/ai-persistence': minor +'@tanstack/ai-client': minor +--- + +Continue a child. With `withPersistence` on the parent `chat()`, an agent tool result gives the model the child's `subagentRunId`, and a later call of the same agent tool with `sessionId` continues that child from its stored transcript. With a host that sets `subagents.binding.start`, `background: true` starts the child and returns at once. These fields sit next to the agent's own fields, and each agent tool has them only when the run can use them. When the run uses `sessionId` or `background`, that name is reserved: `chat()` throws when an agent's `inputSchema` uses it. An agent without `inputSchema` takes the task as text in `prompt`, and the child gets it as the last user message. Agent tools also support the `detach` field of the tool context. + +The card keeps the checked input of an agent with `inputSchema`. On the client, `part.subagent.input` has it, typed from `inputSchema` when the agent declaration goes to `useChat({ subagents })`. + +`@tanstack/ai` adds `LoadChildCapability`, `getLoadChild`, and `provideLoadChild`, so a persistence middleware can load a stored child. A provider lists `LoadChildCapability` in `provides`. `withPersistence` in `@tanstack/ai-persistence` provides it: it stores the task of each child's first call and the thread that started it, and loads only the children of the calling thread. diff --git a/docs/advanced/byok.md b/docs/advanced/byok.md index 3a8f49415..b4f399c49 100644 --- a/docs/advanced/byok.md +++ b/docs/advanced/byok.md @@ -326,6 +326,72 @@ byok.setServerCoverage(true); Then a send with no pasted key still POSTs. The relay uses the env key. If that is also empty, the relay returns `byokMissing` (401). The client sets `snapshot.prompt`. +## Keys inside agents + +Your [subagents](../chat/subagents) call models too, and those calls need the user's key. Pass the key on `subagents.keys`. Each agent builds its adapter with `ctx.keys`. + +1. Define the agent. Wrap its adapter in `keyedAdapter`, then build it with `ctx.keys.adapter`: + +```typescript group=agent-keys +import { defineAgent, keyedAdapter } from "@tanstack/ai"; +import { createOpenaiChat } from "@tanstack/ai-openai"; +import { openaiByok } from "@tanstack/ai-openai/byok"; + +const gpt = keyedAdapter(openaiByok, (key) => + createOpenaiChat("gpt-6.1-sol", key), +); + +export const writer = defineAgent({ + name: "writer", + description: "Writes a short draft", + run: async (ctx) => + ctx.chat({ + adapter: await ctx.keys.adapter(gpt), + messages: [{ role: "user", content: "Write a short draft." }], + stream: false, + }), +}); +``` + +2. In the route, give `subagents.keys` a function. `getByokKey` reads the key that the client sends for each provider: + +```typescript group=agent-keys +import { + chat, + chatParamsFromRequest, + toServerSentEventsResponse, +} from "@tanstack/ai"; +import { byokMissing, getByokKey } from "@tanstack/ai/byok/server"; + +export async function POST(request: Request) { + const params = await chatParamsFromRequest(request); + const apiKey = getByokKey(request, openaiByok); + if (!apiKey) return byokMissing(openaiByok); + + const stream = chat({ + adapter: gpt.create(apiKey), + messages: params.messages, + threadId: params.threadId, + runId: params.runId, + subagents: { + agents: [writer], + keys: (provider) => getByokKey(request, provider), + }, + }); + return toServerSentEventsResponse(stream); +} +``` + +How `ctx.keys` finds a key: + +- It calls the `subagents.keys` function with the provider of the adapter, here `openaiByok`. The function can read any store, for example your database. +- `subagents.keys` can also be an object by provider id, for example `{ openai: apiKey }`. +- If there is no key, it reads the env var of the provider, here `OPENAI_API_KEY`. If that is empty too, the call throws `Missing OpenAI API key. Set OPENAI_API_KEY.` +- A child that calls `ctx.chat({ subagents })` passes the keys on to its own children. Keys that the child sets itself win for those providers. +- `ctx.keys.get(provider)` returns the key or `null`. `ctx.keys.require(provider)` throws when there is no key. + +The writer now calls OpenAI with the key that the user pasted. + ## Image, audio, and other providers For other cases: diff --git a/docs/chat/subagents.md b/docs/chat/subagents.md index 2bf5addb7..c2b6af768 100644 --- a/docs/chat/subagents.md +++ b/docs/chat/subagents.md @@ -154,7 +154,7 @@ const stream = chat({ }) ``` -**Without a router.** The library adds one synthetic server tool per agent. The main model calls that tool. The public stream still emits `SUBAGENT_STARTED` / `SUBAGENT_FINISHED` (or `SUBAGENT_ERROR`) and nested parts. The UI does not treat spawn as a normal tool card. The child's events stream while the tool runs, and the child's text becomes the tool result. The child reads the conversation as it is at that tool call. +**Without a router.** The library adds one synthetic server tool per agent. The main model calls that tool. The public stream still emits `SUBAGENT_STARTED` / `SUBAGENT_FINISHED` (or `SUBAGENT_ERROR`) and nested parts. The UI does not treat spawn as a normal tool card. The child's events stream while the tool runs, and the child's text becomes the tool result. The child reads the conversation as it is at that tool call. The model can also pass the task as `prompt`. The child gets it as the last user message. For all fields of the tool input, see [The tool input](#the-tool-input). ## Let the model write the brief @@ -206,6 +206,71 @@ const stream = chat({ To show the brief in the UI, see [Show the brief on a card](#show-the-brief-on-a-card). +## Continue a child + +Each tool call starts a new child. Sometimes the model must send a follow-up to a child that already answered, like "make it shorter". Add `withPersistence` to the parent `chat()`. Each result then gives the model the `subagentRunId` of the child, and the model continues that child with `sessionId`. + +```ts +import { + chat, + chatParamsFromRequest, + defineAgent, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { openaiText } from '@tanstack/ai-openai' +import { memoryPersistence, withPersistence } from '@tanstack/ai-persistence' + +const writer = defineAgent({ + name: 'writer', + description: 'Writes and edits drafts', + run: (ctx) => ctx.chat({ adapter: openaiText('gpt-6.1-sol') }), +}) + +// A continued child loads from this store. Use a durable backend in production. +const persistence = memoryPersistence() + +export async function POST(request: Request) { + const params = await chatParamsFromRequest(request) + const stream = chat({ + adapter: openaiText('gpt-6.1-sol'), + messages: params.messages, + threadId: params.threadId, + runId: params.runId, + middleware: [withPersistence(persistence)], + subagents: { agents: [writer] }, + }) + return toServerSentEventsResponse(stream) +} +``` + +The model continues the writer like this: + +1. The model calls `writer` with `{ "prompt": "Draft a post about tides" }`. +2. It gets `{ "subagentRunId": "subagent-1767225600000-x7k2m9q", "result": "Tides rise twice a day..." }`. +3. In a later turn, it calls `writer` with `{ "sessionId": "subagent-1767225600000-x7k2m9q", "prompt": "Make it shorter" }`. + +The writer gets the task of its first call, its stored transcript, then `Make it shorter` as the last user message. The stream uses the same `subagentRunId`, so the client shows the new work on the same card. + +### The tool input + +Every agent tool takes the agent's own fields, plus two more: + +- The fields of `inputSchema`, for an agent with one. The tool checks them against the schema. +- `prompt`: the task as text, for an agent without `inputSchema`. The child gets it as the last user message. +- `sessionId`: the `subagentRunId` of an earlier result. The call continues that child. The tool has this field only when `withPersistence` is on the parent `chat()`. +- `background`: `true` starts the child and returns its `subagentRunId` at once. The tool has this field only when the host sets `subagents.binding.start`. + +When the run uses `sessionId` or `background`, that name is reserved. If an `inputSchema` has a field with that name, `chat()` throws before the run starts. + +### Wrong calls + +A wrong call goes back to the model as a tool error, and the run continues. For example: + +- `Input validation failed for tool researcher: ...`: the fields do not match `inputSchema`. +- `Unknown sessionId "..."`: no child of this thread has that id. A thread can continue only its own children. +- `Session "..." belongs to agent "writer", not "researcher".`: the call must use the tool of the agent that ran the child. +- `background needs a host that runs children in the background.` + ## Return a value, or call any activity A child does not have to be a chat. It can make an image, or return a plain value that the parent model reads as the tool result. Call the activity on `ctx` and return what it gives you: @@ -261,6 +326,38 @@ const researcher = defineAgent({ A child that `ctx.chat` starts gets the prompt cache retention of the parent call, unless the call sets its own `promptCache`. +## Limit the tree + +A child can have children of its own, and a model can call a child many times. Set `limits` so one request cannot start an unbounded tree: + +```ts group=subagent-values +const limited = chat({ + adapter: openaiText('gpt-6.1-sol'), + messages: [{ role: 'user', content: 'Research squids in depth' }], + subagents: { + agents: [researcher], + limits: { maxDepth: 2, maxConcurrent: 3, maxCalls: 12, timeoutMs: 120_000 }, + }, +}) +``` + +- `maxDepth`: how deep the tree can grow. The first chat is depth 0. +- `maxCalls`: how many children the whole tree can start. +- `maxConcurrent`: how many children one run can have running at once. +- `timeoutMs`: how long one child can run. A child never gets more time than its parent has left. + +The whole tree shares one budget. A child that calls `ctx.chat({ subagents })` passes it on, so a child cannot reset it. A refused start reaches the model as a tool error, for example `subagent limit reached (maxCalls 12)`. + +- Set `limits` on the first `chat()`. A nested chat in a child uses that budget, so its own `limits` do not apply. +- A child that a router picks does not count against the limits. The children it starts do count, one level deeper. +- A child that runs past `timeoutMs` stops, and the model gets `subagent timed out after 120000 ms`. + +## Keys for agents + +An agent that calls a model needs a provider key. An adapter such as `openaiText` reads its env var, for example `OPENAI_API_KEY`. + +To use the key of the user, pass it on `subagents.keys`. The agent then builds its adapter with `ctx.keys`. See [Keys inside agents](../advanced/byok#keys-inside-agents). + ## Strategy - `exclusive` (default): the chosen child owns the turn. Main does not answer after it. @@ -361,7 +458,7 @@ export function CleanupPanel() { The resume uses the plan that the router picked in the first run. It does not call the router again. While the child waits, its card has `status: 'suspended'` and `interruptIds`. After you approve, the card shows the tool result and the child's reply. Client tools in a child work the same way. -The same flow works without a router. The child's tool call stays open until the resume, then the parent model reads the child's result. The result content carries the child's output or error; `subagentRunId` stays in subagent tracking events. +The same flow works without a router. The child's tool call stays open until the resume, then the parent model reads the child's result. The result content carries the child's output or error. It carries `subagentRunId` only when the model can continue the child: with a persistence store, or for a background child. Else the id stays in the subagent events. ## Middleware @@ -629,38 +726,39 @@ When `run` imports server code, do not import the agent into the browser. Declar ### Show the brief on a card -When an agent has `inputSchema`, the brief is the input of the tool call that started the child. That `tool-call` part is on the same assistant message as the card. `part.subagent.parentToolCallId` is its id: +When an agent has `inputSchema`, the card keeps the checked input that the parent model wrote, as `part.subagent.input`. Put `inputSchema` in the shared declaration, and pass the declaration to `useChat` to type it: -```tsx -import type { UIMessage } from '@tanstack/ai-react' - -function briefOf(message: UIMessage, toolCallId: string) { - for (const part of message.parts) { - if (part.type !== 'tool-call' || part.id !== toolCallId) continue - const input = part.input - if ( - typeof input === 'object' && - input !== null && - 'task' in input && - typeof input.task === 'string' - ) { - return input.task - } - } - return undefined -} +```typescript title="agents.ts" +import { z } from 'zod' + +export const researcher = { + name: 'researcher', + description: 'Researches one focused question and returns sourced findings', + inputSchema: z.object({ task: z.string() }), +} as const +``` + +```tsx title="briefs.tsx" +import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' +import { researcher } from './agents' -function Briefs({ message }: { message: UIMessage }) { - return message.parts.map((part) => { - if (part.type !== 'subagent') return null - const toolCallId = part.subagent.parentToolCallId - if (toolCallId === undefined) return null - return

{briefOf(message, toolCallId)}

+export function Briefs() { + const { messages } = useChat({ + connection: fetchServerSentEvents('/api/chat'), + subagents: [researcher], }) + return messages.flatMap((message) => + message.parts.map((part) => + part.type === 'subagent' ? ( +

{part.subagent.input?.task}

+ ) : null, + ), + ) } ``` -A child that a router started has no `parentToolCallId`, so it has no brief. +- `part.subagent.input` is the same value as `ctx.input` in `run`, with the output type of the schema. Here it is `{ task: string } | undefined`. +- It is `undefined` for an agent without `inputSchema`. ### Style one child's parts diff --git a/docs/config.json b/docs/config.json index b203782dd..6f26f97fa 100644 --- a/docs/config.json +++ b/docs/config.json @@ -160,7 +160,7 @@ "label": "Subagents", "to": "chat/subagents", "addedAt": "2026-09-21", - "updatedAt": "2026-10-12" + "updatedAt": "2026-10-13" }, { "label": "Agentic Cycle", @@ -1128,7 +1128,7 @@ "label": "Bring Your Own Key", "to": "advanced/byok", "addedAt": "2026-08-18", - "updatedAt": "2026-10-09" + "updatedAt": "2026-10-13" }, { "label": "Debug Logging", diff --git a/examples/ts-react-chat/src/lib/researcher-agent.ts b/examples/ts-react-chat/src/lib/researcher-agent.ts new file mode 100644 index 000000000..9529e92f2 --- /dev/null +++ b/examples/ts-react-chat/src/lib/researcher-agent.ts @@ -0,0 +1,17 @@ +import { z } from 'zod' + +// Shared by the server and the browser, so it holds no server code. The +// server adds `run` in `routes/api.subagent-brief.ts`. The browser passes it +// to `useChat`, which types `part.subagent.input` from `inputSchema`. +export const researcher = { + name: 'researcher', + description: + 'Researches one focused question and returns short findings. It cannot see this chat, so the task must hold every detail it needs.', + inputSchema: z.object({ + task: z + .string() + .describe( + 'What to find and what to return. The researcher sees only this text.', + ), + }), +} as const diff --git a/examples/ts-react-chat/src/routes/api.subagent-brief.ts b/examples/ts-react-chat/src/routes/api.subagent-brief.ts index 0f229eb84..3175a98cf 100644 --- a/examples/ts-react-chat/src/routes/api.subagent-brief.ts +++ b/examples/ts-react-chat/src/routes/api.subagent-brief.ts @@ -6,23 +6,14 @@ import { toServerSentEventsResponse, } from '@tanstack/ai' import { openaiText } from '@tanstack/ai-openai' -import { z } from 'zod' +import { researcher as researcherDeclaration } from '../lib/researcher-agent' /** * The parent model writes the researcher's brief. The researcher's only * message is that brief, not the parent conversation. */ const researcher = defineAgent({ - name: 'researcher', - description: - 'Researches one focused question and returns short findings. It cannot see this chat, so the task must hold every detail it needs.', - inputSchema: z.object({ - task: z - .string() - .describe( - 'What to find and what to return. The researcher sees only this text.', - ), - }), + ...researcherDeclaration, run: (ctx) => { console.log( '[subagent-brief] child input:', diff --git a/examples/ts-react-chat/src/routes/subagent-brief.tsx b/examples/ts-react-chat/src/routes/subagent-brief.tsx index 704e83a75..86dd2c8a1 100644 --- a/examples/ts-react-chat/src/routes/subagent-brief.tsx +++ b/examples/ts-react-chat/src/routes/subagent-brief.tsx @@ -1,29 +1,13 @@ import { createFileRoute } from '@tanstack/react-router' import { useState } from 'react' import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' +import { researcher } from '../lib/researcher-agent' import type { UIMessage } from '@tanstack/ai-react' export const Route = createFileRoute('/subagent-brief')({ component: SubagentBriefPage, }) -/** The `task` the parent model wrote on the tool call that started a child. */ -function briefOf(message: UIMessage, toolCallId: string) { - for (const part of message.parts) { - if (part.type !== 'tool-call' || part.id !== toolCallId) continue - const input = part.input - if ( - typeof input === 'object' && - input !== null && - 'task' in input && - typeof input.task === 'string' - ) { - return input.task - } - } - return undefined -} - function childText(messages: ReadonlyArray) { return messages .flatMap((message) => @@ -37,6 +21,8 @@ function childText(messages: ReadonlyArray) { function SubagentBriefPage() { const { messages, sendMessage, isLoading } = useChat({ connection: fetchServerSentEvents('/api/subagent-brief'), + // Types `part.subagent.input` from the researcher's inputSchema. + subagents: [researcher], }) const [input, setInput] = useState( 'We picked three sea animals earlier: squid, octopus, and cuttlefish. Find how many hearts each one has.', @@ -63,7 +49,6 @@ function SubagentBriefPage() { ) } if (part.type !== 'subagent') return null - const toolCallId = part.subagent.parentToolCallId return (
@@ -73,10 +58,7 @@ function SubagentBriefPage() {
- Brief:{' '} - {toolCallId === undefined - ? '(none)' - : (briefOf(message, toolCallId) ?? '(none)')} + Brief: {part.subagent.input?.task ?? '(none)'}

{childText(part.subagent.messages)} diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index a94a3d3b0..4a3c19c54 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -8,6 +8,7 @@ import { mergeMetadata, normalizeToUIMessage, restoreInboundChunk, + subagentInputOf, tanstackMetadata, validateWithStandardSchema, } from '@tanstack/ai/client' @@ -3251,6 +3252,14 @@ export class ChatClient< } private syncSubagentHandles(): void { + // The card metadata carries the checked input of an agent with + // `inputSchema`. The handle shows it as `input`. + // A card copied from a handle can hold an old `input`. Read it fresh. + const fieldsOf = (subagent: SubagentHandle) => { + const { input: _old, ...card } = subagent + const input = subagentInputOf(card.metadata) + return input === undefined ? card : { ...card, input } + } // Every card, nested ones included, gets a handle. const cards = collectSubagentParts(this.processor.getMessages()) const present = new Set(cards.map((part) => part.subagent.id)) @@ -3269,12 +3278,12 @@ export class ChatClient< for (const key of Object.keys(existing)) { Reflect.deleteProperty(existing, key) } - Object.assign(existing, part.subagent, { stop }) + Object.assign(existing, fieldsOf(part.subagent), { stop }) part.subagent = existing continue } const handle: SubagentHandle = { - ...part.subagent, + ...fieldsOf(part.subagent), stop: () => { this.stopSubagent(id) }, diff --git a/packages/ai-client/src/types.ts b/packages/ai-client/src/types.ts index e92d914d8..5cbca7c5f 100644 --- a/packages/ai-client/src/types.ts +++ b/packages/ai-client/src/types.ts @@ -643,6 +643,11 @@ export interface SubagentHandle extends SubagentHandleData { * keeps going. */ stop?: () => void + /** + * The checked input of an agent with `inputSchema`: what the parent model + * wrote when it started this child. Set by ChatClient from the card. + */ + input?: unknown } export interface SubagentPart { @@ -660,6 +665,8 @@ export type SubagentClientAgent = { description?: string tools?: ReadonlyArray<{ name: string }> interrupts?: ReadonlyArray> + /** Types `part.subagent.input`. */ + inputSchema?: unknown outputSchema?: SchemaInput } @@ -687,13 +694,28 @@ type AgentOutputData = TAgent extends { outputSchema?: infer TSchema } : unknown : unknown +// The card holds the checked value, so a Standard Schema gives its output type. +type AgentInputData = TAgent extends { inputSchema?: infer TSchema } + ? [Exclude] extends [never] + ? undefined + : Exclude extends { + '~standard': { types?: { output: infer TOutput } } + } + ? TOutput + : Exclude extends SchemaInput + ? InferSchemaType> + : unknown + : undefined + /** One child handle. `name` is the discriminant. */ export type SubagentHandleOf = Omit< SubagentHandle, - 'name' | 'messages' + 'name' | 'messages' | 'input' > & { name: TAgent['name'] messages: Array, AgentOutputData>> + /** Unset on a card saved before the input was kept. */ + input?: AgentInputData } export type SubagentHandles< diff --git a/packages/ai-client/tests/chat-client-subagents.test.ts b/packages/ai-client/tests/chat-client-subagents.test.ts index cb17450e3..a4c3a0a17 100644 --- a/packages/ai-client/tests/chat-client-subagents.test.ts +++ b/packages/ai-client/tests/chat-client-subagents.test.ts @@ -25,12 +25,13 @@ function runFinished(): StreamChunk { } } -function subagentStarted(): StreamChunk { +function subagentStarted(metadata?: Record): StreamChunk { return { type: EventType.SUBAGENT_STARTED, subagentRunId: 'sub-1', name: 'researcher', timestamp: now(), + ...(metadata && { metadata }), } } @@ -93,6 +94,27 @@ describe('ChatClient subagents', () => { expect(part.subagent.stop).toBe(subagents[0]?.stop) }) + it('gives the handle the checked input from the card metadata', async () => { + const chunks: Array = [ + runStarted(), + subagentStarted({ + tanstack: { subagentInput: { task: 'Find squids' } }, + }), + childTextStart(), + childTextContent('Hello'), + subagentFinished(), + runFinished(), + ] + + const client = new ChatClient({ + connection: createMockConnectionAdapter({ chunks }), + }) + await client.sendMessage('hi') + + // The handle keeps the input after later chunks update the card. + expect(client.getSubagents()[0]?.input).toEqual({ task: 'Find squids' }) + }) + it('keeps status error after stop when SUBAGENT_FINISHED arrives later', async () => { let release = () => {} const gate = new Promise((resolve) => { @@ -253,6 +275,7 @@ describe('ChatClient subagent handles for restored messages', () => { name: 'researcher', status: 'finished', messages: [], + metadata: { tanstack: { subagentInput: { task: 'Saved' } } }, }, }, ], @@ -261,7 +284,11 @@ describe('ChatClient subagent handles for restored messages', () => { }) const [handle] = client.getSubagents() - expect(handle).toMatchObject({ id: 'sub-restored', name: 'researcher' }) + expect(handle).toMatchObject({ + id: 'sub-restored', + name: 'researcher', + input: { task: 'Saved' }, + }) expect(typeof handle?.stop).toBe('function') const part = client.getMessages()[0]?.parts[0] expect(part?.type === 'subagent' ? part.subagent : undefined).toBe(handle) diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index 77b0a90b3..d71982be7 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -3,8 +3,10 @@ import { fromSpecTokenUsage, getDetachableRun, InterruptResumeValidationError, + LoadChildCapability, MetadataCapability, modelMessagesToUIMessages, + provideLoadChild, provideMetadata, readInterruptBinding, validateInterruptResumeBatch, @@ -2086,6 +2088,7 @@ export function withPersistence( const provides = [ PersistenceCapability, PersistenceCompletionCapability, + LoadChildCapability, ...(persistence.stores.metadata ? [MetadataCapability] : []), ...(wantsInterrupts ? [InterruptsCapability] : []), ] @@ -2150,6 +2153,16 @@ export function withPersistence( if (persistence.stores.metadata) { provideMetadata(ctx, persistence.stores.metadata) } + // The agent tools read stored children through this, to continue + // one by `sessionId`. Only the children of this chat can load. + provideLoadChild(ctx, (subagentRunId) => + subagentRuns.loadChild(subagentRunId, { + threadId: ctx.threadId, + ...(ctx.subagentRunId !== undefined && { + subagentRunId: ctx.subagentRunId, + }), + }), + ) let resolveCompletion: () => void = () => undefined let rejectCompletion: (error: unknown) => void = () => undefined @@ -2323,6 +2336,9 @@ export function withPersistence( await subagentRuns.chunk({ threadId: ctx.threadId, runId: ctx.runId, + ...(ctx.subagentRunId !== undefined && { + subagentRunId: ctx.subagentRunId, + }), chunk, }) const current = runState.get(ctx) diff --git a/packages/ai-persistence/src/subagent-runs.ts b/packages/ai-persistence/src/subagent-runs.ts index fcbac8255..e6f3d5be5 100644 --- a/packages/ai-persistence/src/subagent-runs.ts +++ b/packages/ai-persistence/src/subagent-runs.ts @@ -3,6 +3,7 @@ import { convertMessagesToModelMessages, modelMessagesToUIMessages, subagentHostMessageId, + subagentInputOf, wireSubagentInfo, } from '@tanstack/ai' import type { @@ -37,6 +38,55 @@ function childStoreId(subagentRunId: string) { return `subagent:${subagentRunId}` } +/** + * Head metadata key for what a later call of the child needs: + * - `base`: how many stored messages came before the current call. A resume + * of that call keeps the number. + * - `parentThreadId`: the thread that started the child. Only that thread can + * continue it. + * - `prompt`: the task of the first call, the first user message of the child. + */ +const CHILD_KEY = 'tanstack:subagentChild' + +function storedMarker(messages: ReadonlyArray) { + const marker: unknown = messages[0]?.metadata?.[CHILD_KEY] + const read = (key: string) => + marker !== null && typeof marker === 'object' + ? Reflect.get(marker, key) + : undefined + const base = read('base') + const parentThreadId = read('parentThreadId') + const prompt = read('prompt') + return { + ...(typeof base === 'number' && { base }), + ...(typeof parentThreadId === 'string' && { parentThreadId }), + ...(typeof prompt === 'string' && { prompt }), + } +} + +/** The thread a chat writes its children under: its own child thread, if any. */ +function ownerThread(caller: { threadId: string; subagentRunId?: string }) { + return caller.subagentRunId === undefined + ? caller.threadId + : childStoreId(caller.subagentRunId) +} + +/** The `prompt` of an agent tool call, if any. */ +function promptOf(call: { args: string } | undefined) { + if (call === undefined) return undefined + let args: unknown + try { + args = JSON.parse(call.args) + } catch { + return undefined + } + const prompt = + args !== null && typeof args === 'object' + ? Reflect.get(args, 'prompt') + : undefined + return typeof prompt === 'string' ? prompt : undefined +} + function readSubagentRunId(chunk: StreamChunk) { if (!('subagentRunId' in chunk)) return const id = chunk.subagentRunId @@ -170,6 +220,13 @@ export function storedSubagentInfo( return wireSubagentInfo(messages[0]) } +/** A stored child transcript without its card-only placeholder. */ +function withoutPlaceholder(messages: ReadonlyArray) { + return messages.filter( + (message) => storedSubagentInfo([message])?.placeholder !== true, + ) +} + type ChildNote = { name: string /** The thread and run that feed this note now. */ @@ -186,6 +243,12 @@ type ChildNote = { interruptIds?: Array error?: { message: string; code?: string } metadata?: Record + /** Stored messages from before the current call. A resume keeps it. */ + base: number + /** The thread that started this child. */ + parentThreadId: string + /** The task of the first call. */ + prompt?: string } export function createSubagentRunRecorder(stores: { @@ -203,6 +266,15 @@ export function createSubagentRunRecorder(stores: { // suspends. Dropped on abort. const answered = new Map>() const parentSavedAt = new Map() + // Name and arguments of each open tool call, by id. A child that a call + // starts reads its task from here. The run's end drops the calls that are + // still open. + const calls = new Map() + + /** Drop the open tool calls of a run that ended. */ + function dropCalls(runId: string) { + for (const [id, call] of calls) if (call.runId === runId) calls.delete(id) + } async function loadMessages(threadId: string) { return stores.messages.loadThread(threadId) @@ -243,14 +315,27 @@ export function createSubagentRunRecorder(stores: { const transcript = convertMessagesToModelMessages( withoutCards(note.processor.getMessages()), ) + const marker = { + [CHILD_KEY]: { + base: note.base, + parentThreadId: note.parentThreadId, + ...(note.prompt !== undefined && { prompt: note.prompt }), + }, + } const [first, ...rest] = transcript const head: ModelMessage = first - ? { ...first, metadata: withSubagentInfo(first.metadata, info) } + ? { + ...first, + metadata: { ...withSubagentInfo(first.metadata, info), ...marker }, + } : { id: `child:${subagentRunId}`, role: 'assistant', content: '', - metadata: withSubagentInfo(undefined, { ...info, placeholder: true }), + metadata: { + ...withSubagentInfo(undefined, { ...info, placeholder: true }), + ...marker, + }, } await stores.messages.saveThread(childStoreId(subagentRunId), [ head, @@ -297,10 +382,17 @@ export function createSubagentRunRecorder(stores: { } async function startChild( - input: { threadId: string; runId: string }, + input: { threadId: string; runId: string; subagentRunId?: string }, chunk: Extract, ) { const id = chunk.subagentRunId + const call = + chunk.parentToolCallId === undefined + ? undefined + : calls.get(chunk.parentToolCallId) + if (chunk.parentToolCallId !== undefined) { + calls.delete(chunk.parentToolCallId) + } const record = await stores.runs?.createOrResume({ runId: id, threadId: childStoreId(id), @@ -317,10 +409,39 @@ export function createSubagentRunRecorder(stores: { delete existing.interruptIds if (chunk.metadata !== undefined) existing.metadata = chunk.metadata } else { - // A resume continues the stored transcript. - const stored = (await loadMessages(childStoreId(id))).filter( - (message) => storedSubagentInfo([message])?.placeholder !== true, - ) + // A resume or a continue goes on from the stored transcript. + const raw = await loadMessages(childStoreId(id)) + const stored = withoutPlaceholder(raw) + const marker = storedMarker(raw) + // A resume keeps the base of the call it resumes. A new call starts + // after the whole stored transcript. + const resumes = storedSubagentInfo(raw)?.status === 'suspended' + // An agent with `inputSchema` has its checked input on the card. An + // agent without one gets the `prompt` of the call as a user message. + const checked = subagentInputOf(chunk.metadata) + // Only the first call gives the task. A later call keeps it. + const prompt = + raw.length > 0 + ? marker.prompt + : checked === undefined + ? promptOf(call) + : JSON.stringify(checked) + // Core does not stream that user message, so store it here. + const followUp = + raw.length > 0 && !resumes && checked === undefined + ? promptOf(call) + : undefined + const initial: Array = + followUp === undefined + ? stored + : [ + ...stored, + { + id: `${childStoreId(id)}:prompt:${stored.length}`, + role: 'user', + content: followUp, + }, + ] children.set(id, { name: chunk.name, threadId: input.threadId, @@ -337,9 +458,16 @@ export function createSubagentRunRecorder(stores: { ...(chunk.metadata !== undefined && { metadata: chunk.metadata }), processor: new StreamProcessor({ subagentRunId: id, - initialMessages: modelMessagesToUIMessages(stored), + initialMessages: modelMessagesToUIMessages(initial), }), status: 'running', + base: resumes ? (marker.base ?? 0) : stored.length, + // A nested child belongs to the thread of the child that started it. + parentThreadId: ownerThread({ + threadId: input.threadId, + subagentRunId: chunk.parentSubagentRunId ?? input.subagentRunId, + }), + ...(prompt !== undefined && { prompt }), }) } if (record && record.status !== 'running') { @@ -352,6 +480,46 @@ export function createSubagentRunRecorder(stores: { await saveChild(id) } + /** + * The stored transcript of a child, for a call that continues it by + * `sessionId`. It starts with the task of the first call as a user message. + * `undefined` when `caller` did not start a child with this id. While the + * child waits on an interrupt, it is the transcript from before that call: + * the resume adds the newer messages of the child itself. + */ + async function loadChild( + subagentRunId: string, + caller: { threadId: string; subagentRunId?: string }, + ) { + const stored = await loadMessages(childStoreId(subagentRunId)) + const marker = storedMarker(stored) + // The id can come from a prompt injection. Never load the child of an + // other thread. + if (marker.parentThreadId !== ownerThread(caller)) return undefined + const messages = withoutPlaceholder(stored) + const info = storedSubagentInfo(stored) + const waits = info?.status === 'suspended' + const transcript = + waits && marker.base !== undefined + ? messages.slice(0, marker.base) + : messages + const task: Array = + marker.prompt === undefined + ? [] + : [ + { + id: `${childStoreId(subagentRunId)}:prompt`, + role: 'user', + content: marker.prompt, + }, + ] + // The agent the child ran under. Only that agent can continue it. + return { + messages: [...task, ...transcript], + ...(info !== undefined && { agent: info.name }), + } + } + async function settleChild( chunk: Extract< StreamChunk, @@ -467,6 +635,8 @@ export function createSubagentRunRecorder(stores: { } return { + loadChild, + async start(input: { threadId: string runId: string @@ -491,9 +661,23 @@ export function createSubagentRunRecorder(stores: { async chunk(input: { threadId: string runId: string + /** Set when the chat that streams this chunk is a child itself. */ + subagentRunId?: string chunk: StreamChunk }) { const chunk = input.chunk + if (chunk.type === 'TOOL_CALL_START') { + calls.set(chunk.toolCallId, { + name: chunk.toolCallName, + args: '', + runId: input.runId, + }) + } + if (chunk.type === 'TOOL_CALL_ARGS') { + const call = calls.get(chunk.toolCallId) + if (call) call.args += chunk.delta + } + if (chunk.type === 'TOOL_CALL_RESULT') calls.delete(chunk.toolCallId) if (chunk.type === 'SUBAGENT_STARTED') { await startChild(input, chunk) return @@ -551,6 +735,7 @@ export function createSubagentRunRecorder(stores: { }, async finish(input: { threadId: string; runId: string }) { + dropCalls(input.runId) await commitAnswers(input.runId) await saveParents(input.threadId) await settleOpenChildren(input.runId, 'completed') @@ -562,6 +747,7 @@ export function createSubagentRunRecorder(stores: { }, async abort(input: { threadId: string; runId: string; error?: unknown }) { + dropCalls(input.runId) const aborted = input.error instanceof Error && input.error.name === 'AbortError' const message = diff --git a/packages/ai-persistence/tests/subagent-persistence-cards.test.ts b/packages/ai-persistence/tests/subagent-persistence-cards.test.ts index 0809cee0b..0671294ef 100644 --- a/packages/ai-persistence/tests/subagent-persistence-cards.test.ts +++ b/packages/ai-persistence/tests/subagent-persistence-cards.test.ts @@ -2,6 +2,8 @@ import { describe, expect, it, vi } from 'vitest' import { EventType, chat, defineAgent, uiMessagesToWire } from '@tanstack/ai' import type { AnyTextAdapter, + ChatMiddleware, + DefinedAgent, ModelMessage, StreamChunk, SubagentPart, @@ -152,6 +154,151 @@ async function collect(stream: AsyncIterable) { return chunks } +/** One text message, start to end. */ +function textMessage( + messageId: string, + text: string, + ids: { subagentRunId?: string } = {}, +) { + return [ + { + type: EventType.TEXT_MESSAGE_START, + messageId, + role: 'assistant', + ...ids, + timestamp: t, + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId, + delta: text, + ...ids, + timestamp: t, + }, + { type: EventType.TEXT_MESSAGE_END, messageId, ...ids, timestamp: t }, + ] satisfies Array +} + +/** A child start, as the parent stream carries it. */ +function childStarted( + subagentRunId: string, + links: { parentToolCallId?: string; parentSubagentRunId?: string }, +) { + return { + type: EventType.SUBAGENT_STARTED, + subagentRunId, + name: 'writer', + ...links, + timestamp: t, + } satisfies StreamChunk +} + +/** Feed one parent run's chunks on thread `desk` to a recorder. */ +async function feed( + recorder: ReturnType, + runId: string, + chunks: ReadonlyArray, +) { + for (const chunk of chunks) { + await recorder.chunk({ threadId: 'desk', runId, chunk }) + } +} + +/** A child that writes one note per call and keeps the messages it saw. */ +function noteTaker() { + const seen: Array> = [] + const agent = defineAgent({ + name: 'writer', + description: 'Writes notes', + run: async function* (ctx) { + seen.push(ctx.messages) + yield* textMessage(`note-${seen.length}`, `Note ${seen.length}`) + }, + }) + return { agent, seen } +} + +/** A parent model that calls the `writer` agent tool once, then stops. */ +function callingSubagent(toolCallId: string, args: object) { + return scriptedAdapter([ + [ + { + type: EventType.TOOL_CALL_START, + toolCallId, + toolCallName: 'writer', + timestamp: t, + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId, + delta: JSON.stringify(args), + timestamp: t, + }, + { + type: EventType.RUN_FINISHED, + runId: 'model', + threadId: 'desk', + finishReason: 'tool_calls', + timestamp: t, + }, + ], + [ + { + type: EventType.RUN_FINISHED, + runId: 'model', + threadId: 'desk', + finishReason: 'stop', + timestamp: t, + }, + ], + ]) +} + +/** One parent turn that calls the `writer` tool. The thread is `desk`. */ +function subagentTurn( + middleware: ChatMiddleware, + agent: DefinedAgent, + turn: { runId: string; toolCallId: string; args: object; threadId?: string }, +) { + return collect( + chat({ + adapter: callingSubagent(turn.toolCallId, turn.args), + threadId: turn.threadId ?? 'desk', + runId: turn.runId, + messages: [{ id: `u-${turn.runId}`, role: 'user', content: 'Notes' }], + middleware: [middleware], + subagents: { agents: [agent] }, + }), + ) +} + +/** The id of the first child a run started. */ +function startedId(chunks: ReadonlyArray) { + const started = chunks.find( + (chunk) => chunk.type === EventType.SUBAGENT_STARTED, + ) + if (started?.type !== EventType.SUBAGENT_STARTED) { + throw new Error('no child started') + } + return started.subagentRunId +} + +/** The parsed result of one tool call. */ +function toolResult(chunks: ReadonlyArray, toolCallId: string) { + const result = chunks.find( + (chunk) => + chunk.type === EventType.TOOL_CALL_RESULT && + chunk.toolCallId === toolCallId, + ) + if ( + result?.type !== EventType.TOOL_CALL_RESULT || + typeof result.content !== 'string' + ) { + throw new Error(`no text result for ${toolCallId}`) + } + return JSON.parse(result.content) +} + describe('persisted subagent cards', () => { it('keeps reasoning, tool calls, and interrupts across a reload', async () => { const persistence = memoryPersistence() @@ -499,4 +646,220 @@ describe('subagent run recorder', () => { code: 'provider_error', }) }) + + it('gives a nested child to the child that started it', async () => { + const { stores } = memoryPersistence() + const recorder = createSubagentRunRecorder({ + messages: stores.messages, + intervalMs: 0, + }) + await feed(recorder, 'run-1', [ + childStarted('outer', { parentToolCallId: 'call_1' }), + childStarted('inner', { + parentSubagentRunId: 'outer', + parentToolCallId: 'call_n', + }), + ...textMessage('m1', 'Inner notes', { subagentRunId: 'inner' }), + ]) + + // The outer child's chat can continue it. The top thread cannot. + const fromOuter = { threadId: 'desk:writer', subagentRunId: 'outer' } + expect((await recorder.loadChild('inner', fromOuter))?.messages).toEqual([ + expect.objectContaining({ role: 'assistant', content: 'Inner notes' }), + ]) + expect( + await recorder.loadChild('inner', { threadId: 'desk' }), + ).toBeUndefined() + }) + + it('loads a waiting continued child without the messages of the call it waits in', async () => { + const { stores } = memoryPersistence() + const recorder = createSubagentRunRecorder({ + messages: stores.messages, + runs: stores.runs, + intervalMs: 0, + }) + const waiting = { + type: EventType.SUBAGENT_FINISHED, + subagentRunId: 'child', + outcome: { type: 'suspended', interruptIds: ['approval'] }, + timestamp: t, + } satisfies StreamChunk + const old = [{ role: 'assistant', content: 'Old notes' }] + await feed(recorder, 'run-1', [ + childStarted('child', { parentToolCallId: 'call_1' }), + ...textMessage('m1', 'Old notes', { subagentRunId: 'child' }), + { + type: EventType.SUBAGENT_FINISHED, + subagentRunId: 'child', + timestamp: t, + }, + ]) + + // run-2 continues the child. It writes, then waits for an approval. The + // resume brings those newer messages itself. + await feed(recorder, 'run-2', [ + childStarted('child', { parentToolCallId: 'call_2' }), + ...textMessage('m2', 'Delete a?', { subagentRunId: 'child' }), + waiting, + ]) + const desk = { threadId: 'desk' } + expect((await recorder.loadChild('child', desk))?.messages).toMatchObject( + old, + ) + + // run-3 resumes that call, and the child waits again. + await feed(recorder, 'run-3', [ + childStarted('child', { parentToolCallId: 'call_2' }), + ...textMessage('m3', 'Delete b?', { subagentRunId: 'child' }), + waiting, + ]) + expect((await recorder.loadChild('child', desk))?.messages).toMatchObject( + old, + ) + }) + + it('starts a loaded child with the task of its first call, and names its agent', async () => { + const { stores } = memoryPersistence() + const recorder = createSubagentRunRecorder({ + messages: stores.messages, + intervalMs: 0, + }) + await feed(recorder, 'run-1', [ + { + type: EventType.TOOL_CALL_START, + toolCallId: 'call_1', + toolCallName: 'writer', + timestamp: t, + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call_1', + delta: '{"topic":"tides"}', + timestamp: t, + }, + // Core puts the checked input of an agent with inputSchema on the card. + { + ...childStarted('child', { parentToolCallId: 'call_1' }), + metadata: { tanstack: { subagentInput: { topic: 'tides' } } }, + }, + ...textMessage('m1', 'Tide notes', { subagentRunId: 'child' }), + ]) + + const loaded = await recorder.loadChild('child', { threadId: 'desk' }) + expect(loaded?.messages).toMatchObject([ + { role: 'user', content: '{"topic":"tides"}' }, + { role: 'assistant', content: 'Tide notes' }, + ]) + // Only this agent can continue the child. + expect(loaded?.agent).toBe('writer') + }) + + it('stores the input of an agent whose schema has a prompt field as JSON', async () => { + const { stores } = memoryPersistence() + const recorder = createSubagentRunRecorder({ + messages: stores.messages, + intervalMs: 0, + }) + await feed(recorder, 'run-1', [ + { + type: EventType.TOOL_CALL_START, + toolCallId: 'call_1', + toolCallName: 'writer', + timestamp: t, + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call_1', + delta: '{"prompt":"a cat"}', + timestamp: t, + }, + { + ...childStarted('child', { parentToolCallId: 'call_1' }), + metadata: { tanstack: { subagentInput: { prompt: 'a cat' } } }, + }, + ...textMessage('m1', 'Cat image', { subagentRunId: 'child' }), + ]) + + const loaded = await recorder.loadChild('child', { threadId: 'desk' }) + // The child got `ctx.input`, not a user message with the text. + expect(loaded?.messages[0]).toMatchObject({ + role: 'user', + content: '{"prompt":"a cat"}', + }) + }) + + it('loads nothing for an unknown child', async () => { + const { stores } = memoryPersistence() + const recorder = createSubagentRunRecorder({ messages: stores.messages }) + + expect( + await recorder.loadChild('missing', { threadId: 'desk' }), + ).toBeUndefined() + }) +}) + +describe('agent tools with persistence', () => { + it('continues a stored child by sessionId', async () => { + const middleware = withPersistence(memoryPersistence()) + const { agent, seen } = noteTaker() + const first = await subagentTurn(middleware, agent, { + runId: 'run-1', + toolCallId: 'call_1', + args: { prompt: 'Take notes' }, + }) + expect(toolResult(first, 'call_1')).toMatchObject({ + subagentRunId: startedId(first), + }) + const second = await subagentTurn(middleware, agent, { + runId: 'run-2', + toolCallId: 'call_2', + args: { sessionId: startedId(first), prompt: 'More' }, + }) + + expect(startedId(second)).toBe(startedId(first)) + // The child starts from its first task and its stored transcript, then + // the new prompt. The first message is a user message. + expect(seen[1]).toMatchObject([ + { role: 'user', content: 'Take notes' }, + { role: 'assistant', content: 'Note 1' }, + { role: 'user', content: 'More' }, + ]) + + // A third call also sees the prompt of the second call. + await subagentTurn(middleware, agent, { + runId: 'run-3', + toolCallId: 'call_3', + args: { sessionId: startedId(first), prompt: 'Again' }, + }) + expect(seen[2]).toMatchObject([ + { role: 'user', content: 'Take notes' }, + { role: 'assistant', content: 'Note 1' }, + { role: 'user', content: 'More' }, + { role: 'assistant', content: 'Note 2' }, + { role: 'user', content: 'Again' }, + ]) + }) + + it('does not continue a child of another thread', async () => { + const middleware = withPersistence(memoryPersistence()) + const { agent, seen } = noteTaker() + const first = await subagentTurn(middleware, agent, { + runId: 'run-1', + toolCallId: 'call_1', + args: { prompt: 'Take notes' }, + }) + const id = startedId(first) + const other = await subagentTurn(middleware, agent, { + threadId: 'other', + runId: 'run-2', + toolCallId: 'call_2', + args: { sessionId: id, prompt: 'Show me' }, + }) + + expect(toolResult(other, 'call_2')).toEqual({ + error: `Unknown sessionId "${id}".`, + }) + expect(seen).toHaveLength(1) + }) }) diff --git a/packages/ai-react/tests/use-chat-subagent-types.test.ts b/packages/ai-react/tests/use-chat-subagent-types.test.ts index 4e404c9f3..585892670 100644 --- a/packages/ai-react/tests/use-chat-subagent-types.test.ts +++ b/packages/ai-react/tests/use-chat-subagent-types.test.ts @@ -40,3 +40,50 @@ it('useChat types subagent parts from the server agents', () => { } void check }) + +it('types part.subagent.input from the agent inputSchema', () => { + const check = () => { + const inputSchema = { + '~standard': { + version: 1 as const, + vendor: 'test', + types: { input: { task: '' }, output: { task: '' } }, + validate: () => ({ value: { task: '' } }), + jsonSchema: { + input: () => ({ + type: 'object', + properties: { task: { type: 'string' } }, + }), + output: () => ({}), + }, + }, + } + const researcher = defineAgent({ + name: 'researcher', + description: 'Looks up facts', + inputSchema, + run: async function* () {}, + }) + const writer = defineAgent({ + name: 'writer', + description: 'Writes the draft', + run: async function* () {}, + }) + + const chat = useChat({ + connection: { connect: async function* () {} }, + subagents: [researcher, writer], + }) + + const part = chat.messages[0]?.parts[0] + if (part?.type !== 'subagent') return + if (part.subagent.name === 'researcher') { + expectTypeOf(part.subagent.input).toEqualTypeOf< + { task: string } | undefined + >() + } else { + expectTypeOf(part.subagent.input).toEqualTypeOf() + } + } + void check +}) diff --git a/packages/ai/src/activities/chat/agents/bound.ts b/packages/ai/src/activities/chat/agents/bound.ts index a63350bfc..0b4008797 100644 --- a/packages/ai/src/activities/chat/agents/bound.ts +++ b/packages/ai/src/activities/chat/agents/bound.ts @@ -15,6 +15,8 @@ import type { GenerationMiddleware } from '../../middleware/types' import type { AnyChatMiddleware } from '../middleware/types' import type { PromptCacheRetention, RunAgentResumeItem } from '../../../types' import type { SubagentRunInput } from './define-agent' +import type { SubagentBudget } from './limits' +import type { ProviderKeys } from '../../../byok/keyed' /** * The fields a child `chat()` needs, in one spread: @@ -51,11 +53,37 @@ export interface SubagentBinding { * child's `ctx.chat({ subagents })` passes it down to its own children. */ generationMiddleware?: ReadonlyArray + /** + * The subagent tree budget. A child's `ctx.chat({ subagents })` passes it + * down, so limits hold across the whole tree. + */ + budget?: SubagentBudget + /** + * The provider keys an agent reads as `ctx.keys`. `subagents.keys` sets + * them. A child's `ctx.chat({ subagents })` passes them down to its own + * children. Without them, `ctx.keys` reads each provider's `env` names. + */ + keys?: ProviderKeys /** * The prompt cache retention of the parent call. A child's `ctx.chat()` * uses it when the call gives no `promptCache`. */ promptCache?: PromptCacheRetention + /** + * Starts a child without waiting for it and returns its id at once. The + * agent tools call it for `background: true`. With + * `wake: true`, the host starts a new parent turn when the child ends. It + * does not pass down to nested children. + */ + start?: ( + call: { + agent: string + input?: unknown + prompt?: string + parentToolCallId?: string + }, + options: { wake: true }, + ) => Promise<{ subagentRunId: string }> } /** @@ -132,12 +160,14 @@ export function createBoundActivities( ...(binding?.promptCache ? { promptCache: binding.promptCache } : {}), abortController, ...options, - // Nested children get the host middleware before their own. + // Nested children get the host middleware before their own, share + // the tree budget, and get the host keys unless the call sets its own. ...(options.subagents ? { subagents: { ...options.subagents, binding: { + ...(binding?.keys ? { keys: binding.keys } : {}), ...options.subagents.binding, chatMiddleware: [ ...chatMiddleware, @@ -147,6 +177,7 @@ export function createBoundActivities( ...generationMiddleware, ...(options.subagents.binding?.generationMiddleware ?? []), ], + ...(binding?.budget ? { budget: binding.budget } : {}), }, }, } diff --git a/packages/ai/src/activities/chat/agents/define-agent.ts b/packages/ai/src/activities/chat/agents/define-agent.ts index df0a0c0ec..8edf16529 100644 --- a/packages/ai/src/activities/chat/agents/define-agent.ts +++ b/packages/ai/src/activities/chat/agents/define-agent.ts @@ -11,6 +11,7 @@ import type { } from '../../../types' import type { AnyClientTool } from '../tools/tool-definition' import type { BoundActivities, SubagentForward } from './bound' +import type { ProviderKeys } from '../../../byok/keyed' /** * What the library knows about a child run before `run` starts. @@ -55,11 +56,16 @@ export interface SubagentRunInput< * - The activity functions (`ctx.chat`, `ctx.generateImage`, and the rest) * take the same options as the plain functions and fill in the thread id, * a run id, the abort signal, and any middleware a host adds. + * - `keys` finds provider keys. `await ctx.keys.adapter(adapter)` builds a + * `keyedAdapter(...)`. The keys come from `subagents.keys`, else from each + * provider's `env` names. A missing key throws an error that names the env + * var. */ export type SubagentRunContext = SubagentRunInput & BoundActivities & { forward: SubagentForward + keys: ProviderKeys } /** @@ -115,8 +121,11 @@ export interface DefinedAgent< produces?: TProduces /** * The input the parent model writes when it calls this agent's tool, such - * as a short brief. `run` reads it as `ctx.input`. Tool mode only: a - * `subagents.router` cannot start an agent that has `inputSchema`. + * as a short brief. `run` reads it as `ctx.input`, and the client card + * as `part.subagent.input`. When the run uses `sessionId` or + * `background`, that name is reserved. Without it, the tool takes a text + * `prompt`. Tool mode only: a `subagents.router` cannot start an agent + * that has `inputSchema`. */ inputSchema?: TInput tools?: TTools diff --git a/packages/ai/src/activities/chat/agents/limits.ts b/packages/ai/src/activities/chat/agents/limits.ts new file mode 100644 index 000000000..7210880ec --- /dev/null +++ b/packages/ai/src/activities/chat/agents/limits.ts @@ -0,0 +1,75 @@ +/** + * Limits for a whole subagent tree. The root run makes one budget, and every + * child shares it, so a child cannot reset it. + */ +export interface SubagentLimits { + /** How deep the tree may grow. The root run is depth 0. */ + maxDepth?: number + /** How many children one run may have running at once. */ + maxConcurrent?: number + /** How many children the whole tree may start. */ + maxCalls?: number + /** How long one child may run. Never later than its parent's own deadline. */ + timeoutMs?: number +} + +/** The shared budget of one subagent tree. */ +export class SubagentBudget { + private constructor( + readonly limits: SubagentLimits, + /** The depth of the run that owns this budget view. */ + readonly depth: number, + private readonly shared: { calls: number }, + /** Epoch ms after which this run's children must stop, if any. */ + readonly deadline: number | undefined, + ) {} + + /** A budget for a root run. */ + static root(limits: SubagentLimits = {}): SubagentBudget { + return new SubagentBudget(limits, 0, { calls: 0 }, undefined) + } + + /** How many children the tree started so far. */ + get calls(): number { + return this.shared.calls + } + + /** + * Reserve one child spawn. Returns the refusal message the model sees, or + * `undefined` when the child may start. `active` is how many children this + * run has running now. + */ + reserve(active: number): string | undefined { + const { maxDepth, maxCalls, maxConcurrent } = this.limits + if (maxDepth !== undefined && this.depth + 1 > maxDepth) { + return `subagent limit reached (maxDepth ${maxDepth})` + } + if (maxCalls !== undefined && this.shared.calls >= maxCalls) { + return `subagent limit reached (maxCalls ${maxCalls})` + } + if (maxConcurrent !== undefined && active >= maxConcurrent) { + return `subagent limit reached (maxConcurrent ${maxConcurrent})` + } + this.shared.calls += 1 + return undefined + } + + /** Milliseconds a child started now may run, or `undefined` for no limit. */ + childTimeout(now = Date.now()): number | undefined { + const own = this.limits.timeoutMs + const remaining = + this.deadline === undefined ? undefined : Math.max(0, this.deadline - now) + if (own === undefined) return remaining + return remaining === undefined ? own : Math.min(own, remaining) + } + + /** The budget a child's own subagents use: one level deeper, same counters. */ + child(deadline?: number): SubagentBudget { + return new SubagentBudget( + this.limits, + this.depth + 1, + this.shared, + deadline ?? this.deadline, + ) + } +} diff --git a/packages/ai/src/activities/chat/agents/spawn.ts b/packages/ai/src/activities/chat/agents/spawn.ts index 11e0e8ec6..4bd735459 100644 --- a/packages/ai/src/activities/chat/agents/spawn.ts +++ b/packages/ai/src/activities/chat/agents/spawn.ts @@ -10,13 +10,26 @@ import { withTanstackMetadata, } from '../../../utilities/merge-metadata' import { mergeStreams } from '../../../utilities/merge-streams' +import { SUBAGENT_INPUT_KEY } from '../../../utilities/subagent-wire' import { INTERRUPT_BINDING_METADATA_KEY } from '../../../interrupt-resume' import { EMIT_STREAM_CHUNK, SUBAGENT_TOOL } from '../tools/tool-calls' +import { + convertSchemaToJsonSchema, + isStandardSchema, +} from '../tools/schema-converter' +import { getLoadChild } from '../middleware/load-child' +import type { LoadChild } from '../middleware/load-child' import type { SubagentToolOutcome } from '../tools/tool-calls' import type { SpecTokenUsage } from '../../../utilities/ag-ui-usage' +import type { + StandardJSONSchemaV1, + StandardSchemaV1, +} from '@standard-schema/spec' import type { Interrupt, + JSONSchema, ModelMessage, + SchemaInput, RunAgentResumeItem, StreamChunk, SubagentErrorEvent, @@ -24,10 +37,15 @@ import type { SubagentStartedEvent, TokenUsage, Tool, + ToolExecutionContext, UIMessage, } from '../../../types' +import { envProviderKeys } from '../../../byok/env-keys' import { createBoundActivities } from './bound' +import { SubagentBudget } from './limits' import type { SubagentBinding } from './bound' +import type { SubagentLimits } from './limits' +import type { ProviderKeySource } from '../../../byok/env-keys' import type { DefinedAgent, SubagentRunContext, @@ -87,6 +105,78 @@ export interface SubagentsBag< * (`ctx.chat`, `ctx.generateImage`, and the rest). Apps do not set it. */ binding?: SubagentBinding + /** + * Limits for the whole tree of children the model starts through tools: + * depth, total calls, children at once, and time per child. A refused + * start returns to the model as a tool error. Only the root chat's limits + * apply: a nested chat in a child uses the root's budget. + */ + limits?: SubagentLimits + /** + * Provider keys that agents read as `ctx.keys`: an object by provider id, + * for example `{ openai: userKey }`, or a function that returns the key of + * a provider, for example `(provider) => getByokKey(request, provider)`. + * A missing or empty key falls back to the keys of the parent agent, then + * to the provider's `env` names. A child's `ctx.chat({ subagents })` passes + * them on. The keys it sets itself win for those providers. + */ + keys?: ProviderKeySource +} + +/** + * Fields an agent tool takes next to the agent's own fields, when the run + * can use them. An `inputSchema` cannot use a name that the run uses. + * - `sessionId`: the `subagentRunId` of a stored child. The call continues + * it. Needs a middleware that provides `LoadChildCapability`. + * - `background`: start the child and return at once. Needs `binding.start`. + */ +const RESERVED_FIELDS = ['sessionId', 'background'] as const + +type ReservedField = (typeof RESERVED_FIELDS)[number] + +/** The reserved fields that a run uses. */ +function reservedFields( + bag: Pick, + canContinue: boolean, +) { + const can: Record = { + sessionId: canContinue, + background: bag.binding?.start !== undefined, + } + return RESERVED_FIELDS.filter((name) => can[name]) +} + +/** Throws when an agent's `inputSchema` uses a reserved field of the run. */ +export function assertNoReservedFields( + bag: Pick, + canContinue: boolean, +) { + const names = reservedFields(bag, canContinue) + for (const agent of names.length > 0 ? bag.agents : []) { + const schema = convertSchemaToJsonSchema(agent.inputSchema, { io: 'input' }) + const taken = names.find((name) => schema?.properties?.[name]) + if (taken !== undefined) { + throw new Error( + `Agent "${agent.name}" uses the reserved input field "${taken}". Rename that field.`, + ) + } + } +} + +/** Split the reserved fields off a tool call. */ +function splitReserved(value: unknown, names: ReadonlyArray) { + if (typeof value !== 'object' || value === null || Array.isArray(value)) { + return { fields: value, reserved: {} as Record } + } + const fields: Record = {} + const reserved: Record = {} + for (const [key, entry] of Object.entries(value)) { + const target = (names as ReadonlyArray).includes(key) + ? reserved + : fields + target[key] = entry + } + return { fields, reserved } } /** What the children of one parent run left behind for the parent terminal. */ @@ -107,6 +197,13 @@ export interface SpawnEntry { name: string /** The checked tool input for an agent with `inputSchema`. */ input?: unknown + /** Text for an agent without `inputSchema`. One user message at the end. */ + prompt?: string + /** A stored, finished child to run again from its transcript. */ + continued?: { + subagentRunId: string + messages: Array + } resume?: { subagentRunId: string /** The child's own messages from the interrupted run. */ @@ -236,17 +333,42 @@ function agentByName(agents: ReadonlyArray, name: string) { return agent } +/** The root budget of each bag that has no budget from its parent. */ +const rootBudgets = new WeakMap() + +/** + * The subagent tree budget of `bag`: the one its parent passed down, else + * one root budget per bag. The `limits` of a bag apply only at the root. + */ +function treeBudget(bag: SubagentsBag) { + const inherited = bag.binding?.budget + if (inherited) return inherited + let budget = rootBudgets.get(bag) + if (!budget) { + budget = SubagentBudget.root(bag.limits) + rootBudgets.set(bag, budget) + } + return budget +} + function openAgentStream( entry: SpawnEntry, bag: SubagentsBag, ctx: SpawnContext, sink?: SubagentSink, parentToolCallId?: string, - binding: SubagentBinding | undefined = bag.binding, + // A child that a router picks does not count against the budget, but its + // own children do, one level deeper. + binding: SubagentBinding = { + ...bag.binding, + budget: treeBudget(bag).child(), + }, ) { const agent = agentByName(bag.agents, entry.name) const resume = entry.resume - const subagentRunId = resume?.subagentRunId ?? createSubagentId() + const continued = entry.continued + const subagentRunId = + resume?.subagentRunId ?? continued?.subagentRunId ?? createSubagentId() // The parent binds child interrupts to its own run (see rebindInterrupts), // so the resumed child continues from the interrupted parent run id. if (resume !== undefined && ctx.interruptedRunId === undefined) { @@ -254,10 +376,29 @@ function openAgentStream( `Subagent "${entry.name}" has interrupt answers, but the run has no parentRunId. Pass the interrupted run id as parentRunId.`, ) } + // A continued child starts from its stored transcript, not from the parent + // messages. A prompt is one more user message at the end. + const base = continued?.messages ?? ctx.messages + const messages = + entry.prompt === undefined + ? base + : [ + ...base, + { role: 'user', content: entry.prompt } satisfies ModelMessage, + ] + // The card of a continued child also holds its stored messages. Keep only + // the messages that the stored transcript does not have. + const storedIds = new Set(continued?.messages.map((message) => message.id)) const resumed = resume !== undefined ? { - messages: [...ctx.messages, ...resume.messages], + messages: [ + ...messages, + ...resume.messages.filter( + (message) => + message.id === undefined || !storedIds.has(message.id), + ), + ], parentRunId: ctx.interruptedRunId, resume: resume.entries, } @@ -266,7 +407,7 @@ function openAgentStream( agent, { input: entry.input, - messages: resumed?.messages ?? ctx.messages, + messages: resumed?.messages ?? messages, ...(ctx.abortSignal ? { abortSignal: ctx.abortSignal } : {}), threadId: childThreadId(bag.sandbox, ctx.threadId, entry.name), runId: childRunId(ctx.parentRunId, subagentRunId), @@ -494,8 +635,8 @@ function* valueResultChunks( } /** - * Build the context `run` receives: the spawn input plus `forward` and the - * bound activity functions. + * Build the context `run` receives: the spawn input plus `forward`, the + * provider keys, and the bound activity functions. */ function runContext( agentName: string, @@ -518,6 +659,7 @@ function runContext( abortController, ...(binding?.promptCache ? { promptCache: binding.promptCache } : {}), }, + keys: binding?.keys ?? envProviderKeys, ...createBoundActivities(agentName, input, abortController, binding), } satisfies SubagentRunContext } @@ -541,7 +683,7 @@ export async function* spawnAgentStream( // No parentSubagentRunId here, also for a nested child: the chat that // started this child lists it as a direct child. attributeChunk adds the id // one level up, in the stream of the parent child. - yield { + const started = { type: SUBAGENT_STARTED, subagentRunId: id, name: agent.name, @@ -549,6 +691,10 @@ export async function* spawnAgentStream( ...(parentToolCallId !== undefined ? { parentToolCallId } : {}), timestamp: Date.now(), } satisfies SubagentStartedEvent + // The card keeps the checked input, so a UI can show the brief. + yield input.input === undefined + ? started + : withTanstackMetadata(started, { [SUBAGENT_INPUT_KEY]: input.input }) // No yield between here and the `try`, so `finally` always disposes it. const link = linkAbort(input.abortSignal) @@ -800,12 +946,17 @@ function messagesBeforeCall( /** * Record the parent messages when the model calls a subagent tool, so the - * child reads the conversation as it is at that call. + * child reads the conversation as it is at that call. Also read the run's + * `loadChild` service, so a call can continue a stored child. */ export function subagentCallMessages(names: ReadonlySet) { const byCall = new Map>() + let loadChild: LoadChild | undefined const middleware: ChatMiddleware = { name: 'subagent-call-messages', + onStart(ctx) { + loadChild = getLoadChild(ctx, { optional: true }) + }, onBeforeToolCall(ctx, hook) { if (!names.has(hook.toolName)) return undefined byCall.set( @@ -819,6 +970,131 @@ export function subagentCallMessages(names: ReadonlySet) { middleware, messagesFor: (toolCallId: string | undefined) => toolCallId === undefined ? undefined : byCall.get(toolCallId), + childLoader: () => loadChild, + } +} + +/** + * The wire schema of an agent tool: the agent's `inputSchema`, or a `prompt` + * for an agent without one. The reserved fields of the run sit next to + * those fields. + */ +function agentToolSchema( + agent: DefinedAgent, + names: ReadonlyArray, +): JSONSchema { + const schema = convertSchemaToJsonSchema(agent.inputSchema, { + io: 'input', + }) ?? { + type: 'object', + properties: { prompt: { type: 'string' } }, + } + // ponytail: a schema with no top-level properties (a top-level anyOf) gets + // no sessionId or background. Wrap it in allOf if one needs them. + if (schema.properties === undefined || names.length === 0) return schema + return { + ...schema, + properties: { + ...schema.properties, + ...(names.includes('sessionId') && { + sessionId: { + type: 'string', + description: + 'The subagentRunId of an earlier result. The call continues that child.', + }, + }), + ...(names.includes('background') && { + background: { + type: 'boolean', + description: + 'true starts the child and returns its subagentRunId at once.', + }, + }), + }, + } +} + +/** + * The input schema of an agent tool. The tool loop checks a Standard Schema + * before middleware runs, as for any tool. With reserved fields, a wrapper + * checks the agent's fields and passes the reserved ones through. + */ +function agentToolInput( + agent: DefinedAgent, + names: ReadonlyArray, +): SchemaInput { + const inner = agent.inputSchema + if (isStandardSchema(inner) && names.length === 0) return inner + const schema = agentToolSchema(agent, names) + if (!isStandardSchema(inner)) return schema + const toJson = () => schema + const wrapped: StandardSchemaV1 & StandardJSONSchemaV1 = { + '~standard': { + version: 1, + vendor: 'tanstack-ai', + validate: (value) => { + const { fields, reserved } = splitReserved(value, names) + // ponytail: a schema that turns the input into a non-object (a + // top-level transform) loses the reserved fields. + const add = (result: StandardSchemaV1.Result) => + result.issues !== undefined || + Object.keys(reserved).length === 0 || + typeof result.value !== 'object' || + result.value === null || + Array.isArray(result.value) + ? result + : { + value: { + ...(result.value as Record), + ...reserved, + }, + } + const result = inner['~standard'].validate(fields) + return result instanceof Promise ? result.then(add) : add(result) + }, + jsonSchema: { input: toJson, output: toJson }, + }, + } + return wrapped +} + +/** + * Split a checked agent tool call into the agent's input and the reserved + * fields. A reserved field of the wrong type throws a tool error. + */ +function readCall( + agent: DefinedAgent, + args: unknown, + names: ReadonlyArray, +): { + input?: unknown + prompt?: string + sessionId?: string + background: boolean +} { + const { fields, reserved } = splitReserved(args, names) + const { sessionId, background } = reserved + const prompt = + agent.inputSchema === undefined && + typeof fields === 'object' && + fields !== null + ? Reflect.get(fields, 'prompt') + : undefined + const wrongType = + (prompt !== undefined && typeof prompt !== 'string') || + (sessionId !== undefined && typeof sessionId !== 'string') || + (background !== undefined && typeof background !== 'boolean') + if (wrongType) { + throw new Error( + 'prompt and sessionId must be strings, and background a boolean.', + ) + } + return { + // The tool loop checked these fields against `inputSchema`. + ...(agent.inputSchema !== undefined && { input: fields }), + ...(typeof prompt === 'string' && { prompt }), + ...(typeof sessionId === 'string' && { sessionId }), + background: background === true, } } @@ -839,108 +1115,225 @@ export function createSyntheticSubagentTools( abortSignal?: AbortSignal turn?: SubagentTurn sink: SubagentSink + /** The run's `loadChild` service. See subagentCallMessages. */ + childLoader?: () => LoadChild | undefined + /** A middleware provides `loadChild`, so the tools take `sessionId`. */ + canContinue?: boolean }, ): Array { + // One budget per root run. A child run inherits its parent's budget. + const budget = treeBudget(bag) + let active = 0 + + /** Run one child for a tool call. The model gets its text or result. */ + async function runChild(child: SpawnEntry, context: unknown) { + const toolContext = context as + | { + toolCallId?: string + [EMIT_STREAM_CHUNK]?: (chunk: StreamChunk) => void + } + | undefined + const toolCallId = toolContext?.toolCallId + const suspended = parent.turn?.children.find( + (turnChild) => + turnChild.status === 'suspended' && + turnChild.parentToolCallId !== undefined && + turnChild.parentToolCallId === toolCallId, + ) + const entry: SpawnEntry = { + ...child, + ...(suspended && { + resume: { + subagentRunId: suspended.subagentRunId, + messages: suspended.messages, + entries: suspended.resume, + text: suspended.text, + }, + }), + } + const timeout = budget.childTimeout() + // A resumed child already counted when it first started. A parent with + // no time left starts no child. + const refusal = suspended + ? undefined + : timeout === 0 + ? `subagent limit reached (timeoutMs ${budget.limits.timeoutMs ?? 0})` + : budget.reserve(active) + if (refusal !== undefined) { + return { + subagentRunId: '', + text: '', + error: refusal, + } satisfies SubagentToolOutcome + } + active += 1 + const sink = createSubagentSink() + const link = linkAbort(parent.abortSignal) + let timedOut = false + const timer = + timeout === undefined + ? undefined + : setTimeout(() => { + timedOut = true + link.controller.abort( + new Error(`subagent timed out after ${timeout} ms`), + ) + }, timeout) + const childBinding: SubagentBinding = { + ...bag.binding, + budget: budget.child( + timeout === undefined ? undefined : Date.now() + timeout, + ), + } + let subagentRunId = suspended?.subagentRunId ?? '' + let text = suspended?.text ?? '' + let error: string | undefined + let result: unknown + try { + for await (const chunk of openAgentStream( + entry, + bag, + { + messages: parent.messagesFor?.(toolCallId) ?? parent.messages, + abortSignal: link.controller.signal, + threadId: parent.threadId, + parentRunId: parent.runId, + ...(parent.interruptedRunId !== undefined + ? { interruptedRunId: parent.interruptedRunId } + : {}), + ...(parent.parentSubagentRunId !== undefined + ? { parentSubagentRunId: parent.parentSubagentRunId } + : {}), + }, + sink, + toolCallId, + childBinding, + )) { + if (chunk.type === SUBAGENT_STARTED && subagentRunId === '') { + subagentRunId = chunk.subagentRunId + } + if ( + chunk.type === EventType.TEXT_MESSAGE_CONTENT && + 'subagentRunId' in chunk && + chunk.subagentRunId === subagentRunId + ) { + text += chunk.delta + } + if ( + chunk.type === SUBAGENT_ERROR && + chunk.subagentRunId === subagentRunId + ) { + error = chunk.message + } + if ( + chunk.type === SUBAGENT_FINISHED && + chunk.subagentRunId === subagentRunId && + chunk.result !== undefined + ) { + result = chunk.result + } + toolContext?.[EMIT_STREAM_CHUNK]?.(chunk) + } + } finally { + active -= 1 + if (timer !== undefined) clearTimeout(timer) + link.dispose() + } + if (timedOut) error = `subagent timed out after ${timeout} ms` + parent.sink.usage.push(...sink.usage) + if (sink.total) { + parent.sink.total = parent.sink.total + ? addTokenUsage(parent.sink.total, sink.total) + : sink.total + } + return { + subagentRunId, + text, + ...(result !== undefined ? { result } : {}), + ...(error !== undefined ? { error } : {}), + ...(sink.interrupts.length > 0 ? { interrupts: sink.interrupts } : {}), + // Only a stored child can be continued, through a tool with `sessionId`. + ...(subagentRunId !== '' && parent.canContinue === true + ? { keepRunId: true as const } + : {}), + } satisfies SubagentToolOutcome + } + + /** + * The stored child that a `sessionId` names, for `agent`. Throws a tool + * error, also when the child ran under another agent. + */ + async function storedChild(sessionId: string, agent: string) { + const load = parent.childLoader?.() + if (!load) throw new Error('sessionId needs a persistence store.') + const stored = await load(sessionId) + if (!stored) throw new Error(`Unknown sessionId "${sessionId}".`) + if (stored.agent !== undefined && stored.agent !== agent) { + throw new Error( + `Session "${sessionId}" belongs to agent "${stored.agent}", not "${agent}".`, + ) + } + return { subagentRunId: sessionId, messages: stored.messages } + } + + const names = reservedFields(bag, parent.canContinue === true) + // A thrown error in `execute` goes back to the model as a tool error. return bag.agents.map((agent) => ({ name: agent.name, description: agent.description, - ...(agent.inputSchema !== undefined && { inputSchema: agent.inputSchema }), + inputSchema: agentToolInput(agent, names), [SUBAGENT_TOOL]: true, - // The tool loop checks `input` against `inputSchema` before this runs. - execute: async (input: unknown, context?: unknown) => { - const toolContext = context as - | { - toolCallId?: string - [EMIT_STREAM_CHUNK]?: (chunk: StreamChunk) => void - } - | undefined - const toolCallId = toolContext?.toolCallId - const suspended = parent.turn?.children.find( - (child) => - child.status === 'suspended' && - child.parentToolCallId !== undefined && - child.parentToolCallId === toolCallId, - ) - // An agent without `inputSchema` still gets `{}` from the model. Its - // `ctx.input` stays undefined. - const entry: SpawnEntry = { - name: agent.name, - ...(agent.inputSchema !== undefined && { input }), - ...(suspended && { - resume: { - subagentRunId: suspended.subagentRunId, - messages: suspended.messages, - entries: suspended.resume, - text: suspended.text, - }, - }), + execute: async (args: unknown, context?: ToolExecutionContext) => { + const call = readCall(agent, args, names) + const fields = { + ...(call.input !== undefined && { input: call.input }), + ...(call.prompt !== undefined && { prompt: call.prompt }), } - const sink = createSubagentSink() - const link = linkAbort(parent.abortSignal) - let subagentRunId = suspended?.subagentRunId ?? '' - let text = suspended?.text ?? '' - let error: string | undefined - let result: unknown - try { - for await (const chunk of openAgentStream( - entry, - bag, + const child = { name: agent.name, ...fields } + if (call.background) { + if (call.sessionId !== undefined) { + throw new Error('background cannot continue a sessionId.') + } + const start = bag.binding?.start + if (!start) { + throw new Error( + 'background needs a host that runs children in the background.', + ) + } + // A background child counts against the tree budget too. + const refusal = budget.reserve(active) + if (refusal !== undefined) throw new Error(refusal) + const { subagentRunId } = await start( { - messages: parent.messagesFor?.(toolCallId) ?? parent.messages, - abortSignal: link.controller.signal, - threadId: parent.threadId, - parentRunId: parent.runId, - ...(parent.interruptedRunId !== undefined - ? { interruptedRunId: parent.interruptedRunId } - : {}), - ...(parent.parentSubagentRunId !== undefined - ? { parentSubagentRunId: parent.parentSubagentRunId } - : {}), + agent: agent.name, + ...fields, + ...(context?.toolCallId !== undefined && { + parentToolCallId: context.toolCallId, + }), }, - sink, - toolCallId, - )) { - if (chunk.type === SUBAGENT_STARTED && subagentRunId === '') { - subagentRunId = chunk.subagentRunId - } - if ( - chunk.type === EventType.TEXT_MESSAGE_CONTENT && - 'subagentRunId' in chunk && - chunk.subagentRunId === subagentRunId - ) { - text += chunk.delta - } - if ( - chunk.type === SUBAGENT_ERROR && - chunk.subagentRunId === subagentRunId - ) { - error = chunk.message - } - if ( - chunk.type === SUBAGENT_FINISHED && - chunk.subagentRunId === subagentRunId && - chunk.result !== undefined - ) { - result = chunk.result - } - toolContext?.[EMIT_STREAM_CHUNK]?.(chunk) - } - } finally { - link.dispose() - } - parent.sink.usage.push(...sink.usage) - if (sink.total) { - parent.sink.total = parent.sink.total - ? addTokenUsage(parent.sink.total, sink.total) - : sink.total + { wake: true }, + ) + return { + subagentRunId, + text: '', + result: { status: 'started' }, + keepRunId: true, + } satisfies SubagentToolOutcome } - return { - subagentRunId, - text, - ...(result !== undefined ? { result } : {}), - ...(error !== undefined ? { error } : {}), - ...(sink.interrupts.length > 0 ? { interrupts: sink.interrupts } : {}), - } satisfies SubagentToolOutcome + const { sessionId } = call + const work = + sessionId === undefined + ? runChild(child, context) + : storedChild(sessionId, agent.name).then((continued) => + runChild({ ...child, continued }, context), + ) + // The host can move the child to the background. It keeps running. + const detach = context?.detach + if (!detach) return work + const moved = detach(work).then( + (text) => ({ subagentRunId: '', text }) satisfies SubagentToolOutcome, + ) + return Promise.race([work, moved]) }, })) } diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 3ba9ab9bd..691eca150 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -62,7 +62,9 @@ import { transformMessagesForReplay } from '../../utilities/replay-messages' import { LazyToolManager } from './tools/lazy-tool-manager' import { assertUniqueToolNames } from './tools/unique-tool-names' import type { DefinedAgent } from './agents/define-agent' +import { providerKeys } from '../../byok/env-keys' import { + assertNoReservedFields, collectNamedText, createSubagentSink, createSyntheticSubagentTools, @@ -113,6 +115,7 @@ import { } from './messages' import { MiddlewareRunner } from './middleware/compose' import { getRunDetached } from './middleware/run-store' +import { LoadChildCapability } from './middleware/load-child' import { publishRunDetachedSignal } from '../../delivery-detach' import { publishRunDisconnectHandler } from '../../delivery-disconnect' import { provideSandboxRuntime } from './middleware/sandbox-runtime' @@ -5650,6 +5653,12 @@ export function chat< if (options.tools) { assertUniqueToolNames(options.tools) } + if (options.subagents) { + assertNoReservedFields( + options.subagents, + continuesChildren(readRuntimeMiddleware(options.middleware)), + ) + } const { outputSchema, stream } = options @@ -5902,11 +5911,20 @@ async function* streamTextChunks( engineRef: DeliveryEngineRef, ): AsyncIterable { // Children inherit the retention of this call. Their key is their own - // threadId. + // threadId. The keys of this call win over the keys of a parent agent, + // one provider at a time. const bag = options.subagents && { ...options.subagents, binding: { ...options.subagents.binding, + ...(options.subagents.keys + ? { + keys: providerKeys( + options.subagents.keys, + options.subagents.binding?.keys, + ), + } + : {}), promptCache: options.subagents.binding?.promptCache ?? options.promptCache.retention, }, @@ -5930,6 +5948,8 @@ async function* streamTextChunks( const synthetic = createSyntheticSubagentTools(bag, { messages: turn?.before ?? messages, messagesFor: calls.messagesFor, + childLoader: calls.childLoader, + canContinue: continuesChildren(options.middleware), threadId, runId, ...(options.parentRunId !== undefined && { @@ -6313,6 +6333,15 @@ function withPlan( return withTanstackMetadata(chunk, { [SUBAGENT_PLAN_KEY]: plan }) } +/** A middleware provides `loadChild`, so agent tools take `sessionId`. */ +function continuesChildren( + middleware: ReadonlyArray<{ provides?: ReadonlyArray }> | undefined, +) { + return (middleware ?? []).some((item) => + item.provides?.includes(LoadChildCapability), + ) +} + function readRoutedSubagentPersistence( middleware: | ReadonlyArray<{ diff --git a/packages/ai/src/activities/chat/middleware/load-child.ts b/packages/ai/src/activities/chat/middleware/load-child.ts new file mode 100644 index 000000000..bd806011b --- /dev/null +++ b/packages/ai/src/activities/chat/middleware/load-child.ts @@ -0,0 +1,23 @@ +import { createCapability } from './capabilities' +import type { ModelMessage, UIMessage } from '../../../types' + +/** + * Loads the stored transcript of a finished subagent run by its + * `subagentRunId`. Returns `undefined` for an unknown id. A persistence + * middleware provides it. The agent tools read it to continue a + * child by `sessionId`. When `agent` is set, only that agent can continue + * the child. A middleware that provides it must list `LoadChildCapability` + * in `provides`. Only then do the agent tools take `sessionId`. + */ +export type LoadChild = (subagentRunId: string) => Promise< + | { + messages: Array + /** The name of the agent the child ran under. */ + agent?: string + } + | undefined +> + +export const LoadChildCapability = createCapability()('load-child') + +export const [getLoadChild, provideLoadChild] = LoadChildCapability diff --git a/packages/ai/src/activities/chat/tools/tool-calls.ts b/packages/ai/src/activities/chat/tools/tool-calls.ts index 82983b4a4..2b889b293 100644 --- a/packages/ai/src/activities/chat/tools/tool-calls.ts +++ b/packages/ai/src/activities/chat/tools/tool-calls.ts @@ -70,6 +70,12 @@ export interface SubagentToolOutcome { error?: string /** Set when the child stopped for outside input. The tool call stays open. */ interrupts?: Array + /** + * The model gets `subagentRunId` in the result: it can continue the child + * (`sessionId`, with a persistence store) or a background child. Else the + * id stays in the subagent events only. + */ + keepRunId?: true } function isSubagentTool(tool: AnyTool): boolean { @@ -780,9 +786,13 @@ export async function* executeServerTool( subagentInterrupts?.push(...outcome.interrupts) return } + const runId = outcome.keepRunId + ? { subagentRunId: outcome.subagentRunId } + : {} const modelResult = outcome.error - ? { error: outcome.error } + ? { ...runId, error: outcome.error } : { + ...runId, // A string result is also the child's text: the model gets all of it. result: outcome.result === undefined || typeof outcome.result === 'string' diff --git a/packages/ai/src/activities/index.ts b/packages/ai/src/activities/index.ts index 5e9322e87..cb8568b91 100644 --- a/packages/ai/src/activities/index.ts +++ b/packages/ai/src/activities/index.ts @@ -53,6 +53,7 @@ export type { SubagentBinding, SubagentForward, } from './chat/agents/bound' +export { SubagentBudget, type SubagentLimits } from './chat/agents/limits' export { subagentRoute, type SubagentRouteOptions } from './chat/agents/route' export type { SubagentOrder, diff --git a/packages/ai/src/byok.ts b/packages/ai/src/byok.ts index 550a12891..010672f55 100644 --- a/packages/ai/src/byok.ts +++ b/packages/ai/src/byok.ts @@ -16,4 +16,8 @@ export { } from './byok/errors' export { maskKey, scrubSecrets } from './byok/scrub' export { keyedAdapter, keyedAdapters, isKeyedAdapter } from './byok/keyed' -export type { KeyedAdapter, KeyedAdapterResult } from './byok/keyed' +export type { + KeyedAdapter, + KeyedAdapterResult, + ProviderKeys, +} from './byok/keyed' diff --git a/packages/ai/src/byok/env-keys.ts b/packages/ai/src/byok/env-keys.ts new file mode 100644 index 000000000..ef068002f --- /dev/null +++ b/packages/ai/src/byok/env-keys.ts @@ -0,0 +1,66 @@ +import { ByokMissingError } from './errors' +import { envKey } from './get-key' +import { isKeyedAdapter } from './keyed' +import { resolveProviderId } from './providers' +import type { ByokProvider } from './define-provider' +import type { ProviderKeys } from './keyed' +import type { ProviderId } from './providers' + +function missingKeyMessage(provider: ByokProvider | ProviderId) { + if (typeof provider === 'string') return `Missing ${provider} API key.` + const names = provider.env ?? [] + const hint = names.length > 0 ? ` Set ${names.join(' or ')}.` : '' + return `Missing ${provider.label} API key.${hint}` +} + +/** + * Where `subagents.keys` finds a key: an object by provider id, or a function + * that gets the provider (a BYOK descriptor or an id) and returns its key. + */ +export type ProviderKeySource = + | Readonly> + | (( + provider: ByokProvider | ProviderId, + ) => string | null | undefined | Promise) + +/** + * The `ctx.keys` of an agent: `keys`, else `fallback` (the keys of a parent + * agent), else each provider's `env` names. Not in `@tanstack/ai/byok`, so + * the client entry has no env access. + */ +export function providerKeys( + keys: ProviderKeySource = {}, + fallback?: ProviderKeys, +): ProviderKeys { + const get = async (provider: ByokProvider | ProviderId) => { + const own = + typeof keys === 'function' + ? await keys(provider) + : keys[resolveProviderId(provider)] + return ( + own || + (fallback ? await fallback.get(provider) : envKey(provider)) || + null + ) + } + const require = async (provider: ByokProvider | ProviderId) => { + const key = await get(provider) + if (!key) { + const error = new ByokMissingError(resolveProviderId(provider)) + error.message = missingKeyMessage(provider) + throw error + } + return key + } + return { + get, + require, + adapter: async (adapter) => + isKeyedAdapter(adapter) + ? adapter.create(await require(adapter.provider)) + : adapter, + } +} + +/** The `ctx.keys` of an agent with no keys set: each provider's `env` names. */ +export const envProviderKeys = providerKeys() diff --git a/packages/ai/src/byok/get-key.ts b/packages/ai/src/byok/get-key.ts index dec1a5b66..8b2318c16 100644 --- a/packages/ai/src/byok/get-key.ts +++ b/packages/ai/src/byok/get-key.ts @@ -3,6 +3,21 @@ import type { KeyedAdapter, KeyedAdapterResult } from './keyed' import type { ByokProvider } from './define-provider' import type { ProviderId } from './providers' +/** + * Read `provider.env` in order and return the first value that is set. A + * slug has no env names. Where `process` is missing (a browser) it returns + * `null`. + */ +export function envKey(provider: ProviderId | ByokProvider) { + if (typeof provider === 'string') return null + const env = globalThis.process?.env + for (const name of provider.env ?? []) { + const value = env?.[name] + if (typeof value === 'string' && value.length > 0) return value + } + return null +} + /** * Read a key on the relay. Import from `@tanstack/ai/byok/server` so this * `process.env` access is not in the client graph. @@ -19,12 +34,7 @@ export function getByokKey( const trimmed = value.trim() if (trimmed.length > 0) return trimmed } - if (typeof provider === 'string') return null - for (const name of provider.env ?? []) { - const envValue = process.env[name] - if (typeof envValue === 'string' && envValue.length > 0) return envValue - } - return null + return envKey(provider) } /** diff --git a/packages/ai/src/byok/keyed.ts b/packages/ai/src/byok/keyed.ts index 51ec7c099..f0c4b5d42 100644 --- a/packages/ai/src/byok/keyed.ts +++ b/packages/ai/src/byok/keyed.ts @@ -17,6 +17,25 @@ export interface KeyedAdapter { readonly create: (key: string) => TAdapter } +/** + * Finds provider keys for a run. An agent reads it as `ctx.keys`. A host sets + * it on the binding and reads each user's own keys. Without a host, + * `ctx.keys` reads each provider's `env` names. + */ +export interface ProviderKeys { + /** The key for `provider`, or `null` when there is none. */ + get: (provider: ByokProvider | ProviderId) => Promise + /** The key for `provider`. Throws when there is none. */ + require: (provider: ByokProvider | ProviderId) => Promise + /** + * The adapter to call. A plain adapter comes back unchanged. A + * {@link KeyedAdapter} is built with the key from `require`. + */ + adapter: ( + adapter: TAdapter | KeyedAdapter, + ) => Promise +} + /** * Wrap an adapter factory that needs a provider key. The host finds the key * and calls `create` just before each call, so the key is not in your code diff --git a/packages/ai/src/client.ts b/packages/ai/src/client.ts index 19ba1ecb9..e3bcbf4c0 100644 --- a/packages/ai/src/client.ts +++ b/packages/ai/src/client.ts @@ -323,6 +323,7 @@ export type { } from './activities/chat/stream/index' export { uiMessagesToWire } from './utilities/ag-ui-wire' +export { subagentInputOf } from './utilities/subagent-wire' export { mergeMetadata, tanstackMetadata, diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 8a6514f44..45985eb4f 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -99,13 +99,19 @@ export type { SubagentBinding, SubagentForward, } from './activities/chat/agents/bound' +export { + SubagentBudget, + type SubagentLimits, +} from './activities/chat/agents/limits' export { keyedAdapter, keyedAdapters, isKeyedAdapter, type KeyedAdapter, type KeyedAdapterResult, + type ProviderKeys, } from './byok/keyed' +export type { ProviderKeySource } from './byok/env-keys' export { subagentRoute, type SubagentRouteOptions, @@ -377,6 +383,13 @@ export type { AnyChatMiddleware, MetadataStore, } from './activities/chat/middleware/index' +// A persistence store gives the agent tools the stored children. +export { + LoadChildCapability, + getLoadChild, + provideLoadChild, +} from './activities/chat/middleware/load-child' +export type { LoadChild } from './activities/chat/middleware/load-child' // Locks are a distributed-mutex primitive — coordination, not chat state — and // live behind their own subpath: `@tanstack/ai/locks` (see ./locks.ts). @@ -569,6 +582,7 @@ export { uiMessagesToWire } from './utilities/ag-ui-wire' export type { SubagentWireInfo } from './utilities/subagent-wire' export { subagentHostMessageId, + subagentInputOf, wireSubagentInfo, } from './utilities/subagent-wire' export { mergeMetadata, withTanstackMetadata } from './utilities/merge-metadata' diff --git a/packages/ai/src/utilities/subagent-wire.ts b/packages/ai/src/utilities/subagent-wire.ts index db2a79e04..1df463d5e 100644 --- a/packages/ai/src/utilities/subagent-wire.ts +++ b/packages/ai/src/utilities/subagent-wire.ts @@ -34,6 +34,18 @@ function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null && !Array.isArray(value) } +/** Key of a child's checked input in `SUBAGENT_STARTED` metadata (`metadata.tanstack`). */ +export const SUBAGENT_INPUT_KEY = 'subagentInput' + +/** + * The checked input of a child that an agent with `inputSchema` got, read + * from the `metadata` of its card or of its `SUBAGENT_STARTED` event. + */ +export function subagentInputOf(metadata: unknown): unknown { + if (!isRecord(metadata) || !isRecord(metadata.tanstack)) return undefined + return metadata.tanstack[SUBAGENT_INPUT_KEY] +} + export function wireSubagentRunId(message: unknown): string | undefined { if (!isRecord(message) || 'parts' in message) return undefined const id = message.subagentRunId diff --git a/packages/ai/tests/define-agent.test.ts b/packages/ai/tests/define-agent.test.ts index ba676a054..b8e2a87d7 100644 --- a/packages/ai/tests/define-agent.test.ts +++ b/packages/ai/tests/define-agent.test.ts @@ -870,8 +870,10 @@ describe('defineAgent inputSchema', () => { ) const tool = calls[0]?.tools?.find((entry) => entry.name === 'researcher') - // The engine gives a tool with no schema an empty object schema. - expect(tool?.inputSchema?.properties ?? {}).toEqual({}) + // An agent without inputSchema takes only a prompt. + expect(tool?.inputSchema?.properties).toEqual({ + prompt: { type: 'string' }, + }) expect(inputs).toEqual([undefined]) }) diff --git a/packages/ai/tests/provider-keys.test.ts b/packages/ai/tests/provider-keys.test.ts index 2bf0285b2..4ec9eac8b 100644 --- a/packages/ai/tests/provider-keys.test.ts +++ b/packages/ai/tests/provider-keys.test.ts @@ -7,7 +7,15 @@ import { keyedAdapters, } from '../src/byok' import { getByokKey, keyedAdapterFromRequest } from '../src/byok/server' -import { createMockAdapter } from './test-utils' +import { defineAgent } from '../src/activities/chat/agents/define-agent' +import { spawnAgentStream } from '../src/activities/chat/agents/spawn' +import { chat } from '../src/activities/chat' +import { collectChunks, createMockAdapter } from './test-utils' +import type { AnyTextAdapter } from '../src/activities/chat/adapter' +import type { SubagentBinding } from '../src/activities/chat/agents/bound' +import type { DefinedAgent } from '../src/activities/chat/agents/define-agent' +import type { KeyedAdapter, ProviderKeys } from '../src/byok' +import type { ProviderKeySource } from '../src/byok/env-keys' const acme = defineByokProvider({ id: 'acme', label: 'Acme' }) @@ -102,3 +110,262 @@ describe('keyedAdapters', () => { expect(built).toEqual({ kind: 'image', key: 'sk-user' }) }) }) + +const envName = 'TANSTACK_AI_PROVIDER_KEYS_TEST' +const acmeEnv = defineByokProvider({ id: 'acme', label: 'Acme', env: envName }) + +const input = { + input: undefined, + messages: [], + threadId: 'thread-1', + runId: 'run-1', + parentRunId: 'parent-1', + subagentRunId: 'child-1', +} + +/** Provider keys a host would pass: every provider gets `key`. */ +function hostKeys(key: string): ProviderKeys { + return { + get: async () => key, + require: async () => key, + adapter: async (adapter) => + isKeyedAdapter(adapter) ? adapter.create(key) : adapter, + } +} + +/** An agent that records the `ctx.keys` its run gets. */ +function probe() { + const seen: Array = [] + const agent = defineAgent({ + name: 'probe', + description: 'Records ctx.keys', + run: (ctx) => { + seen.push(ctx.keys) + return Promise.resolve('ok') + }, + }) + return { agent, seen } +} + +/** Run `agent` the way a host does, outside a parent chat turn. */ +function run(agent: DefinedAgent, binding?: SubagentBinding) { + return collectChunks( + spawnAgentStream(agent, input, undefined, undefined, binding), + ) +} + +/** The `ctx.keys` of one probe run. */ +async function runKeys(binding?: SubagentBinding) { + const { agent, seen } = probe() + await run(agent, binding) + const keys = seen[0] + if (!keys) throw new Error('the probe did not run') + return keys +} + +/** + * The `ctx.keys` of a probe that a parent agent starts through its own + * `ctx.chat({ subagents })`. `own` is the binding that nested chat sets. + */ +async function nestedKeys(host: SubagentBinding, own?: SubagentBinding) { + const { agent: child, seen } = probe() + const { adapter } = createMockAdapter({ iterations: [] }) + const lead = defineAgent({ + name: 'lead', + description: 'Starts the probe', + run: (ctx) => + ctx.chat({ + adapter, + subagents: { + agents: [child], + router: () => 'probe', + ...(own ? { binding: own } : {}), + }, + }), + }) + await run(lead, host) + return seen[0] +} + +describe('ctx.keys', () => { + const acmeEnvImage = keyedAdapter(acmeEnv, (key) => ({ + kind: 'image' as const, + key, + })) + + afterEach(() => vi.unstubAllEnvs()) + + it('is the provider keys of the binding', async () => { + const keys = hostKeys('sk-host') + expect(await runKeys({ keys })).toBe(keys) + }) + + it('reads the provider env var when the host sets no keys', async () => { + vi.stubEnv(envName, 'sk-env') + const keys = await runKeys() + expect(await keys.get(acmeEnv)).toBe('sk-env') + expect(await keys.require(acmeEnv)).toBe('sk-env') + // A bare provider id has no env names. + expect(await keys.get('acme')).toBeNull() + }) + + it('throws an error that names the env var when the key is missing', async () => { + const keys = await runKeys() + expect(await keys.get(acmeEnv)).toBeNull() + await expect(keys.require(acmeEnv)).rejects.toThrow( + `Missing Acme API key. Set ${envName}.`, + ) + await expect(keys.require('acme')).rejects.toThrow('Missing acme API key.') + }) + + it('returns no key where process is missing', async () => { + const keys = await runKeys() + vi.stubGlobal('process', undefined) + // `get` reads the env before its first await, so the stub is gone + // before vitest needs `process` again. + const pending = keys.get(acmeEnv) + vi.unstubAllGlobals() + expect(await pending).toBeNull() + }) + + it('builds a keyed adapter with the key and returns a plain adapter unchanged', async () => { + vi.stubEnv(envName, 'sk-env') + const keys = await runKeys() + const { adapter: text } = createMockAdapter({}) + const choice: AnyTextAdapter | KeyedAdapter = keyedAdapter( + acmeEnv, + () => text, + ) + + expect(await keys.adapter(acmeEnvImage)).toEqual({ + kind: 'image', + key: 'sk-env', + }) + expect(await keys.adapter(text)).toBe(text) + expect(await keys.adapter(choice)).toBe(text) + expectTypeOf(keys.adapter(acmeEnvImage)).resolves.toEqualTypeOf<{ + kind: 'image' + key: string + }>() + expectTypeOf(keys.adapter(choice)).resolves.toEqualTypeOf() + }) + + it('gives a nested child the keys of the parent binding', async () => { + const keys = hostKeys('sk-host') + expect(await nestedKeys({ keys })).toBe(keys) + }) + + it('lets a nested chat pass its own keys', async () => { + const own = hostKeys('sk-own') + expect(await nestedKeys({ keys: hostKeys('sk-host') }, { keys: own })).toBe( + own, + ) + }) +}) + +/** + * The `ctx.keys` of a probe that `chat({ subagents: { keys } })` starts. + * With `nested`, a lead agent starts the probe through its own + * `ctx.chat({ subagents: { ...nested } })`. + */ +async function chatKeys( + keys: ProviderKeySource, + nested?: { keys?: ProviderKeySource }, +) { + const { agent: target, seen } = probe() + const { adapter } = createMockAdapter({ iterations: [] }) + const lead = defineAgent({ + name: 'lead', + description: 'Starts the probe', + run: (ctx) => + ctx.chat({ + adapter, + subagents: { agents: [target], router: () => 'probe', ...nested }, + }), + }) + const agents: Array = nested ? [lead] : [target] + const name = nested ? 'lead' : 'probe' + await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + subagents: { agents, router: () => name, keys }, + }), + ) + const found = seen[0] + if (!found) throw new Error('the probe did not run') + return found +} + +describe('subagents.keys', () => { + afterEach(() => vi.unstubAllEnvs()) + + it('gives agents the key of each provider id', async () => { + vi.stubEnv(envName, 'sk-env') + const keys = await chatKeys({ acme: 'sk-user' }) + const image = keyedAdapter(acmeEnv, (key) => ({ kind: 'image', key })) + + expect(await keys.require(acmeEnv)).toBe('sk-user') + expect(await keys.get('acme')).toBe('sk-user') + expect(await keys.adapter(image)).toEqual({ kind: 'image', key: 'sk-user' }) + }) + + it('falls back to the env var for a missing or null key', async () => { + vi.stubEnv(envName, 'sk-env') + expect(await (await chatKeys({ acme: null })).require(acmeEnv)).toBe( + 'sk-env', + ) + expect(await (await chatKeys({ other: 'sk-other' })).get(acmeEnv)).toBe( + 'sk-env', + ) + }) + + it('throws an error that names the env var when no key is found', async () => { + const keys = await chatKeys({ acme: '' }) + await expect(keys.require(acmeEnv)).rejects.toThrow( + `Missing Acme API key. Set ${envName}.`, + ) + }) + + it('gives a nested child the keys of the chat', async () => { + const keys = await chatKeys({ acme: 'sk-user' }, {}) + expect(await keys.require(acmeEnv)).toBe('sk-user') + }) + + it('lets a nested chat set its own keys', async () => { + const keys = await chatKeys( + { acme: 'sk-user' }, + { keys: { acme: 'sk-own' } }, + ) + expect(await keys.require(acmeEnv)).toBe('sk-own') + }) + + it('reads keys from a function, with the provider it gets', async () => { + vi.stubEnv(envName, 'sk-env') + const asked: Array = [] + const keys = await chatKeys(async (provider) => { + asked.push(provider) + return provider === acmeEnv ? 'sk-fn' : null + }) + + expect(await keys.require(acmeEnv)).toBe('sk-fn') + expect(asked).toEqual([acmeEnv]) + // The function has no key for a bare id, so the env name is not read. + expect(await keys.get('acme')).toBeNull() + }) + + it('falls back to the parent keys when a nested function has no key', async () => { + const keys = await chatKeys({ acme: 'sk-user' }, { keys: () => null }) + expect(await keys.require(acmeEnv)).toBe('sk-user') + }) + + it('keeps the parent keys for providers that a nested chat does not set', async () => { + vi.stubEnv(envName, 'sk-env') + const keys = await chatKeys( + { acme: 'sk-user' }, + { keys: { other: 'sk-other' } }, + ) + expect(await keys.require(acmeEnv)).toBe('sk-user') + expect(await keys.get('other')).toBe('sk-other') + }) +}) diff --git a/packages/ai/tests/subagent-limits.test.ts b/packages/ai/tests/subagent-limits.test.ts new file mode 100644 index 000000000..9e74ad74c --- /dev/null +++ b/packages/ai/tests/subagent-limits.test.ts @@ -0,0 +1,186 @@ +import { describe, expect, it, vi } from 'vitest' +import { defineAgent } from '../src/activities/chat/agents/define-agent' +import { SubagentBudget } from '../src/activities/chat/agents/limits' +import { chat } from '../src/activities/chat' +import { collectChunks, createMockAdapter, ev } from './test-utils' +import type { StreamChunk } from '../src/types' + +/** A model that calls `name` once for each id, all in one turn, then stops. */ +function callsInOneTurn(name: string, ids: Array) { + return createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ...ids.flatMap((id) => [ev.toolStart(id, name), ev.toolArgs(id, '{}')]), + ev.runFinished('tool_calls'), + ], + [ev.runStarted(), ev.runFinished('stop')], + ], + }) +} + +function results(chunks: Array): Record { + const out: Record = {} + for (const chunk of chunks) { + if ( + chunk.type === 'TOOL_CALL_RESULT' && + typeof chunk.content === 'string' + ) { + out[chunk.toolCallId] = JSON.parse(chunk.content) + } + } + return out +} + +describe('SubagentBudget', () => { + it('counts calls across the tree and checks depth', () => { + const root = SubagentBudget.root({ maxCalls: 2, maxDepth: 1 }) + expect(root.reserve(0)).toBeUndefined() + const child = root.child() + expect(child.reserve(0)).toBe('subagent limit reached (maxDepth 1)') + expect(root.reserve(0)).toBeUndefined() + expect(root.reserve(0)).toBe('subagent limit reached (maxCalls 2)') + expect(child.calls).toBe(2) + }) + + it('never gives a child more time than its parent has left', () => { + const root = SubagentBudget.root({ timeoutMs: 10_000 }) + expect(root.childTimeout(0)).toBe(10_000) + const child = root.child(5_000) + expect(child.childTimeout(1_000)).toBe(4_000) + }) +}) + +describe('subagent limits', () => { + it('refuses children past maxCalls with a tool error', async () => { + const run = vi.fn(async () => 'done') + const worker = defineAgent({ name: 'worker', description: 'Works', run }) + const { adapter } = callsInOneTurn('worker', ['c1', 'c2', 'c3']) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + subagents: { agents: [worker], limits: { maxCalls: 2 } }, + }) as AsyncIterable, + ) + const byId = results(chunks) + const refused = Object.values(byId).filter( + (result) => result.error !== undefined, + ) + expect(refused).toHaveLength(1) + expect(refused[0].error).toBe('subagent limit reached (maxCalls 2)') + expect(run).toHaveBeenCalledTimes(2) + }) + + it('passes the budget to nested children through ctx.chat', async () => { + const grandchildRun = vi.fn(async () => 'deep') + const grandchild = defineAgent({ + name: 'grandchild', + description: 'Deep', + run: grandchildRun, + }) + const child = defineAgent({ + name: 'child', + description: 'Calls a grandchild', + run: (ctx) => + ctx.chat({ + adapter: callsInOneTurn('grandchild', ['g1']).adapter, + subagents: { agents: [grandchild] }, + }), + }) + const { adapter } = callsInOneTurn('child', ['c1']) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + subagents: { agents: [child], limits: { maxDepth: 1 } }, + }) as AsyncIterable, + ) + expect(grandchildRun).not.toHaveBeenCalled() + const grandchildResult = chunks.find( + (chunk) => chunk.type === 'TOOL_CALL_RESULT' && chunk.toolCallId === 'g1', + ) + expect(JSON.stringify(grandchildResult)).toContain('maxDepth 1') + }) + + it('stops a child that runs past timeoutMs', async () => { + const slow = defineAgent({ + name: 'slow', + description: 'Never finishes on its own', + run: (ctx) => + new Promise((resolve) => { + ctx.abortSignal?.addEventListener('abort', () => resolve('stopped')) + }), + }) + const { adapter } = callsInOneTurn('slow', ['c1']) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + subagents: { agents: [slow], limits: { timeoutMs: 20 } }, + }) as AsyncIterable, + ) + expect(results(chunks).c1.error).toBe('subagent timed out after 20 ms') + }) + + it('refuses children past maxConcurrent in one turn', async () => { + let release!: () => void + const gate = new Promise((resolve) => { + release = resolve + }) + const waiting = defineAgent({ + name: 'waiting', + description: 'Waits', + run: async () => { + await gate + return 'ok' + }, + }) + setTimeout(() => release(), 20) + const { adapter } = callsInOneTurn('waiting', ['c1', 'c2', 'c3']) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + subagents: { agents: [waiting], limits: { maxConcurrent: 2 } }, + toolExecution: 'parallel', + }) as AsyncIterable, + ) + const errors = Object.values(results(chunks)) + .map((result) => result.error) + .filter(Boolean) + // The three calls start together. The third finds two running. + expect(errors).toEqual(['subagent limit reached (maxConcurrent 2)']) + }) + + it('counts the children of a routed child against the root limits', async () => { + const grandchildRun = vi.fn(async () => 'deep') + const grandchild = defineAgent({ + name: 'grandchild', + description: 'Deep', + run: grandchildRun, + }) + const child = defineAgent({ + name: 'child', + description: 'Calls a grandchild', + run: (ctx) => + ctx.chat({ + adapter: callsInOneTurn('grandchild', ['g1']).adapter, + subagents: { agents: [grandchild] }, + }), + }) + const chunks = await collectChunks( + chat({ + adapter: createMockAdapter({ iterations: [] }).adapter, + messages: [{ role: 'user', content: 'Go' }], + subagents: { + agents: [child], + router: () => 'child', + limits: { maxDepth: 1 }, + }, + }) as AsyncIterable, + ) + expect(grandchildRun).not.toHaveBeenCalled() + expect(JSON.stringify(chunks)).toContain('maxDepth 1') + }) +}) diff --git a/packages/ai/tests/subagent-tool-input.test.ts b/packages/ai/tests/subagent-tool-input.test.ts new file mode 100644 index 000000000..ac2f30bac --- /dev/null +++ b/packages/ai/tests/subagent-tool-input.test.ts @@ -0,0 +1,512 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { defineAgent } from '../src/activities/chat/agents/define-agent' +import { chat } from '../src/activities/chat' +import { + LoadChildCapability, + provideLoadChild, +} from '../src/activities/chat/middleware/load-child' +import { + createSubagentSink, + createSyntheticSubagentTools, +} from '../src/activities/chat/agents/spawn' +import { EventType } from '../src/types' +import { collectChunks, createMockAdapter, ev } from './test-utils' +import type { SubagentTurn } from '../src/activities/chat/agents/turn' +import type { SubagentBinding } from '../src/activities/chat/agents/bound' +import type { ChatMiddleware } from '../src/activities/chat/middleware/types' +import type { ModelMessage, StreamChunk, UIMessage } from '../src/types' + +/** What one child run saw. */ +interface ChildRun { + agent: string + input: unknown + messages: Array +} + +const runs: Array = [] + +const researcher = defineAgent({ + name: 'researcher', + description: 'Looks up facts', + inputSchema: z.object({ topic: z.string() }), + run: async (ctx) => { + runs.push({ agent: 'researcher', input: ctx.input, messages: ctx.messages }) + return `facts about ${ctx.input.topic}` + }, +}) + +const writer = defineAgent({ + name: 'writer', + description: 'Writes text', + run: async (ctx) => { + runs.push({ agent: 'writer', input: ctx.input, messages: ctx.messages }) + return 'draft' + }, +}) + +const agents = [researcher, writer] + +beforeEach(() => { + runs.length = 0 +}) + +/** + * A parent model that writes a line of text, calls the tool of `agent` once + * with `args`, then stops. + */ +function parentCalling(agent: string, args: object) { + return createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.textStart(), + ev.textContent('Asking a subagent.'), + ev.textEnd(), + ev.toolStart('call_1', agent), + ev.toolArgs('call_1', JSON.stringify(args)), + ev.runFinished('tool_calls'), + ], + [ev.runStarted(), ev.runFinished('stop')], + ], + }) +} + +/** The parsed result of tool call `call_1`. */ +function toolResult(chunks: Array) { + for (const chunk of chunks) { + if ( + chunk.type === EventType.TOOL_CALL_RESULT && + chunk.toolCallId === 'call_1' && + typeof chunk.content === 'string' + ) { + return JSON.parse(chunk.content) + } + } + throw new Error('No result for call_1') +} + +/** Run a parent chat. The model calls the tool of `agent` with `args`. */ +async function callSubagent( + agent: string, + args: object, + options: { + middleware?: Array + binding?: SubagentBinding + runId?: string + } = {}, +) { + const { adapter, calls } = parentCalling(agent, args) + const chunks = await collectChunks( + chat({ + adapter, + ...(options.runId !== undefined && { runId: options.runId }), + messages: [{ role: 'user', content: 'Go' }], + subagents: { + agents, + ...(options.binding && { binding: options.binding }), + }, + ...(options.middleware && { middleware: options.middleware }), + }) as AsyncIterable, + ) + return { chunks, calls, result: toolResult(chunks) } +} + +/** + * A fake persistence middleware that provides the `loadChild` service. Each + * stored child ran under the `writer` agent. + */ +function storedChildren( + stored: Map>, +): ChatMiddleware { + return { + name: 'stored-children', + provides: [LoadChildCapability], + setup(ctx) { + provideLoadChild(ctx, async (subagentRunId) => { + const messages = stored.get(subagentRunId) + return messages === undefined + ? undefined + : { messages, agent: 'writer' } + }) + }, + } +} + +const oldChild = new Map>([ + [ + 'subagent-old', + [ + { role: 'user', content: 'Draft a tide report' }, + { role: 'assistant', content: 'Tides rise twice a day.' }, + ], + ], +]) + +/** The wire schema of each agent tool, by name. */ +function toolSchemas( + calls: Array<{ tools?: Array<{ name: string; inputSchema?: unknown }> }>, +) { + return Object.fromEntries( + (calls[0]?.tools ?? []).map((tool) => [tool.name, tool.inputSchema]), + ) +} + +describe('agent tool input', () => { + it('takes the agent schema fields, and gives an agent without one a prompt', async () => { + const { calls, chunks, result } = await callSubagent('researcher', { + topic: 'tides', + }) + + const schemas = toolSchemas(calls) + expect(Object.keys(schemas)).toEqual(['researcher', 'writer']) + expect(schemas.researcher).toMatchObject({ + type: 'object', + properties: { topic: { type: 'string' } }, + required: ['topic'], + }) + expect(schemas.writer).toEqual({ + type: 'object', + properties: { prompt: { type: 'string' } }, + }) + expect(runs).toMatchObject([ + { agent: 'researcher', input: { topic: 'tides' } }, + ]) + expect(result.result).toBe('facts about tides') + // The card keeps the checked input for the UI. + expect( + chunks.find((chunk) => chunk.type === EventType.SUBAGENT_STARTED), + ).toMatchObject({ + metadata: { tanstack: { subagentInput: { topic: 'tides' } } }, + }) + }) + + it('takes sessionId next to the fields of a strict schema', async () => { + const strict = defineAgent({ + name: 'strict', + description: 'Takes a strict input', + inputSchema: z.object({ topic: z.string() }).strict(), + run: async (ctx) => { + runs.push({ agent: 'strict', input: ctx.input, messages: ctx.messages }) + return 'ok' + }, + }) + const { adapter } = parentCalling('strict', { + topic: 'tides', + sessionId: 'subagent-missing', + }) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + middleware: [storedChildren(oldChild)], + subagents: { agents: [strict] }, + }) as AsyncIterable, + ) + + // The schema check passes. The unknown sessionId is the only error. + expect(toolResult(chunks)).toEqual({ + error: 'Unknown sessionId "subagent-missing".', + }) + }) + + const taken = defineAgent({ + name: 'taken', + description: 'Uses sessionId', + inputSchema: z.object({ sessionId: z.string() }), + run: async (ctx) => { + runs.push({ agent: 'taken', input: ctx.input, messages: ctx.messages }) + return 'ok' + }, + }) + + it('checks the agent fields before middleware, with sessionId in the args', async () => { + const deep = defineAgent({ + name: 'deep', + description: 'Takes a depth', + inputSchema: z.object({ + topic: z.string(), + depth: z.number().default(2), + }), + run: async (ctx) => { + runs.push({ agent: 'deep', input: ctx.input, messages: ctx.messages }) + return 'ok' + }, + }) + const seen: Array = [] + const spy: ChatMiddleware = { + name: 'spy', + onBeforeToolCall: (_ctx, hook) => { + seen.push(hook.args) + }, + } + const { adapter } = parentCalling('deep', { topic: 'tides' }) + await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + middleware: [storedChildren(oldChild), spy], + subagents: { agents: [deep] }, + }) as AsyncIterable, + ) + + // The schema default is in place when middleware runs. + expect(seen).toEqual([{ topic: 'tides', depth: 2 }]) + expect(runs[0]?.input).toEqual({ topic: 'tides', depth: 2 }) + }) + + it('throws when an inputSchema uses a reserved field that the run uses', () => { + const { adapter } = parentCalling('taken', {}) + + expect(() => + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + middleware: [storedChildren(oldChild)], + subagents: { agents: [taken] }, + }), + ).toThrow( + 'Agent "taken" uses the reserved input field "sessionId". Rename that field.', + ) + }) + + it('keeps a reserved name as an agent field when the run does not use it', async () => { + const { adapter } = parentCalling('taken', { sessionId: 'mine' }) + await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + subagents: { agents: [taken] }, + }) as AsyncIterable, + ) + + expect(runs[0]?.input).toEqual({ sessionId: 'mine' }) + }) + + it('has sessionId and background only when the run can use them', async () => { + const plain = await callSubagent('writer', {}) + const able = await callSubagent( + 'writer', + {}, + { + middleware: [storedChildren(oldChild)], + binding: { start: async () => ({ subagentRunId: 'unused' }) }, + }, + ) + + expect(toolSchemas(plain.calls).writer).not.toHaveProperty( + 'properties.sessionId', + ) + expect(toolSchemas(plain.calls).writer).not.toHaveProperty( + 'properties.background', + ) + expect(toolSchemas(able.calls).writer).toMatchObject({ + properties: { + sessionId: { type: 'string' }, + background: { type: 'boolean' }, + }, + }) + }) + + it('reads the conversation when an agent without inputSchema gets no prompt', async () => { + const { result } = await callSubagent('writer', {}) + + expect(runs[0]?.input).toBeUndefined() + expect(runs[0]?.messages).toMatchObject([ + { role: 'user', content: 'Go' }, + { role: 'assistant', content: 'Asking a subagent.' }, + ]) + // No persistence store, so the model gets no subagentRunId. + expect(result).toEqual({ result: 'draft' }) + }) + + it('adds the prompt as the last user message of a child without inputSchema', async () => { + await callSubagent('writer', { prompt: 'Write a haiku' }) + + expect(runs).toHaveLength(1) + expect(runs[0]?.input).toBeUndefined() + expect(runs[0]?.messages).toMatchObject([ + { role: 'user', content: 'Go' }, + { role: 'assistant', content: 'Asking a subagent.' }, + { role: 'user', content: 'Write a haiku' }, + ]) + }) + + it('continues a stored child by sessionId', async () => { + const { result } = await callSubagent( + 'writer', + { sessionId: 'subagent-old', prompt: 'Make it shorter' }, + { middleware: [storedChildren(oldChild)] }, + ) + + expect(runs[0]?.messages).toEqual([ + { role: 'user', content: 'Draft a tide report' }, + { role: 'assistant', content: 'Tides rise twice a day.' }, + { role: 'user', content: 'Make it shorter' }, + ]) + expect(result).toEqual({ subagentRunId: 'subagent-old', result: 'draft' }) + }) + + it('gives the text of each continue its own message', async () => { + const ids: Array = [] + for (const runId of ['run-1', 'run-2']) { + const { chunks } = await callSubagent( + 'writer', + { sessionId: 'subagent-old' }, + { middleware: [storedChildren(oldChild)], runId }, + ) + for (const chunk of chunks) { + const isChildText = + chunk.type === EventType.TEXT_MESSAGE_START && + 'subagentRunId' in chunk && + chunk.subagentRunId === 'subagent-old' + if (isChildText) ids.push(chunk.messageId) + } + } + + expect(ids).toHaveLength(2) + expect(ids[0]).not.toBe(ids[1]) + }) + + it('gives a resumed continued child each stored message once', async () => { + const stored: Array = [ + { id: 'old-user', role: 'user', content: 'Draft a tide report' }, + { id: 'old-reply', role: 'assistant', content: 'Tides rise twice.' }, + ] + const waiting: ModelMessage = { + id: 'new-reply', + role: 'assistant', + content: 'Shorter like this?', + } + const turn: SubagentTurn = { + before: [], + children: [ + { + subagentRunId: 'subagent-old', + name: 'writer', + status: 'suspended', + parentToolCallId: 'call_1', + // The card holds the whole child: the stored messages, then the + // call that waits. + messages: [...stored, waiting], + text: '', + resume: [ + { interruptId: 'approval_1', status: 'resolved', payload: true }, + ], + }, + ], + rest: [], + } + const [, tool] = createSyntheticSubagentTools( + { agents }, + { + messages: [], + threadId: 'thread-1', + runId: 'run-2', + interruptedRunId: 'run-1', + turn, + sink: createSubagentSink(), + childLoader: () => async () => ({ messages: stored }), + canContinue: true, + }, + ) + await tool?.execute?.( + { sessionId: 'subagent-old', prompt: 'Make it shorter' }, + { toolCallId: 'call_1', emitCustomEvent: () => {} }, + ) + + expect(runs[0]?.messages).toEqual([ + ...stored, + { role: 'user', content: 'Make it shorter' }, + waiting, + ]) + }) + + it('starts a background child through the start hook', async () => { + const start = vi.fn(async () => ({ subagentRunId: 'subagent-bg' })) + const { result } = await callSubagent( + 'writer', + { prompt: 'Draft', background: true }, + { binding: { start } }, + ) + + expect(result).toEqual({ + subagentRunId: 'subagent-bg', + result: { status: 'started' }, + }) + expect(start).toHaveBeenCalledWith( + { agent: 'writer', prompt: 'Draft', parentToolCallId: 'call_1' }, + { wake: true }, + ) + expect(runs).toEqual([]) + }) + + const host: SubagentBinding = { + start: async () => ({ subagentRunId: 'unused' }), + } + const toolErrors: Array< + [ + string, + string, + object, + { middleware?: Array; binding?: SubagentBinding }, + string, + ] + > = [ + [ + 'a missing field of the agent schema', + 'researcher', + {}, + {}, + 'Input validation failed for tool researcher: Validation failed: Invalid input: expected string, received undefined', + ], + [ + 'a field that fails the agent schema', + 'researcher', + { topic: 3 }, + {}, + 'Input validation failed for tool researcher: Validation failed: Invalid input: expected string, received number', + ], + [ + 'an unknown sessionId', + 'writer', + { sessionId: 'subagent-missing' }, + { middleware: [storedChildren(oldChild)] }, + 'Unknown sessionId "subagent-missing".', + ], + [ + 'a sessionId of a child of another agent', + 'researcher', + { sessionId: 'subagent-old', topic: 'x' }, + { middleware: [storedChildren(oldChild)] }, + 'Session "subagent-old" belongs to agent "writer", not "researcher".', + ], + [ + 'a background that is not a boolean', + 'writer', + { background: 'true' }, + { binding: host }, + 'prompt and sessionId must be strings, and background a boolean.', + ], + [ + 'background with a sessionId', + 'writer', + { background: true, sessionId: 'subagent-old' }, + { middleware: [storedChildren(oldChild)], binding: host }, + 'background cannot continue a sessionId.', + ], + ] + it.each(toolErrors)( + 'returns a tool error for %s', + async (_label, agent, args, options, error) => { + const { chunks, result } = await callSubagent(agent, args, options) + + expect(result).toEqual({ error }) + expect(runs).toEqual([]) + // The parent run goes on after the tool error. + expect(chunks.some((chunk) => chunk.type === EventType.RUN_ERROR)).toBe( + false, + ) + }, + ) +}) diff --git a/testing/e2e/fixtures/subagents/basic.json b/testing/e2e/fixtures/subagents/basic.json index 284943524..df15dad37 100644 --- a/testing/e2e/fixtures/subagents/basic.json +++ b/testing/e2e/fixtures/subagents/basic.json @@ -140,6 +140,72 @@ "response": { "content": "Pricing is ready." } + }, + { + "match": { + "userMessage": "[subagent-limits] research two topics", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "researcher", + "arguments": "{\"task\":\"[subagent-limits-task] squids\"}" + }, + { + "name": "researcher", + "arguments": "{\"task\":\"[subagent-limits-task] octopuses\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[subagent-limits-task] squids" + }, + "response": { + "content": "Squids have three hearts." + } + }, + { + "match": { + "userMessage": "[subagent-limits] research two topics", + "sequenceIndex": 1 + }, + "response": { + "content": "One topic is done." + } + }, + { + "match": { + "userMessage": "[subagent-keys] write a draft", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "writer", + "arguments": "{\"task\":\"[subagent-keys-task] Write a draft.\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[subagent-keys-task] Write a draft." + }, + "response": { + "content": "The draft is written." + } + }, + { + "match": { + "userMessage": "[subagent-keys] write a draft", + "sequenceIndex": 1 + }, + "response": { + "content": "The writer is done." + } } ] } diff --git a/testing/e2e/fixtures/subagents/session.json b/testing/e2e/fixtures/subagents/session.json new file mode 100644 index 000000000..17e8241d8 --- /dev/null +++ b/testing/e2e/fixtures/subagents/session.json @@ -0,0 +1,59 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[subagent-session] take squid notes", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "researcher", + "arguments": "{\"prompt\":\"[subagent-session-task] Note one squid fact.\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[subagent-session-task] Note one squid fact.", + "sequenceIndex": 0 + }, + "response": { + "content": "Squids have three hearts." + } + }, + { + "match": { + "userMessage": "[subagent-session] take squid notes", + "sequenceIndex": 1 + }, + "response": { + "toolCalls": [ + { + "name": "researcher", + "arguments": "{\"sessionId\":\"first-child\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[subagent-session-task] Note one squid fact.", + "sequenceIndex": 1 + }, + "response": { + "content": "Squids also have blue blood." + } + }, + { + "match": { + "userMessage": "[subagent-session] take squid notes", + "sequenceIndex": 2 + }, + "response": { + "content": "Both notes are ready." + } + } + ] +} diff --git a/testing/e2e/src/lib/providers.ts b/testing/e2e/src/lib/providers.ts index 6b2fe8636..bac63a8b2 100644 --- a/testing/e2e/src/lib/providers.ts +++ b/testing/e2e/src/lib/providers.ts @@ -28,7 +28,8 @@ import type { AnyTextAdapter } from '@tanstack/ai' import type { BytePlusChatModel } from '@tanstack/ai-byteplus' import type { Feature, Provider } from '@/lib/types' -const LLMOCK_DEFAULT_BASE = process.env.LLMOCK_URL || 'http://127.0.0.1:4010' +export const LLMOCK_DEFAULT_BASE = + process.env.LLMOCK_URL || 'http://127.0.0.1:4010' const DUMMY_KEY = 'sk-e2e-test-dummy-key' /** diff --git a/testing/e2e/src/lib/subagents-test.ts b/testing/e2e/src/lib/subagents-test.ts index 1b50135a9..2748d182f 100644 --- a/testing/e2e/src/lib/subagents-test.ts +++ b/testing/e2e/src/lib/subagents-test.ts @@ -24,6 +24,12 @@ export type SubagentScenario = | 'tool' | 'brief' | 'result' + | 'limits' + | 'session' + | 'keys' + +/** The key the `keys` scenario passes in `subagents.keys`. */ +export const SUBAGENT_E2E_KEY = 'sk-e2e-subagent-key' /** The user message for each scenario. It matches `fixtures/subagents`. */ export const SUBAGENT_PROMPTS: Record = { @@ -32,6 +38,9 @@ export const SUBAGENT_PROMPTS: Record = { tool: '[subagent-tool] research squids', brief: '[subagent-brief] compare squid facts', result: '[subagent-result] price vendor a', + limits: '[subagent-limits] research two topics', + session: '[subagent-session] take squid notes', + keys: '[subagent-keys] write a draft', } export function isSubagentScenario(value: unknown): value is SubagentScenario { @@ -40,6 +49,9 @@ export function isSubagentScenario(value: unknown): value is SubagentScenario { value === 'approval' || value === 'tool' || value === 'brief' || - value === 'result' + value === 'result' || + value === 'limits' || + value === 'session' || + value === 'keys' ) } diff --git a/testing/e2e/src/routes/api.subagents-test.ts b/testing/e2e/src/routes/api.subagents-test.ts index b6306745c..0e17b63a3 100644 --- a/testing/e2e/src/routes/api.subagents-test.ts +++ b/testing/e2e/src/routes/api.subagents-test.ts @@ -3,12 +3,23 @@ import { chat, chatParamsFromRequestBody, defineAgent, + keyedAdapter, maxIterations, toServerSentEventsResponse, } from '@tanstack/ai' +import { memoryPersistence, withPersistence } from '@tanstack/ai-persistence' import { z } from 'zod' -import type { DefinedAgent, ReasoningRequest, Tool } from '@tanstack/ai' -import { createTextAdapter } from '@/lib/providers' +import type { + ChatMiddleware, + DefinedAgent, + ReasoningRequest, + Tool, +} from '@tanstack/ai' +import { createOpenaiChat } from '@tanstack/ai-openai' +import { openaiByok } from '@tanstack/ai-openai/byok' +import { getByokKey } from '@tanstack/ai/byok/server' +import type { ByokProvider } from '@tanstack/ai/byok' +import { LLMOCK_DEFAULT_BASE, createTextAdapter } from '@/lib/providers' import { deleteLogs, isSubagentScenario, @@ -30,11 +41,29 @@ import type { SubagentScenario } from '@/lib/subagents-test' * child answers, and the parent reads its result. * - `brief`: no router. The parent model writes a `task` for `researcher`, * and the child's only message is that task. + * - `limits`: like `brief`, with `limits: { maxCalls: 1 }`. The parent model + * calls `researcher` twice in one turn. The second call gets a tool error. + * - `keys`: like `result`, but the page sends an `x-byok-openai` header. + * `subagents.keys` reads it with `getByokKey`, and `writer` builds its + * adapter with `ctx.keys.adapter(...)`. The child request carries that + * key. aimock hides `Authorization` in its journal, so the request also + * sends it as `x-e2e-authorization`. * - `result`: no router. `pricer` calls the bound `ctx.chat` with * `stream: false`, so its `run` resolves to a string. The string streams as * the child's text and goes back to the parent model as the result. + * - `session`: no router. The parent model starts `researcher`, then a + * second `researcher` call continues it by its `sessionId`. + * `withPersistence` stores the child, so that call loads its transcript. */ +/** Send `Authorization` again as `x-e2e-authorization`, which aimock keeps. */ +const echoAuthorization: typeof fetch = (input, init) => { + const headers = new Headers(init?.headers) + headers.set('x-e2e-authorization', headers.get('authorization') ?? '') + return fetch(input, { ...init, headers }) +} + function subagentsFor( + request: Request, scenario: SubagentScenario, aimockPort: number | undefined, testId: string | undefined, @@ -78,7 +107,7 @@ function subagentsFor( ]) return { agents: [cleaner], router: () => 'cleaner' } } - if (scenario === 'brief') { + if (scenario === 'brief' || scenario === 'limits') { const researcher = defineAgent({ name: 'researcher', description: 'Researches one focused task', @@ -96,7 +125,9 @@ function subagentsFor( }) // Each branch returns its own bag, so give them one agent type. const agents: Array = [researcher] - return { agents } + return scenario === 'limits' + ? { agents, limits: { maxCalls: 1 } } + : { agents } } if (scenario === 'result') { const pricer = defineAgent({ @@ -113,9 +144,70 @@ function subagentsFor( const agents: Array = [pricer] return { agents } } + if (scenario === 'keys') { + const gpt = keyedAdapter(openaiByok, (key) => + createOpenaiChat('gpt-4o', key, { + baseURL: `${LLMOCK_DEFAULT_BASE}/v1`, + fetch: echoAuthorization, + ...(testId ? { defaultHeaders: { 'X-Test-Id': testId } } : {}), + }), + ) + const writer = defineAgent({ + name: 'writer', + description: 'Writes a short draft', + inputSchema: z.object({ task: z.string() }), + run: async (ctx) => + ctx.chat({ + adapter: await ctx.keys.adapter(gpt), + messages: [{ role: 'user', content: ctx.input.task }], + stream: false, + }), + }) + const agents: Array = [writer] + return { + agents, + keys: (provider: ByokProvider | string) => getByokKey(request, provider), + } + } return { agents: [child('researcher', [])] } } +/** + * aimock replies are fixed, so the second `researcher` call in + * `fixtures/subagents/session.json` cannot hold the random id of the + * first child. It sends `sessionId: 'first-child'`. This middleware does what + * a real model does: it copies the `subagentRunId` of the first result. + */ +function fillSessionId(): ChatMiddleware { + let first: string | undefined + return { + name: 'fill-session-id', + onAfterToolCall(_ctx, { result }) { + if ( + first === undefined && + typeof result === 'object' && + result !== null && + 'subagentRunId' in result && + typeof result.subagentRunId === 'string' + ) { + first = result.subagentRunId + } + }, + onBeforeToolCall(_ctx, { args }) { + if ( + first === undefined || + typeof args !== 'object' || + args === null || + !('sessionId' in args) || + args.sessionId !== 'first-child' + ) { + return undefined + } + return { type: 'transformArgs', args: { ...args, sessionId: first } } + }, + } +} + export const Route = createFileRoute('/api/subagents-test')({ server: { handlers: { @@ -146,7 +238,16 @@ export const Route = createFileRoute('/api/subagents-test')({ runId: params.runId, ...(params.parentRunId ? { parentRunId: params.parentRunId } : {}), ...(params.resume ? { resume: params.resume } : {}), - subagents: subagentsFor(scenario, aimockPort, testId), + subagents: subagentsFor(request, scenario, aimockPort, testId), + // One run holds both calls, so a store per request is enough. + ...(scenario === 'session' + ? { + middleware: [ + withPersistence(memoryPersistence()), + fillSessionId(), + ], + } + : {}), agentLoopStrategy: maxIterations(4), abortController, }) diff --git a/testing/e2e/src/routes/subagents-test.tsx b/testing/e2e/src/routes/subagents-test.tsx index 560aa680a..717b07a20 100644 --- a/testing/e2e/src/routes/subagents-test.tsx +++ b/testing/e2e/src/routes/subagents-test.tsx @@ -6,6 +6,7 @@ import { SUBAGENT_PROMPTS, deleteLogs, isSubagentScenario, + SUBAGENT_E2E_KEY, } from '@/lib/subagents-test' const clientTools = [deleteLogs.client()] as const @@ -27,7 +28,13 @@ function SubagentsTestPage() { const { testId, aimockPort, scenario } = Route.useSearch() const { messages, sendMessage, interrupts, isLoading } = useChat({ threadId: `subagents-${testId ?? 'manual'}-${scenario}`, - connection: fetchServerSentEvents('/api/subagents-test'), + // The `keys` scenario sends the user's key the way a BYOK store does. + connection: fetchServerSentEvents( + '/api/subagents-test', + scenario === 'keys' + ? { headers: { 'x-byok-openai': SUBAGENT_E2E_KEY } } + : {}, + ), body: { scenario, testId, aimockPort }, tools: clientTools, }) @@ -64,6 +71,9 @@ function SubagentsTestPage() { {cards.map((card) => (

{card.status} + + {card.input === undefined ? '' : JSON.stringify(card.input)} +
    {describeParts(card.messages).map((line, index) => (
  • {line}
    • diff --git a/testing/e2e/tests/subagents.spec.ts b/testing/e2e/tests/subagents.spec.ts index 2461f4e8a..90f3709e9 100644 --- a/testing/e2e/tests/subagents.spec.ts +++ b/testing/e2e/tests/subagents.spec.ts @@ -14,7 +14,15 @@ async function run(page: Page) { function open( page: Page, - scenario: 'route' | 'approval' | 'tool' | 'brief' | 'result', + scenario: + | 'route' + | 'approval' + | 'tool' + | 'brief' + | 'result' + | 'limits' + | 'session' + | 'keys', testId: string, aimockPort: number, ) { @@ -154,6 +162,10 @@ test.describe('subagents', () => { await expect(page.getByTestId('card-researcher')).toContainText( 'text:Squids have three hearts and blue blood.', ) + // The card keeps the checked input the parent model wrote. + await expect(page.getByTestId('card-input-researcher')).toHaveText( + JSON.stringify({ task: '[subagent-brief-task] List two squid facts.' }), + ) }) test('a child whose run resolves to a string reports it to the parent', async ({ @@ -172,4 +184,104 @@ test.describe('subagents', () => { 'text:Vendor A costs 42 EUR.', ) }) + + test('subagents.limits refuses a call past maxCalls with a tool error', async ({ + page, + request, + testId, + aimockPort, + }) => { + await open(page, 'limits', testId, aimockPort) + await run(page) + + await expect(page.getByTestId('parent-text')).toHaveText( + 'One topic is done.', + ) + await expect(page.getByTestId('card-researcher')).toHaveCount(1) + await expect(page.getByTestId('card-researcher')).toContainText( + 'text:Squids have three hearts.', + ) + + const journal = await request.get( + `http://127.0.0.1:${aimockPort}/v1/_requests`, + ) + const entries = (await journal.json()) as Array<{ + headers?: Record + body?: { + messages?: Array<{ role?: string; content?: string | null }> + } | null + }> + const toolMessages = entries + .filter((entry) => entry.headers?.['x-test-id'] === testId) + .flatMap((entry) => entry.body?.messages ?? []) + .filter((message) => message.role === 'tool') + expect( + toolMessages.some((message) => + message.content?.includes('subagent limit reached (maxCalls 1)'), + ), + ).toBe(true) + }) + + test('an agent tool continues a child by its sessionId', async ({ + page, + testId, + aimockPort, + }) => { + await open(page, 'session', testId, aimockPort) + await run(page) + + // The parent model runs a third time only after the second call gave it + // the child's result. + await expect(page.getByTestId('parent-text')).toHaveText( + 'Both notes are ready.', + ) + await expect(page.getByTestId('card-status-researcher')).toHaveText( + 'finished', + ) + const card = page.getByTestId('card-researcher') + await expect(card).toContainText('text:Squids have three hearts.') + // The second call has no prompt, so the child's last user message is the + // task it stored from the first call. Only that fixture gives this reply. + // A child that did not get its earlier messages would get no reply. + await expect(card).toContainText('text:Squids also have blue blood.') + // One card: the second call continued the same child. + await expect(card).toHaveCount(1) + }) + + test('subagents.keys gives the child its key', async ({ + page, + request, + testId, + aimockPort, + }) => { + await open(page, 'keys', testId, aimockPort) + await run(page) + + await expect(page.getByTestId('parent-text')).toHaveText( + 'The writer is done.', + ) + await expect(page.getByTestId('card-writer')).toContainText( + 'text:The draft is written.', + ) + + const journal = await request.get( + `http://127.0.0.1:${aimockPort}/v1/_requests`, + ) + const entries = (await journal.json()) as Array<{ + headers?: Record + body?: { + messages?: Array<{ role?: string; content?: string | null }> + } | null + }> + const child = entries.find( + (entry) => + entry.headers?.['x-test-id'] === testId && + entry.body?.messages?.some((message) => + message.content?.includes('[subagent-keys-task]'), + ), + ) + expect(child?.headers?.['x-e2e-authorization']).toBe( + 'Bearer sk-e2e-subagent-key', + ) + }) })