From 7d35b6ed25f1035463bb557026f60205536d8b25 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 12 Oct 2026 19:12:01 +0200 Subject: [PATCH 1/4] feat(ai): subagent limits, ctx.keys, ctx.step, and ctx.agents subagents.limits caps depth, calls, concurrency, and time across the whole subagent tree with one shared SubagentBudget. Agents get ctx.keys (ProviderKeys) from the binding or the provider env names, plus ctx.step and ctx.agents that a host can bind. --- .changeset/agent-limits-keys-step.md | 10 + docs/advanced/byok.md | 29 ++ docs/chat/subagents.md | 45 +++ docs/config.json | 2 +- .../ai/src/activities/chat/agents/bound.ts | 88 +++++- .../activities/chat/agents/define-agent.ts | 19 +- .../ai/src/activities/chat/agents/limits.ts | 75 +++++ .../ai/src/activities/chat/agents/spawn.ts | 295 ++++++++++++------ packages/ai/src/activities/index.ts | 4 + packages/ai/src/byok.ts | 6 +- packages/ai/src/byok/env-keys.ts | 37 +++ packages/ai/src/byok/get-key.ts | 22 +- packages/ai/src/byok/keyed.ts | 19 ++ packages/ai/src/client.ts | 3 + packages/ai/src/index.ts | 8 + packages/ai/tests/agent-start.test.ts | 77 +++++ packages/ai/tests/agent-step.test.ts | 65 ++++ packages/ai/tests/provider-keys.test.ts | 160 +++++++++- packages/ai/tests/subagent-limits.test.ts | 186 +++++++++++ testing/e2e/fixtures/subagents/basic.json | 35 +++ testing/e2e/src/lib/subagents-test.ts | 5 +- testing/e2e/src/routes/api.subagents-test.ts | 8 +- testing/e2e/tests/subagents.spec.ts | 39 ++- 23 files changed, 1121 insertions(+), 116 deletions(-) create mode 100644 .changeset/agent-limits-keys-step.md create mode 100644 packages/ai/src/activities/chat/agents/limits.ts create mode 100644 packages/ai/src/byok/env-keys.ts create mode 100644 packages/ai/tests/agent-start.test.ts create mode 100644 packages/ai/tests/agent-step.test.ts create mode 100644 packages/ai/tests/subagent-limits.test.ts diff --git a/.changeset/agent-limits-keys-step.md b/.changeset/agent-limits-keys-step.md new file mode 100644 index 0000000000..857a6d5ac7 --- /dev/null +++ b/.changeset/agent-limits-keys-step.md @@ -0,0 +1,10 @@ +--- +'@tanstack/ai': minor +--- + +Subagent limits, provider keys for agents, and host hooks for agent steps and background 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. +- **`ctx.keys`** (`ProviderKeys`: `get`, `require`, and `adapter`) finds provider keys in agent code. A host sets them with `SubagentBinding.keys`. Without a host, `ctx.keys` reads each provider's `env` names, and a missing key throws an error that names the env var. +- **`ctx.step.do(name, fn)`** (`AgentStep`). A host can set `SubagentBinding.step` to keep the value of `fn`, so a side effect does not run again when the host runs the agent again. Without a host, `fn` runs each time. +- **`ctx.agents.start()`** (`AgentStarter`, `AgentRunHandle`) starts an agent in the background when a host binds it with `SubagentBinding.agents`. Without a host, it throws. diff --git a/docs/advanced/byok.md b/docs/advanced/byok.md index 3a8f49415f..b740c70273 100644 --- a/docs/advanced/byok.md +++ b/docs/advanced/byok.md @@ -326,6 +326,35 @@ 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 + +An agent that calls a model needs a key too. Wrap the adapter in `keyedAdapter`. Before the call, the agent makes it with `ctx.keys`: + +```typescript +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, + }), +}); +``` + +- Without a host, `ctx.keys` reads the env var of the provider, here `OPENAI_API_KEY`. If it is empty, the call throws `Missing OpenAI API key. Set OPENAI_API_KEY.` +- A host can give its own keys in `subagents.binding.keys`. A child agent gives them to its own children. +- `ctx.keys.get(provider)` returns the key or `null`. `ctx.keys.require(provider)` throws when there is no key. + ## Image, audio, and other providers For other cases: diff --git a/docs/chat/subagents.md b/docs/chat/subagents.md index 2bf5addb78..0bf9d32bfb 100644 --- a/docs/chat/subagents.md +++ b/docs/chat/subagents.md @@ -261,6 +261,51 @@ 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`. + +## Run context for hosts + +A host that runs agents for a long session can give each run more on `subagents.binding`: + +- `keys`: the provider keys that the agent reads as `ctx.keys`. See [Keys inside agents](../advanced/byok#keys-inside-agents). +- `step`: what `ctx.step.do(name, fn)` uses. A host can store the value of `fn`, so that a side effect does not run again when the host runs the agent again after a crash. Without a host, `fn` runs each time. +- `agents`: what `ctx.agents.start(agent, input)` uses to start an agent in the background. Without a host, `ctx.agents.start` throws. + +```ts group=subagent-values +const charger = defineAgent({ + name: 'charger', + description: 'Charges the card once', + run: async (ctx) => { + const receipt = await ctx.step.do('charge', () => ({ id: 'rcpt-1' })) + return receipt + }, +}) +``` + ## Strategy - `exclusive` (default): the chosen child owns the turn. Main does not answer after it. diff --git a/docs/config.json b/docs/config.json index b203782dd1..57f3c1b723 100644 --- a/docs/config.json +++ b/docs/config.json @@ -1128,7 +1128,7 @@ "label": "Bring Your Own Key", "to": "advanced/byok", "addedAt": "2026-08-18", - "updatedAt": "2026-10-09" + "updatedAt": "2026-10-12" }, { "label": "Debug Logging", diff --git a/packages/ai/src/activities/chat/agents/bound.ts b/packages/ai/src/activities/chat/agents/bound.ts index c859a5822a..c7746fc640 100644 --- a/packages/ai/src/activities/chat/agents/bound.ts +++ b/packages/ai/src/activities/chat/agents/bound.ts @@ -13,8 +13,14 @@ import { rerank } from '../../rerank/index' import { decide } from '../../evaluate/index' 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 { + ContentPart, + PromptCacheRetention, + RunAgentResumeItem, +} from '../../../types' +import type { DefinedAgent, SubagentRunInput } from './define-agent' +import type { SubagentBudget } from './limits' +import type { ProviderKeys } from '../../../byok/keyed' /** * The fields a child `chat()` needs, in one spread: @@ -32,11 +38,73 @@ export interface SubagentForward { promptCache?: PromptCacheRetention } +/** The steps of an agent run, as `ctx.step`. */ +export interface AgentStep { + /** + * Run `fn` for `name`, and return its value. When a host runs a stopped + * agent again, a finished step returns its stored value and `fn` does not + * run again. Without such a host, `fn` runs each time. The value must be + * JSON. Use a name that is the same on every run. + */ + do: (name: string, fn: () => T | Promise) => Promise +} + +/** + * An agent run that `ctx.agents.start` started in the background. Await it + * for the run's result. + */ +export interface AgentRunHandle< + TResult = unknown, +> extends PromiseLike { + readonly id: string + /** + * Add a message to the run: `'steer'` (default) for its next model call, + * or `'followUp'` to run the agent again after the run ends. The host + * decides the details. + */ + send: ( + message: string | Array, + options?: { mode?: 'steer' | 'followUp'; inputId?: string }, + ) => Promise + /** Stop the run. */ + cancel: () => Promise +} + +/** `ctx.agents`: start agents in the background. A host gives it. */ +export interface AgentStarter { + /** + * Start `agent` (a definition, or a name the host knows) in the + * background, and return its run at once. Throws when the host refuses + * it, for example over the subagent tree budget. + */ + start: ( + agent: DefinedAgent | string, + input?: unknown, + options?: { + wake?: boolean + attach?: 'reference' | 'none' + resume?: boolean + }, + ) => AgentRunHandle +} + /** * What a host adds to every activity call an agent makes through `ctx`. Apps * using plain `chat({ subagents })` do not set it. */ export interface SubagentBinding { + /** + * The steps the agent reads as `ctx.step`. Only this agent run gets them, + * not the children it starts. Without them, `ctx.step.do` runs `fn` each + * time. + */ + step?: AgentStep + /** + * What the agent reads as `ctx.agents`: it starts agents in the + * background. Only this agent run gets it, not the children it starts in + * a `ctx.chat({ subagents })`. Without it, `ctx.agents.start` throws. + */ + agents?: AgentStarter /** * Added before the call's own middleware on every `ctx.chat` call. A * child's `ctx.chat({ subagents })` passes it down to its own children. @@ -47,6 +115,17 @@ 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`. 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`. @@ -128,12 +207,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, @@ -143,6 +224,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 df0a0c0ec2..17dfdb1769 100644 --- a/packages/ai/src/activities/chat/agents/define-agent.ts +++ b/packages/ai/src/activities/chat/agents/define-agent.ts @@ -10,7 +10,13 @@ import type { UIMessage, } from '../../../types' import type { AnyClientTool } from '../tools/tool-definition' -import type { BoundActivities, SubagentForward } from './bound' +import type { + AgentStarter, + AgentStep, + BoundActivities, + SubagentForward, +} from './bound' +import type { ProviderKeys } from '../../../byok/keyed' /** * What the library knows about a child run before `run` starts. @@ -55,11 +61,22 @@ 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(...)`. A host sets the keys. Without a host, they come + * from each provider's `env` names, and a missing key throws an error that + * names the env var. + * - `step.do(name, fn)` keeps a side effect from running twice when a host + * runs the agent again after a crash. See {@link AgentStep}. + * - `agents.start(agent, input, options)` starts an agent in the + * background, when a host runs this agent. Without such a host it throws. */ export type SubagentRunContext = SubagentRunInput & BoundActivities & { forward: SubagentForward + keys: ProviderKeys + step: AgentStep + agents: AgentStarter } /** 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 0000000000..7210880ec0 --- /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 11e0e8ec64..5c633f5e20 100644 --- a/packages/ai/src/activities/chat/agents/spawn.ts +++ b/packages/ai/src/activities/chat/agents/spawn.ts @@ -26,8 +26,11 @@ import type { Tool, UIMessage, } from '../../../types' +import { envProviderKeys } from '../../../byok/env-keys' import { createBoundActivities } from './bound' -import type { SubagentBinding } from './bound' +import { SubagentBudget } from './limits' +import type { AgentStarter, AgentStep, SubagentBinding } from './bound' +import type { SubagentLimits } from './limits' import type { DefinedAgent, SubagentRunContext, @@ -87,6 +90,13 @@ 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 } /** What the children of one parent run left behind for the parent terminal. */ @@ -236,13 +246,36 @@ 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 @@ -494,8 +527,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, the steps, `agents`, and the bound activity functions. */ function runContext( agentName: string, @@ -518,10 +551,25 @@ function runContext( abortController, ...(binding?.promptCache ? { promptCache: binding.promptCache } : {}), }, + keys: binding?.keys ?? envProviderKeys, + step: binding?.step ?? runEachTime, + agents: binding?.agents ?? noAgents, ...createBoundActivities(agentName, input, abortController, binding), } satisfies SubagentRunContext } +/** The steps of an agent run that no host can run again: `fn` runs each time. */ +const runEachTime: AgentStep = { do: async (_name, fn) => fn() } + +/** `ctx.agents` of a run that no host binds: a start throws. */ +const noAgents: AgentStarter = { + start: () => { + throw new Error( + 'ctx.agents.start needs a host that runs agents in the background.', + ) + }, +} + /** * Run one agent and stream its chunks, from `SUBAGENT_STARTED` to * `SUBAGENT_FINISHED` or `SUBAGENT_ERROR`. The package root exports it as @@ -841,106 +889,155 @@ export function createSyntheticSubagentTools( sink: SubagentSink }, ): 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 } : {}), + } satisfies SubagentToolOutcome + } + return bag.agents.map((agent) => ({ name: agent.name, description: agent.description, ...(agent.inputSchema !== undefined && { inputSchema: agent.inputSchema }), [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, - }, - }), - } - 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, - { - 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, - )) { - 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 - } - return { - subagentRunId, - text, - ...(result !== undefined ? { result } : {}), - ...(error !== undefined ? { error } : {}), - ...(sink.interrupts.length > 0 ? { interrupts: sink.interrupts } : {}), - } satisfies SubagentToolOutcome - }, + // An agent without `inputSchema` still gets `{}` from the model. Its + // `ctx.input` stays undefined. + execute: (input: unknown, context?: unknown) => + runChild( + { + name: agent.name, + ...(agent.inputSchema !== undefined && { input }), + }, + context, + ), })) } diff --git a/packages/ai/src/activities/index.ts b/packages/ai/src/activities/index.ts index 5e9322e87a..8aa956d489 100644 --- a/packages/ai/src/activities/index.ts +++ b/packages/ai/src/activities/index.ts @@ -50,9 +50,13 @@ export { } from './chat/agents/define-agent' export type { BoundActivities, + AgentRunHandle, + AgentStarter, + AgentStep, 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 550a12891e..010672f553 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 0000000000..869bde495f --- /dev/null +++ b/packages/ai/src/byok/env-keys.ts @@ -0,0 +1,37 @@ +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}` +} + +async function requireEnvKey(provider: ByokProvider | ProviderId) { + const key = envKey(provider) + if (key === null) { + const error = new ByokMissingError(resolveProviderId(provider)) + error.message = missingKeyMessage(provider) + throw error + } + return key +} + +/** + * The `ctx.keys` of an agent whose host sets no keys: each provider's `env` + * names. Not in `@tanstack/ai/byok`, so the client entry has no env access. + */ +export const envProviderKeys: ProviderKeys = { + get: async (provider) => envKey(provider), + require: requireEnvKey, + adapter: async (adapter) => + isKeyedAdapter(adapter) + ? adapter.create(await requireEnvKey(adapter.provider)) + : adapter, +} diff --git a/packages/ai/src/byok/get-key.ts b/packages/ai/src/byok/get-key.ts index dec1a5b667..8b2318c168 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 51ec7c0992..f0c4b5d429 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 19ba1ecb99..9e94b0d8f6 100644 --- a/packages/ai/src/client.ts +++ b/packages/ai/src/client.ts @@ -236,6 +236,9 @@ export { } from './activities/chat/agents/define-agent' export type { BoundActivities, + AgentRunHandle, + AgentStarter, + AgentStep, SubagentBinding, SubagentForward, } from './activities/chat/agents/bound' diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 8a6514f44e..5b70896670 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -96,15 +96,23 @@ export { } from './activities/chat/agents/define-agent' export type { BoundActivities, + AgentRunHandle, + AgentStarter, + AgentStep, 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 { subagentRoute, diff --git a/packages/ai/tests/agent-start.test.ts b/packages/ai/tests/agent-start.test.ts new file mode 100644 index 0000000000..53a436e9fe --- /dev/null +++ b/packages/ai/tests/agent-start.test.ts @@ -0,0 +1,77 @@ +import { describe, expect, it } from 'vitest' +import { EventType } from '../src/types' +import { defineAgent } from '../src/activities/chat/agents/define-agent' +import { spawnAgentStream } from '../src/activities/chat/agents/spawn' +import type { AgentStarter } from '../src/activities/chat/agents/bound' +import type { SubagentRunInput } from '../src/activities/chat/agents/define-agent' +import type { StreamChunk } from '../src/types' + +const input: SubagentRunInput = { + input: undefined, + messages: [], + threadId: 't', + runId: 'r', + parentRunId: 'p', + subagentRunId: 's', +} + +async function collect(stream: AsyncIterable) { + const chunks: Array = [] + for await (const chunk of stream) chunks.push(chunk) + return chunks +} + +describe('ctx.agents of an agent run', () => { + it('throws a clear error when no host binds it', async () => { + const agent = defineAgent({ + name: 'lead', + description: 'Starts a helper', + run: async (ctx) => ctx.agents.start('helper'), + }) + + const chunks = await collect(spawnAgentStream(agent, input)) + + expect( + chunks.find((chunk) => chunk.type === EventType.SUBAGENT_ERROR), + ).toMatchObject({ + message: expect.stringContaining('ctx.agents.start needs a host'), + }) + }) + + it('starts agents through the agents of the host binding', async () => { + const started: Array = [] + const agents: AgentStarter = { + start: (agent, childInput, options) => { + started.push({ agent, input: childInput, options }) + return Object.assign(Promise.resolve('helped'), { + id: 'op-helper', + send: async () => undefined, + cancel: async () => undefined, + }) + }, + } + const agent = defineAgent({ + name: 'lead', + description: 'Starts a helper', + run: async (ctx) => { + const child = ctx.agents.start( + 'helper', + { topic: 'tea' }, + { wake: true }, + ) + return { id: child.id, result: await child } + }, + }) + + const chunks = await collect( + spawnAgentStream(agent, input, undefined, undefined, { agents }), + ) + + expect(started).toEqual([ + { agent: 'helper', input: { topic: 'tea' }, options: { wake: true } }, + ]) + expect( + chunks.find((chunk) => chunk.type === EventType.SUBAGENT_FINISHED), + ).toMatchObject({ result: { id: 'op-helper', result: 'helped' } }) + }) +}) diff --git a/packages/ai/tests/agent-step.test.ts b/packages/ai/tests/agent-step.test.ts new file mode 100644 index 0000000000..9569ba434d --- /dev/null +++ b/packages/ai/tests/agent-step.test.ts @@ -0,0 +1,65 @@ +import { describe, expect, it } from 'vitest' +import { defineAgent } from '../src/activities/chat/agents/define-agent' +import { spawnAgentStream } from '../src/activities/chat/agents/spawn' +import type { AgentStep } from '../src/activities/chat/agents/bound' +import type { SubagentRunInput } from '../src/activities/chat/agents/define-agent' + +const input: SubagentRunInput = { + input: undefined, + messages: [], + threadId: 't', + runId: 'r', + parentRunId: 'p', + subagentRunId: 's', +} + +async function drain(stream: AsyncIterable) { + for await (const _chunk of stream) { + // Run the agent to its end. + } +} + +describe('ctx.step of an agent run', () => { + it('runs fn each time when no host gives steps', async () => { + let runs = 0 + const agent = defineAgent({ + name: 'counter', + description: 'Counts', + run: async (ctx) => ctx.step.do('count', () => (runs += 1)), + }) + + await drain(spawnAgentStream(agent, input)) + await drain(spawnAgentStream(agent, input)) + + expect(runs).toBe(2) + }) + + it('uses the steps of the host binding', async () => { + const names: Array = [] + const step: AgentStep = { + do: async (name: string) => { + names.push(name) + return 'stored' as T + }, + } + let ran = false + let value: unknown + const agent = defineAgent({ + name: 'charger', + description: 'Charges once', + run: async (ctx) => { + value = await ctx.step.do('charge', () => { + ran = true + return 'fresh' + }) + return value + }, + }) + + await drain(spawnAgentStream(agent, input, undefined, undefined, { step })) + + expect(names).toEqual(['charge']) + expect(ran).toBe(false) + expect(value).toBe('stored') + }) +}) diff --git a/packages/ai/tests/provider-keys.test.ts b/packages/ai/tests/provider-keys.test.ts index 2bf0285b2c..0c0a6d4a99 100644 --- a/packages/ai/tests/provider-keys.test.ts +++ b/packages/ai/tests/provider-keys.test.ts @@ -7,7 +7,13 @@ 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 { 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' const acme = defineByokProvider({ id: 'acme', label: 'Acme' }) @@ -102,3 +108,155 @@ 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, + ) + }) +}) diff --git a/packages/ai/tests/subagent-limits.test.ts b/packages/ai/tests/subagent-limits.test.ts new file mode 100644 index 0000000000..9e74ad74c4 --- /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/testing/e2e/fixtures/subagents/basic.json b/testing/e2e/fixtures/subagents/basic.json index 284943524d..ac46395da5 100644 --- a/testing/e2e/fixtures/subagents/basic.json +++ b/testing/e2e/fixtures/subagents/basic.json @@ -140,6 +140,41 @@ "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." + } } ] } diff --git a/testing/e2e/src/lib/subagents-test.ts b/testing/e2e/src/lib/subagents-test.ts index 1b50135a93..3bd8b62dc6 100644 --- a/testing/e2e/src/lib/subagents-test.ts +++ b/testing/e2e/src/lib/subagents-test.ts @@ -24,6 +24,7 @@ export type SubagentScenario = | 'tool' | 'brief' | 'result' + | 'limits' /** The user message for each scenario. It matches `fixtures/subagents`. */ export const SUBAGENT_PROMPTS: Record = { @@ -32,6 +33,7 @@ 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', } export function isSubagentScenario(value: unknown): value is SubagentScenario { @@ -40,6 +42,7 @@ export function isSubagentScenario(value: unknown): value is SubagentScenario { value === 'approval' || value === 'tool' || value === 'brief' || - value === 'result' + value === 'result' || + value === 'limits' ) } diff --git a/testing/e2e/src/routes/api.subagents-test.ts b/testing/e2e/src/routes/api.subagents-test.ts index b6306745ce..29c4224265 100644 --- a/testing/e2e/src/routes/api.subagents-test.ts +++ b/testing/e2e/src/routes/api.subagents-test.ts @@ -30,6 +30,8 @@ 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. * - `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. @@ -78,7 +80,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 +98,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({ diff --git a/testing/e2e/tests/subagents.spec.ts b/testing/e2e/tests/subagents.spec.ts index 2461f4e8ab..6e08f62473 100644 --- a/testing/e2e/tests/subagents.spec.ts +++ b/testing/e2e/tests/subagents.spec.ts @@ -14,7 +14,7 @@ async function run(page: Page) { function open( page: Page, - scenario: 'route' | 'approval' | 'tool' | 'brief' | 'result', + scenario: 'route' | 'approval' | 'tool' | 'brief' | 'result' | 'limits', testId: string, aimockPort: number, ) { @@ -172,4 +172,41 @@ 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) + }) }) From d3ed0ac330ebe726873dfb77c179e0bb797cc3ea Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 19:46:16 +0200 Subject: [PATCH 2/4] feat(ai): subagents.keys for ctx.keys, drop ctx.step and ctx.agents ctx.keys now works in a plain chat(): pass provider keys by id on subagents.keys, for example { openai: userKey }. An agent builds a keyed adapter with ctx.keys.adapter. A missing key falls back to the keys of the parent agent, then to the provider's env names. A nested chat's own keys win for those providers only. ctx.step and ctx.agents did nothing without a host (step.do ran fn each time, agents.start threw). They move to the harness work in #1555. --- .changeset/agent-limits-keys-step.md | 6 +- docs/advanced/byok.md | 42 ++++++++- docs/chat/subagents.md | 19 +--- docs/config.json | 4 +- .../ai/src/activities/chat/agents/bound.ts | 76 ++-------------- .../activities/chat/agents/define-agent.ts | 19 +--- .../ai/src/activities/chat/agents/spawn.ts | 26 +++--- packages/ai/src/activities/chat/index.ts | 12 ++- packages/ai/src/activities/index.ts | 3 - packages/ai/src/byok/env-keys.ts | 50 +++++++---- packages/ai/src/client.ts | 3 - packages/ai/src/index.ts | 3 - packages/ai/tests/agent-start.test.ts | 77 ---------------- packages/ai/tests/agent-step.test.ts | 65 -------------- packages/ai/tests/provider-keys.test.ts | 89 +++++++++++++++++++ testing/e2e/fixtures/subagents/basic.json | 31 +++++++ testing/e2e/src/lib/providers.ts | 3 +- testing/e2e/src/lib/subagents-test.ts | 8 +- testing/e2e/src/routes/api.subagents-test.ts | 39 +++++++- testing/e2e/tests/subagents.spec.ts | 46 +++++++++- 20 files changed, 318 insertions(+), 303 deletions(-) delete mode 100644 packages/ai/tests/agent-start.test.ts delete mode 100644 packages/ai/tests/agent-step.test.ts diff --git a/.changeset/agent-limits-keys-step.md b/.changeset/agent-limits-keys-step.md index 857a6d5ac7..459351e760 100644 --- a/.changeset/agent-limits-keys-step.md +++ b/.changeset/agent-limits-keys-step.md @@ -2,9 +2,7 @@ '@tanstack/ai': minor --- -Subagent limits, provider keys for agents, and host hooks for agent steps and background agents. +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. -- **`ctx.keys`** (`ProviderKeys`: `get`, `require`, and `adapter`) finds provider keys in agent code. A host sets them with `SubagentBinding.keys`. Without a host, `ctx.keys` reads each provider's `env` names, and a missing key throws an error that names the env var. -- **`ctx.step.do(name, fn)`** (`AgentStep`). A host can set `SubagentBinding.step` to keep the value of `fn`, so a side effect does not run again when the host runs the agent again. Without a host, `fn` runs each time. -- **`ctx.agents.start()`** (`AgentStarter`, `AgentRunHandle`) starts an agent in the background when a host binds it with `SubagentBinding.agents`. Without a host, it throws. +- **`subagents.keys`** and **`ctx.keys`** (`ProviderKeys`: `get`, `require`, and `adapter`). Pass keys by provider id, for example `{ openai: userKey }`. 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/docs/advanced/byok.md b/docs/advanced/byok.md index b740c70273..6674d9b217 100644 --- a/docs/advanced/byok.md +++ b/docs/advanced/byok.md @@ -328,9 +328,11 @@ Then a send with no pasted key still POSTs. The relay uses the env key. If that ## Keys inside agents -An agent that calls a model needs a key too. Wrap the adapter in `keyedAdapter`. Before the call, the agent makes it with `ctx.keys`: +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`. -```typescript +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"; @@ -351,10 +353,42 @@ export const writer = defineAgent({ }); ``` -- Without a host, `ctx.keys` reads the env var of the provider, here `OPENAI_API_KEY`. If it is empty, the call throws `Missing OpenAI API key. Set OPENAI_API_KEY.` -- A host can give its own keys in `subagents.binding.keys`. A child agent gives them to its own children. +2. In the route, read the user's key. Pass it on `subagents.keys`, by provider id: + +```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: { openai: apiKey } }, + }); + return toServerSentEventsResponse(stream); +} +``` + +How `ctx.keys` finds a key: + +- It reads `subagents.keys` by provider id. The id of `openaiByok` is `openai`. +- If `subagents.keys` has no key for the provider, it reads the env var of the provider, here `OPENAI_API_KEY`. +- If both are empty, 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 0bf9d32bfb..d65931a860 100644 --- a/docs/chat/subagents.md +++ b/docs/chat/subagents.md @@ -287,24 +287,11 @@ The whole tree shares one budget. A child that calls `ctx.chat({ subagents })` p - 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`. -## Run context for hosts +## Keys for agents -A host that runs agents for a long session can give each run more on `subagents.binding`: +An agent that calls a model needs a provider key. An adapter such as `openaiText` reads its env var, for example `OPENAI_API_KEY`. -- `keys`: the provider keys that the agent reads as `ctx.keys`. See [Keys inside agents](../advanced/byok#keys-inside-agents). -- `step`: what `ctx.step.do(name, fn)` uses. A host can store the value of `fn`, so that a side effect does not run again when the host runs the agent again after a crash. Without a host, `fn` runs each time. -- `agents`: what `ctx.agents.start(agent, input)` uses to start an agent in the background. Without a host, `ctx.agents.start` throws. - -```ts group=subagent-values -const charger = defineAgent({ - name: 'charger', - description: 'Charges the card once', - run: async (ctx) => { - const receipt = await ctx.step.do('charge', () => ({ id: 'rcpt-1' })) - return receipt - }, -}) -``` +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 diff --git a/docs/config.json b/docs/config.json index 57f3c1b723..6f26f97fa2 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-12" + "updatedAt": "2026-10-13" }, { "label": "Debug Logging", diff --git a/packages/ai/src/activities/chat/agents/bound.ts b/packages/ai/src/activities/chat/agents/bound.ts index c807037683..1655fd993e 100644 --- a/packages/ai/src/activities/chat/agents/bound.ts +++ b/packages/ai/src/activities/chat/agents/bound.ts @@ -13,12 +13,8 @@ import { rerank } from '../../rerank/index' import { decide } from '../../evaluate/index' import type { GenerationMiddleware } from '../../middleware/types' import type { AnyChatMiddleware } from '../middleware/types' -import type { - ContentPart, - PromptCacheRetention, - RunAgentResumeItem, -} from '../../../types' -import type { DefinedAgent, SubagentRunInput } from './define-agent' +import type { PromptCacheRetention, RunAgentResumeItem } from '../../../types' +import type { SubagentRunInput } from './define-agent' import type { SubagentBudget } from './limits' import type { ProviderKeys } from '../../../byok/keyed' @@ -42,73 +38,11 @@ export interface SubagentForward { promptCache?: PromptCacheRetention } -/** The steps of an agent run, as `ctx.step`. */ -export interface AgentStep { - /** - * Run `fn` for `name`, and return its value. When a host runs a stopped - * agent again, a finished step returns its stored value and `fn` does not - * run again. Without such a host, `fn` runs each time. The value must be - * JSON. Use a name that is the same on every run. - */ - do: (name: string, fn: () => T | Promise) => Promise -} - -/** - * An agent run that `ctx.agents.start` started in the background. Await it - * for the run's result. - */ -export interface AgentRunHandle< - TResult = unknown, -> extends PromiseLike { - readonly id: string - /** - * Add a message to the run: `'steer'` (default) for its next model call, - * or `'followUp'` to run the agent again after the run ends. The host - * decides the details. - */ - send: ( - message: string | Array, - options?: { mode?: 'steer' | 'followUp'; inputId?: string }, - ) => Promise - /** Stop the run. */ - cancel: () => Promise -} - -/** `ctx.agents`: start agents in the background. A host gives it. */ -export interface AgentStarter { - /** - * Start `agent` (a definition, or a name the host knows) in the - * background, and return its run at once. Throws when the host refuses - * it, for example over the subagent tree budget. - */ - start: ( - agent: DefinedAgent | string, - input?: unknown, - options?: { - wake?: boolean - attach?: 'reference' | 'none' - resume?: boolean - }, - ) => AgentRunHandle -} - /** * What a host adds to every activity call an agent makes through `ctx`. Apps * using plain `chat({ subagents })` do not set it. */ export interface SubagentBinding { - /** - * The steps the agent reads as `ctx.step`. Only this agent run gets them, - * not the children it starts. Without them, `ctx.step.do` runs `fn` each - * time. - */ - step?: AgentStep - /** - * What the agent reads as `ctx.agents`: it starts agents in the - * background. Only this agent run gets it, not the children it starts in - * a `ctx.chat({ subagents })`. Without it, `ctx.agents.start` throws. - */ - agents?: AgentStarter /** * Added before the call's own middleware on every `ctx.chat` call. A * child's `ctx.chat({ subagents })` passes it down to its own children. @@ -125,9 +59,9 @@ export interface SubagentBinding { */ budget?: SubagentBudget /** - * The provider keys an agent reads as `ctx.keys`. A child's - * `ctx.chat({ subagents })` passes them down to its own children. Without - * them, `ctx.keys` reads each provider's `env` names. + * 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 /** diff --git a/packages/ai/src/activities/chat/agents/define-agent.ts b/packages/ai/src/activities/chat/agents/define-agent.ts index 17dfdb1769..7a2417749f 100644 --- a/packages/ai/src/activities/chat/agents/define-agent.ts +++ b/packages/ai/src/activities/chat/agents/define-agent.ts @@ -10,12 +10,7 @@ import type { UIMessage, } from '../../../types' import type { AnyClientTool } from '../tools/tool-definition' -import type { - AgentStarter, - AgentStep, - BoundActivities, - SubagentForward, -} from './bound' +import type { BoundActivities, SubagentForward } from './bound' import type { ProviderKeys } from '../../../byok/keyed' /** @@ -62,21 +57,15 @@ export interface SubagentRunInput< * 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(...)`. A host sets the keys. Without a host, they come - * from each provider's `env` names, and a missing key throws an error that - * names the env var. - * - `step.do(name, fn)` keeps a side effect from running twice when a host - * runs the agent again after a crash. See {@link AgentStep}. - * - `agents.start(agent, input, options)` starts an agent in the - * background, when a host runs this agent. Without such a host it throws. + * `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 - step: AgentStep - agents: AgentStarter } /** diff --git a/packages/ai/src/activities/chat/agents/spawn.ts b/packages/ai/src/activities/chat/agents/spawn.ts index 5c633f5e20..fe2607aa08 100644 --- a/packages/ai/src/activities/chat/agents/spawn.ts +++ b/packages/ai/src/activities/chat/agents/spawn.ts @@ -29,8 +29,9 @@ import type { import { envProviderKeys } from '../../../byok/env-keys' import { createBoundActivities } from './bound' import { SubagentBudget } from './limits' -import type { AgentStarter, AgentStep, SubagentBinding } from './bound' +import type { SubagentBinding } from './bound' import type { SubagentLimits } from './limits' +import type { ProviderId } from '../../../byok/providers' import type { DefinedAgent, SubagentRunContext, @@ -97,6 +98,13 @@ export interface SubagentsBag< * apply: a nested chat in a child uses the root's budget. */ limits?: SubagentLimits + /** + * Provider keys by provider id, for example `{ openai: userKey }`. Agents + * read them as `ctx.keys`. A missing or empty key falls back 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?: Readonly> } /** What the children of one parent run left behind for the parent terminal. */ @@ -528,7 +536,7 @@ function* valueResultChunks( /** * Build the context `run` receives: the spawn input plus `forward`, the - * provider keys, the steps, `agents`, and the bound activity functions. + * provider keys, and the bound activity functions. */ function runContext( agentName: string, @@ -552,24 +560,10 @@ function runContext( ...(binding?.promptCache ? { promptCache: binding.promptCache } : {}), }, keys: binding?.keys ?? envProviderKeys, - step: binding?.step ?? runEachTime, - agents: binding?.agents ?? noAgents, ...createBoundActivities(agentName, input, abortController, binding), } satisfies SubagentRunContext } -/** The steps of an agent run that no host can run again: `fn` runs each time. */ -const runEachTime: AgentStep = { do: async (_name, fn) => fn() } - -/** `ctx.agents` of a run that no host binds: a start throws. */ -const noAgents: AgentStarter = { - start: () => { - throw new Error( - 'ctx.agents.start needs a host that runs agents in the background.', - ) - }, -} - /** * Run one agent and stream its chunks, from `SUBAGENT_STARTED` to * `SUBAGENT_FINISHED` or `SUBAGENT_ERROR`. The package root exports it as diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 3ba9ab9bd4..ac7ef6d1c2 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -62,6 +62,7 @@ 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 { collectNamedText, createSubagentSink, @@ -5902,11 +5903,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, }, diff --git a/packages/ai/src/activities/index.ts b/packages/ai/src/activities/index.ts index 8aa956d489..cb8568b91c 100644 --- a/packages/ai/src/activities/index.ts +++ b/packages/ai/src/activities/index.ts @@ -50,9 +50,6 @@ export { } from './chat/agents/define-agent' export type { BoundActivities, - AgentRunHandle, - AgentStarter, - AgentStep, SubagentBinding, SubagentForward, } from './chat/agents/bound' diff --git a/packages/ai/src/byok/env-keys.ts b/packages/ai/src/byok/env-keys.ts index 869bde495f..9c6f208e8a 100644 --- a/packages/ai/src/byok/env-keys.ts +++ b/packages/ai/src/byok/env-keys.ts @@ -13,25 +13,37 @@ function missingKeyMessage(provider: ByokProvider | ProviderId) { return `Missing ${provider.label} API key.${hint}` } -async function requireEnvKey(provider: ByokProvider | ProviderId) { - const key = envKey(provider) - if (key === null) { - const error = new ByokMissingError(resolveProviderId(provider)) - error.message = missingKeyMessage(provider) - throw error - } - return key -} - /** - * The `ctx.keys` of an agent whose host sets no keys: each provider's `env` - * names. Not in `@tanstack/ai/byok`, so the client entry has no env access. + * The `ctx.keys` of an agent: `keys` by provider id, 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 const envProviderKeys: ProviderKeys = { - get: async (provider) => envKey(provider), - require: requireEnvKey, - adapter: async (adapter) => - isKeyedAdapter(adapter) - ? adapter.create(await requireEnvKey(adapter.provider)) - : adapter, +export function providerKeys( + keys: Readonly> = {}, + fallback?: ProviderKeys, +): ProviderKeys { + const get = async (provider: ByokProvider | ProviderId) => + keys[resolveProviderId(provider)] || + (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/client.ts b/packages/ai/src/client.ts index 9e94b0d8f6..19ba1ecb99 100644 --- a/packages/ai/src/client.ts +++ b/packages/ai/src/client.ts @@ -236,9 +236,6 @@ export { } from './activities/chat/agents/define-agent' export type { BoundActivities, - AgentRunHandle, - AgentStarter, - AgentStep, SubagentBinding, SubagentForward, } from './activities/chat/agents/bound' diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 5b70896670..0d2bbcd380 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -96,9 +96,6 @@ export { } from './activities/chat/agents/define-agent' export type { BoundActivities, - AgentRunHandle, - AgentStarter, - AgentStep, SubagentBinding, SubagentForward, } from './activities/chat/agents/bound' diff --git a/packages/ai/tests/agent-start.test.ts b/packages/ai/tests/agent-start.test.ts deleted file mode 100644 index 53a436e9fe..0000000000 --- a/packages/ai/tests/agent-start.test.ts +++ /dev/null @@ -1,77 +0,0 @@ -import { describe, expect, it } from 'vitest' -import { EventType } from '../src/types' -import { defineAgent } from '../src/activities/chat/agents/define-agent' -import { spawnAgentStream } from '../src/activities/chat/agents/spawn' -import type { AgentStarter } from '../src/activities/chat/agents/bound' -import type { SubagentRunInput } from '../src/activities/chat/agents/define-agent' -import type { StreamChunk } from '../src/types' - -const input: SubagentRunInput = { - input: undefined, - messages: [], - threadId: 't', - runId: 'r', - parentRunId: 'p', - subagentRunId: 's', -} - -async function collect(stream: AsyncIterable) { - const chunks: Array = [] - for await (const chunk of stream) chunks.push(chunk) - return chunks -} - -describe('ctx.agents of an agent run', () => { - it('throws a clear error when no host binds it', async () => { - const agent = defineAgent({ - name: 'lead', - description: 'Starts a helper', - run: async (ctx) => ctx.agents.start('helper'), - }) - - const chunks = await collect(spawnAgentStream(agent, input)) - - expect( - chunks.find((chunk) => chunk.type === EventType.SUBAGENT_ERROR), - ).toMatchObject({ - message: expect.stringContaining('ctx.agents.start needs a host'), - }) - }) - - it('starts agents through the agents of the host binding', async () => { - const started: Array = [] - const agents: AgentStarter = { - start: (agent, childInput, options) => { - started.push({ agent, input: childInput, options }) - return Object.assign(Promise.resolve('helped'), { - id: 'op-helper', - send: async () => undefined, - cancel: async () => undefined, - }) - }, - } - const agent = defineAgent({ - name: 'lead', - description: 'Starts a helper', - run: async (ctx) => { - const child = ctx.agents.start( - 'helper', - { topic: 'tea' }, - { wake: true }, - ) - return { id: child.id, result: await child } - }, - }) - - const chunks = await collect( - spawnAgentStream(agent, input, undefined, undefined, { agents }), - ) - - expect(started).toEqual([ - { agent: 'helper', input: { topic: 'tea' }, options: { wake: true } }, - ]) - expect( - chunks.find((chunk) => chunk.type === EventType.SUBAGENT_FINISHED), - ).toMatchObject({ result: { id: 'op-helper', result: 'helped' } }) - }) -}) diff --git a/packages/ai/tests/agent-step.test.ts b/packages/ai/tests/agent-step.test.ts deleted file mode 100644 index 9569ba434d..0000000000 --- a/packages/ai/tests/agent-step.test.ts +++ /dev/null @@ -1,65 +0,0 @@ -import { describe, expect, it } from 'vitest' -import { defineAgent } from '../src/activities/chat/agents/define-agent' -import { spawnAgentStream } from '../src/activities/chat/agents/spawn' -import type { AgentStep } from '../src/activities/chat/agents/bound' -import type { SubagentRunInput } from '../src/activities/chat/agents/define-agent' - -const input: SubagentRunInput = { - input: undefined, - messages: [], - threadId: 't', - runId: 'r', - parentRunId: 'p', - subagentRunId: 's', -} - -async function drain(stream: AsyncIterable) { - for await (const _chunk of stream) { - // Run the agent to its end. - } -} - -describe('ctx.step of an agent run', () => { - it('runs fn each time when no host gives steps', async () => { - let runs = 0 - const agent = defineAgent({ - name: 'counter', - description: 'Counts', - run: async (ctx) => ctx.step.do('count', () => (runs += 1)), - }) - - await drain(spawnAgentStream(agent, input)) - await drain(spawnAgentStream(agent, input)) - - expect(runs).toBe(2) - }) - - it('uses the steps of the host binding', async () => { - const names: Array = [] - const step: AgentStep = { - do: async (name: string) => { - names.push(name) - return 'stored' as T - }, - } - let ran = false - let value: unknown - const agent = defineAgent({ - name: 'charger', - description: 'Charges once', - run: async (ctx) => { - value = await ctx.step.do('charge', () => { - ran = true - return 'fresh' - }) - return value - }, - }) - - await drain(spawnAgentStream(agent, input, undefined, undefined, { step })) - - expect(names).toEqual(['charge']) - expect(ran).toBe(false) - expect(value).toBe('stored') - }) -}) diff --git a/packages/ai/tests/provider-keys.test.ts b/packages/ai/tests/provider-keys.test.ts index 0c0a6d4a99..f35faaed2e 100644 --- a/packages/ai/tests/provider-keys.test.ts +++ b/packages/ai/tests/provider-keys.test.ts @@ -9,6 +9,7 @@ import { import { getByokKey, keyedAdapterFromRequest } from '../src/byok/server' 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' @@ -260,3 +261,91 @@ describe('ctx.keys', () => { ) }) }) + +/** + * 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: Record, + nested?: { keys?: Record }, +) { + 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('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/testing/e2e/fixtures/subagents/basic.json b/testing/e2e/fixtures/subagents/basic.json index ac46395da5..df15dad377 100644 --- a/testing/e2e/fixtures/subagents/basic.json +++ b/testing/e2e/fixtures/subagents/basic.json @@ -175,6 +175,37 @@ "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/src/lib/providers.ts b/testing/e2e/src/lib/providers.ts index 6b2fe8636e..bac63a8b25 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 3bd8b62dc6..fe05f18890 100644 --- a/testing/e2e/src/lib/subagents-test.ts +++ b/testing/e2e/src/lib/subagents-test.ts @@ -25,6 +25,10 @@ export type SubagentScenario = | 'brief' | 'result' | 'limits' + | '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 = { @@ -34,6 +38,7 @@ export const SUBAGENT_PROMPTS: Record = { brief: '[subagent-brief] compare squid facts', result: '[subagent-result] price vendor a', limits: '[subagent-limits] research two topics', + keys: '[subagent-keys] write a draft', } export function isSubagentScenario(value: unknown): value is SubagentScenario { @@ -43,6 +48,7 @@ export function isSubagentScenario(value: unknown): value is SubagentScenario { value === 'tool' || value === 'brief' || value === 'result' || - value === 'limits' + value === 'limits' || + value === 'keys' ) } diff --git a/testing/e2e/src/routes/api.subagents-test.ts b/testing/e2e/src/routes/api.subagents-test.ts index 29c4224265..328efdf0ca 100644 --- a/testing/e2e/src/routes/api.subagents-test.ts +++ b/testing/e2e/src/routes/api.subagents-test.ts @@ -3,16 +3,20 @@ import { chat, chatParamsFromRequestBody, defineAgent, + keyedAdapter, maxIterations, toServerSentEventsResponse, } from '@tanstack/ai' import { z } from 'zod' import type { DefinedAgent, ReasoningRequest, Tool } from '@tanstack/ai' -import { createTextAdapter } from '@/lib/providers' +import { createOpenaiChat } from '@tanstack/ai-openai' +import { openaiByok } from '@tanstack/ai-openai/byok' +import { LLMOCK_DEFAULT_BASE, createTextAdapter } from '@/lib/providers' import { deleteLogs, isSubagentScenario, lookupFacts, + SUBAGENT_E2E_KEY, } from '@/lib/subagents-test' import type { SubagentScenario } from '@/lib/subagents-test' @@ -32,10 +36,21 @@ import type { SubagentScenario } from '@/lib/subagents-test' * 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 `writer` builds its adapter with + * `ctx.keys.adapter(...)` from `subagents.keys`. 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. */ +/** 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( scenario: SubagentScenario, aimockPort: number | undefined, @@ -117,6 +132,28 @@ 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: { openai: SUBAGENT_E2E_KEY } } + } return { agents: [child('researcher', [])] } } diff --git a/testing/e2e/tests/subagents.spec.ts b/testing/e2e/tests/subagents.spec.ts index 6e08f62473..471417d587 100644 --- a/testing/e2e/tests/subagents.spec.ts +++ b/testing/e2e/tests/subagents.spec.ts @@ -14,7 +14,14 @@ async function run(page: Page) { function open( page: Page, - scenario: 'route' | 'approval' | 'tool' | 'brief' | 'result' | 'limits', + scenario: + | 'route' + | 'approval' + | 'tool' + | 'brief' + | 'result' + | 'limits' + | 'keys', testId: string, aimockPort: number, ) { @@ -209,4 +216,41 @@ test.describe('subagents', () => { ), ).toBe(true) }) + + 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', + ) + }) }) From bb1bba706a3047f4f1b7872b497fb49e6e6dbcad Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 20:09:27 +0200 Subject: [PATCH 3/4] feat(ai): subagents.keys takes a key resolver function subagents.keys now also takes a function that gets the provider and returns its key, so any store plugs in. For BYOK headers: keys: (provider) => getByokKey(request, provider). A missing key still falls back to the parent agent's keys, then to env. ProviderKeySource is exported. --- .changeset/agent-limits-keys-step.md | 2 +- docs/advanced/byok.md | 13 +++++--- .../ai/src/activities/chat/agents/spawn.ts | 14 ++++---- packages/ai/src/byok/env-keys.ts | 33 ++++++++++++++----- packages/ai/src/index.ts | 1 + packages/ai/tests/provider-keys.test.ts | 24 ++++++++++++-- testing/e2e/src/routes/api.subagents-test.ts | 20 +++++++---- testing/e2e/src/routes/subagents-test.tsx | 9 ++++- 8 files changed, 86 insertions(+), 30 deletions(-) diff --git a/.changeset/agent-limits-keys-step.md b/.changeset/agent-limits-keys-step.md index 459351e760..8af37a4554 100644 --- a/.changeset/agent-limits-keys-step.md +++ b/.changeset/agent-limits-keys-step.md @@ -5,4 +5,4 @@ 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 keys by provider id, for example `{ openai: userKey }`. 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. +- **`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/docs/advanced/byok.md b/docs/advanced/byok.md index 6674d9b217..b4f399c499 100644 --- a/docs/advanced/byok.md +++ b/docs/advanced/byok.md @@ -353,7 +353,7 @@ export const writer = defineAgent({ }); ``` -2. In the route, read the user's key. Pass it on `subagents.keys`, by provider id: +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 { @@ -373,7 +373,10 @@ export async function POST(request: Request) { messages: params.messages, threadId: params.threadId, runId: params.runId, - subagents: { agents: [writer], keys: { openai: apiKey } }, + subagents: { + agents: [writer], + keys: (provider) => getByokKey(request, provider), + }, }); return toServerSentEventsResponse(stream); } @@ -381,9 +384,9 @@ export async function POST(request: Request) { How `ctx.keys` finds a key: -- It reads `subagents.keys` by provider id. The id of `openaiByok` is `openai`. -- If `subagents.keys` has no key for the provider, it reads the env var of the provider, here `OPENAI_API_KEY`. -- If both are empty, the call throws `Missing OpenAI API key. Set OPENAI_API_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. diff --git a/packages/ai/src/activities/chat/agents/spawn.ts b/packages/ai/src/activities/chat/agents/spawn.ts index fe2607aa08..e6309333ed 100644 --- a/packages/ai/src/activities/chat/agents/spawn.ts +++ b/packages/ai/src/activities/chat/agents/spawn.ts @@ -31,7 +31,7 @@ import { createBoundActivities } from './bound' import { SubagentBudget } from './limits' import type { SubagentBinding } from './bound' import type { SubagentLimits } from './limits' -import type { ProviderId } from '../../../byok/providers' +import type { ProviderKeySource } from '../../../byok/env-keys' import type { DefinedAgent, SubagentRunContext, @@ -99,12 +99,14 @@ export interface SubagentsBag< */ limits?: SubagentLimits /** - * Provider keys by provider id, for example `{ openai: userKey }`. Agents - * read them as `ctx.keys`. A missing or empty key falls back to the - * provider's `env` names. A child's `ctx.chat({ subagents })` passes them - * on. The keys it sets itself win for those providers. + * 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?: Readonly> + keys?: ProviderKeySource } /** What the children of one parent run left behind for the parent terminal. */ diff --git a/packages/ai/src/byok/env-keys.ts b/packages/ai/src/byok/env-keys.ts index 9c6f208e8a..ef068002f2 100644 --- a/packages/ai/src/byok/env-keys.ts +++ b/packages/ai/src/byok/env-keys.ts @@ -14,18 +14,35 @@ function missingKeyMessage(provider: ByokProvider | ProviderId) { } /** - * The `ctx.keys` of an agent: `keys` by provider id, 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. + * 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: Readonly> = {}, + keys: ProviderKeySource = {}, fallback?: ProviderKeys, ): ProviderKeys { - const get = async (provider: ByokProvider | ProviderId) => - keys[resolveProviderId(provider)] || - (fallback ? await fallback.get(provider) : envKey(provider)) || - null + 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) { diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 0d2bbcd380..a5ad53c842 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -111,6 +111,7 @@ export { type KeyedAdapterResult, type ProviderKeys, } from './byok/keyed' +export type { ProviderKeySource } from './byok/env-keys' export { subagentRoute, type SubagentRouteOptions, diff --git a/packages/ai/tests/provider-keys.test.ts b/packages/ai/tests/provider-keys.test.ts index f35faaed2e..4ec9eac8b1 100644 --- a/packages/ai/tests/provider-keys.test.ts +++ b/packages/ai/tests/provider-keys.test.ts @@ -15,6 +15,7 @@ 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' }) @@ -268,8 +269,8 @@ describe('ctx.keys', () => { * `ctx.chat({ subagents: { ...nested } })`. */ async function chatKeys( - keys: Record, - nested?: { keys?: Record }, + keys: ProviderKeySource, + nested?: { keys?: ProviderKeySource }, ) { const { agent: target, seen } = probe() const { adapter } = createMockAdapter({ iterations: [] }) @@ -339,6 +340,25 @@ describe('subagents.keys', () => { 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( diff --git a/testing/e2e/src/routes/api.subagents-test.ts b/testing/e2e/src/routes/api.subagents-test.ts index 328efdf0ca..fda6e11164 100644 --- a/testing/e2e/src/routes/api.subagents-test.ts +++ b/testing/e2e/src/routes/api.subagents-test.ts @@ -11,12 +11,13 @@ import { z } from 'zod' import type { 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, lookupFacts, - SUBAGENT_E2E_KEY, } from '@/lib/subagents-test' import type { SubagentScenario } from '@/lib/subagents-test' @@ -36,10 +37,11 @@ import type { SubagentScenario } from '@/lib/subagents-test' * 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 `writer` builds its adapter with - * `ctx.keys.adapter(...)` from `subagents.keys`. The child request carries - * that key. aimock hides `Authorization` in its journal, so the request - * also sends it as `x-e2e-authorization`. + * - `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. @@ -52,6 +54,7 @@ const echoAuthorization: typeof fetch = (input, init) => { } function subagentsFor( + request: Request, scenario: SubagentScenario, aimockPort: number | undefined, testId: string | undefined, @@ -152,7 +155,10 @@ function subagentsFor( }), }) const agents: Array = [writer] - return { agents, keys: { openai: SUBAGENT_E2E_KEY } } + return { + agents, + keys: (provider: ByokProvider | string) => getByokKey(request, provider), + } } return { agents: [child('researcher', [])] } } @@ -187,7 +193,7 @@ 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), agentLoopStrategy: maxIterations(4), abortController, }) diff --git a/testing/e2e/src/routes/subagents-test.tsx b/testing/e2e/src/routes/subagents-test.tsx index 560aa680ac..5d8e63c1e8 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, }) From 122b22d0733ec63c66ba3270192f727391ac9cc0 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 20:29:33 +0200 Subject: [PATCH 4/4] feat(ai, ai-client, ai-persistence): continue a child by sessionId, and a typed part.subagent.input (#1731) * feat(ai, ai-persistence): one subagent tool with sessionId continue subagents.tool 'single' gives the model one subagent tool that names the agent, takes input or a prompt, continues a stored child by sessionId through the new LoadChild capability, and starts a background child through binding.start. withPersistence provides LoadChild and stores the task and owner thread of each child. * feat(ai, ai-persistence)!: every agent tool takes input fields, drop tool: 'single' Each per-agent tool now takes { input } (agent with inputSchema) or { prompt } (agent without), plus sessionId when a middleware provides LoadChildCapability and background when the host sets binding.start. This replaces the separate `subagents: { tool: 'single' }` mode, so continue and background work on the normal agent tools. - The agent schema sits under `input`. Its refs are rewritten and its definitions move to the top, so recursive schemas still resolve. - A result carries subagentRunId only when the tool takes sessionId. - ai-persistence reads prompt/input from any agent tool call. - Breaking: the tool-call shape changes. New migration page docs/migration/subagent-tool-input.md. * feat(ai, ai-client, ai-persistence): flat agent tool fields and a typed part.subagent.input Agent tools take the agent's own fields again, so the tool-call shape does not change for existing apps. sessionId and background sit next to them, only when the run uses them (a LoadChildCapability middleware, or binding.start). Only then are those names reserved in an inputSchema. - With reserved fields on, a Standard Schema wrapper checks the agent fields in the tool loop, before middleware, as for any tool, and passes the reserved fields through. Without them, the tool uses the agent schema as is. - SUBAGENT_STARTED carries the checked input in metadata.tanstack.subagentInput. ChatClient puts it on the card handle as `input`, typed from the inputSchema of the agent declaration given to useChat({ subagents }). - ai-persistence uses that card input to tell an input agent from a prompt agent when it stores the task and follow-up prompts. - Drop the migration page. The brief docs and the example read part.subagent.input. --- .changeset/subagent-tool-input.md | 11 + docs/chat/subagents.md | 122 ++++- .../ts-react-chat/src/lib/researcher-agent.ts | 17 + .../src/routes/api.subagent-brief.ts | 13 +- .../src/routes/subagent-brief.tsx | 26 +- packages/ai-client/src/chat-client.ts | 13 +- packages/ai-client/src/types.ts | 24 +- .../tests/chat-client-subagents.test.ts | 31 +- packages/ai-persistence/src/middleware.ts | 16 + packages/ai-persistence/src/subagent-runs.ts | 202 ++++++- .../tests/subagent-persistence-cards.test.ts | 363 +++++++++++++ .../tests/use-chat-subagent-types.test.ts | 47 ++ .../ai/src/activities/chat/agents/bound.ts | 15 + .../activities/chat/agents/define-agent.ts | 7 +- .../ai/src/activities/chat/agents/spawn.ts | 334 +++++++++++- packages/ai/src/activities/chat/index.ts | 19 + .../activities/chat/middleware/load-child.ts | 23 + .../src/activities/chat/tools/tool-calls.ts | 12 +- packages/ai/src/client.ts | 1 + packages/ai/src/index.ts | 8 + packages/ai/src/utilities/subagent-wire.ts | 12 + packages/ai/tests/define-agent.test.ts | 6 +- packages/ai/tests/subagent-tool-input.test.ts | 512 ++++++++++++++++++ testing/e2e/fixtures/subagents/session.json | 59 ++ testing/e2e/src/lib/subagents-test.ts | 3 + testing/e2e/src/routes/api.subagents-test.ts | 56 +- testing/e2e/src/routes/subagents-test.tsx | 3 + testing/e2e/tests/subagents.spec.ts | 31 ++ 28 files changed, 1889 insertions(+), 97 deletions(-) create mode 100644 .changeset/subagent-tool-input.md create mode 100644 examples/ts-react-chat/src/lib/researcher-agent.ts create mode 100644 packages/ai/src/activities/chat/middleware/load-child.ts create mode 100644 packages/ai/tests/subagent-tool-input.test.ts create mode 100644 testing/e2e/fixtures/subagents/session.json diff --git a/.changeset/subagent-tool-input.md b/.changeset/subagent-tool-input.md new file mode 100644 index 0000000000..a45cbc8dda --- /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/chat/subagents.md b/docs/chat/subagents.md index d65931a860..c2b6af7682 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: @@ -393,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 @@ -661,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/examples/ts-react-chat/src/lib/researcher-agent.ts b/examples/ts-react-chat/src/lib/researcher-agent.ts new file mode 100644 index 0000000000..9529e92f2c --- /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 0f229eb84c..3175a98cf6 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 704e83a75b..86dd2c8a13 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 a94a3d3b0e..4a3c19c544 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 e92d914d88..5cbca7c5f3 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 cb17450e3f..a4c3a0a175 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 77b0a90b39..d71982be7b 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 fcbac82553..e6f3d5be51 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 0809cee0b4..0671294ef5 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 4e404c9f34..5858926701 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 1655fd993e..0b4008797c 100644 --- a/packages/ai/src/activities/chat/agents/bound.ts +++ b/packages/ai/src/activities/chat/agents/bound.ts @@ -69,6 +69,21 @@ export interface SubagentBinding { * 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 }> } /** diff --git a/packages/ai/src/activities/chat/agents/define-agent.ts b/packages/ai/src/activities/chat/agents/define-agent.ts index 7a2417749f..8edf16529f 100644 --- a/packages/ai/src/activities/chat/agents/define-agent.ts +++ b/packages/ai/src/activities/chat/agents/define-agent.ts @@ -121,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/spawn.ts b/packages/ai/src/activities/chat/agents/spawn.ts index e6309333ed..4bd7354591 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,6 +37,7 @@ import type { SubagentStartedEvent, TokenUsage, Tool, + ToolExecutionContext, UIMessage, } from '../../../types' import { envProviderKeys } from '../../../byok/env-keys' @@ -109,6 +123,62 @@ export interface SubagentsBag< 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. */ export interface SubagentSink { interrupts: Array @@ -127,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. */ @@ -289,7 +366,9 @@ function openAgentStream( ) { 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) { @@ -297,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, } @@ -309,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), @@ -585,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, @@ -593,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) @@ -844,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( @@ -863,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, } } @@ -883,6 +1115,10 @@ 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. @@ -1016,24 +1252,88 @@ export function createSyntheticSubagentTools( ...(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. - // An agent without `inputSchema` still gets `{}` from the model. Its - // `ctx.input` stays undefined. - execute: (input: unknown, context?: unknown) => - runChild( - { - name: agent.name, - ...(agent.inputSchema !== undefined && { input }), - }, - context, - ), + 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 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( + { + agent: agent.name, + ...fields, + ...(context?.toolCallId !== undefined && { + parentToolCallId: context.toolCallId, + }), + }, + { wake: true }, + ) + return { + subagentRunId, + text: '', + result: { status: 'started' }, + keepRunId: true, + } 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 ac7ef6d1c2..691eca150f 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -64,6 +64,7 @@ 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, @@ -114,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' @@ -5651,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 @@ -5940,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 && { @@ -6323,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 0000000000..bd806011bf --- /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 82983b4a46..2b889b2934 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/client.ts b/packages/ai/src/client.ts index 19ba1ecb99..e3bcbf4c0e 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 a5ad53c842..45985eb4f7 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -383,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). @@ -575,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 db2a79e042..1df463d5e4 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 ba676a0549..b8e2a87d75 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/subagent-tool-input.test.ts b/packages/ai/tests/subagent-tool-input.test.ts new file mode 100644 index 0000000000..ac2f30bacb --- /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/session.json b/testing/e2e/fixtures/subagents/session.json new file mode 100644 index 0000000000..17e8241d82 --- /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/subagents-test.ts b/testing/e2e/src/lib/subagents-test.ts index fe05f18890..2748d182ff 100644 --- a/testing/e2e/src/lib/subagents-test.ts +++ b/testing/e2e/src/lib/subagents-test.ts @@ -25,6 +25,7 @@ export type SubagentScenario = | 'brief' | 'result' | 'limits' + | 'session' | 'keys' /** The key the `keys` scenario passes in `subagents.keys`. */ @@ -38,6 +39,7 @@ export const SUBAGENT_PROMPTS: Record = { 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', } @@ -49,6 +51,7 @@ export function isSubagentScenario(value: unknown): value is SubagentScenario { value === 'brief' || 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 fda6e11164..0e17b63a30 100644 --- a/testing/e2e/src/routes/api.subagents-test.ts +++ b/testing/e2e/src/routes/api.subagents-test.ts @@ -7,8 +7,14 @@ import { 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 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' @@ -45,6 +51,9 @@ import type { SubagentScenario } from '@/lib/subagents-test' * - `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) => { @@ -163,6 +172,42 @@ function subagentsFor( 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: { @@ -194,6 +239,15 @@ export const Route = createFileRoute('/api/subagents-test')({ ...(params.parentRunId ? { parentRunId: params.parentRunId } : {}), ...(params.resume ? { resume: params.resume } : {}), 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 5d8e63c1e8..717b07a20f 100644 --- a/testing/e2e/src/routes/subagents-test.tsx +++ b/testing/e2e/src/routes/subagents-test.tsx @@ -71,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 471417d587..90f3709e91 100644 --- a/testing/e2e/tests/subagents.spec.ts +++ b/testing/e2e/tests/subagents.spec.ts @@ -21,6 +21,7 @@ function open( | 'brief' | 'result' | 'limits' + | 'session' | 'keys', testId: string, aimockPort: number, @@ -161,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 ({ @@ -217,6 +222,32 @@ test.describe('subagents', () => { ).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,