From 909126af0af434949b1d19f04ec393fb80733c4b Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 25 Sep 2026 17:44:12 +0200 Subject: [PATCH 001/286] feat(ai): add inputSchema to defineAgent The parent model can now write the child's input when it calls the agent's tool, for example a short brief. run() reads the checked input as a typed ctx.input. A bad input goes back to the model as a tool error, and the child does not start. A router cannot write input, so chat() throws when subagents.router meets an agent with inputSchema. Refs #1482 --- .changeset/define-agent-input-schema.md | 5 + .../activities/chat/agents/define-agent.ts | 23 ++- .../ai/src/activities/chat/agents/spawn.ts | 32 +-- packages/ai/src/activities/chat/index.ts | 10 + packages/ai/tests/define-agent.test.ts | 190 +++++++++++++++++- packages/ai/tests/subagent-interrupts.test.ts | 135 ++++++++++--- 6 files changed, 347 insertions(+), 48 deletions(-) create mode 100644 .changeset/define-agent-input-schema.md diff --git a/.changeset/define-agent-input-schema.md b/.changeset/define-agent-input-schema.md new file mode 100644 index 0000000000..c46e444124 --- /dev/null +++ b/.changeset/define-agent-input-schema.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai': minor +--- + +Add an optional `inputSchema` to `defineAgent`. The agent's tool shows the schema to the model, so the parent model writes the child's input, such as a short brief. `run` gets the checked input as a typed `ctx.input`. A bad input goes back to the model as a tool error, and the child does not start. `inputSchema` needs tool mode: `chat()` throws when `subagents.router` meets an agent that has one. Agents without `inputSchema` do not change. diff --git a/packages/ai/src/activities/chat/agents/define-agent.ts b/packages/ai/src/activities/chat/agents/define-agent.ts index 5137d0a148..e9822f64b8 100644 --- a/packages/ai/src/activities/chat/agents/define-agent.ts +++ b/packages/ai/src/activities/chat/agents/define-agent.ts @@ -2,6 +2,7 @@ import type { SubagentInfo as AGUISubagentInfo } from '@ag-ui/core' import type { InterruptDefinition } from '../../../interrupt-definition' import type { AnyTool, + InferSchemaType, ModelMessage, RunAgentResumeItem, SchemaInput, @@ -12,8 +13,16 @@ import type { AnyClientTool } from '../tools/tool-definition' /** * Context the library passes into {@link defineAgent} `run`. + * `TInput` is the agent's `inputSchema`. */ -export interface SubagentRunContext { +export interface SubagentRunContext< + TInput extends SchemaInput | undefined = any, +> { + /** + * The input the parent model wrote for this child, checked against + * `inputSchema`. `undefined` when the agent has no `inputSchema`. + */ + input: TInput extends SchemaInput ? InferSchemaType : undefined messages: Array abortSignal?: AbortSignal threadId: string @@ -53,13 +62,20 @@ export interface DefinedAgent< TSchema extends SchemaInput | undefined = SchemaInput | undefined, TInterrupts extends ReadonlyArray> = ReadonlyArray>, + TInput extends SchemaInput | undefined = any, > extends AGUISubagentInfo { name: TName /** Required here: the router and the synthetic tool both read it. */ description: string run: ( - ctx: SubagentRunContext, + ctx: SubagentRunContext, ) => AsyncIterable | Promise> + /** + * 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`. + */ + inputSchema?: TInput tools?: TTools interrupts?: TInterrupts outputSchema?: TSchema @@ -104,7 +120,8 @@ export function defineAgent< const TInterrupts extends ReadonlyArray< InterruptDefinition > = readonly [], ->(agent: DefinedAgent) { + TInput extends SchemaInput | undefined = undefined, +>(agent: DefinedAgent) { if (agent.name.trim() === '') { throw new Error('defineAgent requires a non-empty name') } diff --git a/packages/ai/src/activities/chat/agents/spawn.ts b/packages/ai/src/activities/chat/agents/spawn.ts index 9f16e34363..18e712007a 100644 --- a/packages/ai/src/activities/chat/agents/spawn.ts +++ b/packages/ai/src/activities/chat/agents/spawn.ts @@ -93,6 +93,8 @@ export function createSubagentSink(): SubagentSink { /** One child to start, or a suspended child to continue. */ export interface SpawnEntry { name: string + /** The checked tool input for an agent with `inputSchema`. */ + input?: unknown resume?: { subagentRunId: string /** The child's own messages from the interrupted run. */ @@ -248,6 +250,7 @@ function openAgentStream( return spawnAgentStream( agent, { + input: entry.input, messages: resumed?.messages ?? ctx.messages, ...(ctx.abortSignal ? { abortSignal: ctx.abortSignal } : {}), threadId: childThreadId(bag.sandbox, ctx.threadId, entry.name), @@ -721,8 +724,10 @@ export function createSyntheticSubagentTools( return bag.agents.map((agent) => ({ name: agent.name, description: agent.description, + ...(agent.inputSchema !== undefined && { inputSchema: agent.inputSchema }), [SUBAGENT_TOOL]: true, - execute: async (_input: unknown, context?: unknown) => { + // The tool loop checks `input` against `inputSchema` before this runs. + execute: async (input: unknown, context?: unknown) => { const toolContext = context as | { toolCallId?: string @@ -736,17 +741,20 @@ export function createSyntheticSubagentTools( child.parentToolCallId !== undefined && child.parentToolCallId === toolCallId, ) - const entry: SpawnEntry = suspended - ? { - name: agent.name, - resume: { - subagentRunId: suspended.subagentRunId, - messages: suspended.messages, - entries: suspended.resume, - text: suspended.text, - }, - } - : { name: agent.name } + // 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 ?? '' diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 940390490d..29b8fc599f 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -5319,6 +5319,16 @@ async function* runRoutedSubagents( yield* runChatEngine(options, engineRef) return } + // ponytail: a router cannot write input yet. Add it to the router pick if + // someone needs it. + const withInput = bag.agents.find( + (agent: DefinedAgent) => agent.inputSchema !== undefined, + ) + if (withInput) { + throw new Error( + `Subagent "${withInput.name}" has an inputSchema. inputSchema needs tool mode: remove subagents.router so the model writes the input.`, + ) + } const threadId = options.threadId ?? `thread-${Date.now()}` const runId = options.runId ?? `run-${Date.now()}` const messages = options.messages ?? [] diff --git a/packages/ai/tests/define-agent.test.ts b/packages/ai/tests/define-agent.test.ts index 2311652925..ba676a0549 100644 --- a/packages/ai/tests/define-agent.test.ts +++ b/packages/ai/tests/define-agent.test.ts @@ -1,4 +1,5 @@ -import { describe, expect, it } from 'vitest' +import { describe, expect, expectTypeOf, it } from 'vitest' +import { z } from 'zod' import { defineAgent, type DefinedAgent, @@ -710,6 +711,193 @@ describe('chat({ subagents }) synthetic tools', () => { }) }) +describe('defineAgent inputSchema', () => { + const briefSchema = z.object({ task: z.string() }) + + /** A parent model that calls `researcher` once per args string, then stops. */ + function parentCalling(...args: Array) { + return createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ...args.flatMap((json, index) => [ + ev.toolStart(`call_${index}`, 'researcher'), + ev.toolArgs(`call_${index}`, json), + ]), + ev.runFinished('tool_calls'), + ], + [ev.runStarted(), ev.runFinished('stop')], + ], + }) + } + + /** A researcher that records every `ctx.input` it gets. */ + function briefResearcher() { + const inputs: Array<{ task: string }> = [] + const agent = defineAgent({ + name: 'researcher', + description: 'Looks up facts', + inputSchema: briefSchema, + run: async function* (ctx) { + expectTypeOf(ctx.input).toEqualTypeOf<{ task: string }>() + inputs.push(ctx.input) + }, + }) + return { agent, inputs } + } + + it('shows the schema to the model and passes the input to run', async () => { + const { agent, inputs } = briefResearcher() + const { adapter, calls } = parentCalling('{"task":"Price Vendor A"}') + + await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Compare pricing' }], + subagents: { agents: [agent] }, + }) as AsyncIterable, + ) + + const tool = calls[0]?.tools?.find((entry) => entry.name === 'researcher') + expect(tool?.inputSchema).toMatchObject({ + type: 'object', + properties: { task: { type: 'string' } }, + }) + expect(inputs).toEqual([{ task: 'Price Vendor A' }]) + }) + + it('gives each call to the same agent its own input', async () => { + const { agent, inputs } = briefResearcher() + const { adapter } = parentCalling( + '{"task":"Price Vendor A"}', + '{"task":"Price Vendor B"}', + ) + + await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Compare pricing' }], + subagents: { agents: [agent] }, + }) as AsyncIterable, + ) + + expect(inputs).toHaveLength(2) + expect(inputs).toEqual( + expect.arrayContaining([ + { task: 'Price Vendor A' }, + { task: 'Price Vendor B' }, + ]), + ) + }) + + it('returns a bad input to the model and does not start the child', async () => { + const { agent, inputs } = briefResearcher() + const { adapter } = parentCalling('{}') + + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Compare pricing' }], + subagents: { agents: [agent] }, + }) as AsyncIterable, + ) + + const result = chunks.find( + (chunk) => + chunk.type === EventType.TOOL_CALL_RESULT && + chunk.toolCallId === 'call_0', + ) + expect(result).toMatchObject({ + content: expect.stringContaining( + 'Input validation failed for tool researcher', + ), + }) + expect(chunks.some((chunk) => chunk.type === 'SUBAGENT_STARTED')).toBe( + false, + ) + expect(inputs).toEqual([]) + }) + + it('passes the parsed args through a plain JSON Schema', async () => { + const inputs: Array = [] + const agent = defineAgent({ + name: 'researcher', + description: 'Looks up facts', + inputSchema: { + type: 'object', + properties: { task: { type: 'string' } }, + required: ['task'], + }, + run: async function* (ctx) { + inputs.push(ctx.input) + }, + }) + const { adapter, calls } = parentCalling('{"task":"Price Vendor A"}') + + await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Compare pricing' }], + subagents: { agents: [agent] }, + }) as AsyncIterable, + ) + + const tool = calls[0]?.tools?.find((entry) => entry.name === 'researcher') + expect(tool?.inputSchema).toMatchObject({ + properties: { task: { type: 'string' } }, + }) + expect(inputs).toEqual([{ task: 'Price Vendor A' }]) + }) + + it('keeps ctx.input undefined for an agent without inputSchema', async () => { + const inputs: Array = [] + const agent = defineAgent({ + name: 'researcher', + description: 'Looks up facts', + run: async function* (ctx) { + expectTypeOf(ctx.input).toEqualTypeOf() + inputs.push(ctx.input) + }, + }) + const { adapter, calls } = parentCalling('{}') + + await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Research this' }], + subagents: { agents: [agent] }, + }) as AsyncIterable, + ) + + 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({}) + expect(inputs).toEqual([undefined]) + }) + + it('throws when a router meets an agent with inputSchema', async () => { + const { agent } = briefResearcher() + let routerCalls = 0 + + await expect( + collectChunks( + chat({ + adapter: parentAdapter().adapter, + messages: [{ role: 'user', content: 'Compare pricing' }], + subagents: { + agents: [agent], + router: () => { + routerCalls += 1 + return 'main' + }, + }, + }) as AsyncIterable, + ), + ).rejects.toThrow('Subagent "researcher" has an inputSchema') + expect(routerCalls).toBe(0) + }) +}) + describe('subagent guards', () => { it('rejects the reserved name main', () => { expect(() => namedAgent('main')).toThrow("cannot use the name 'main'") diff --git a/packages/ai/tests/subagent-interrupts.test.ts b/packages/ai/tests/subagent-interrupts.test.ts index e86b3d9e01..f7eb5408ee 100644 --- a/packages/ai/tests/subagent-interrupts.test.ts +++ b/packages/ai/tests/subagent-interrupts.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' import { chat } from '../src/activities/chat' import { defineAgent } from '../src/activities/chat/agents/define-agent' import { StreamProcessor } from '../src/activities/chat/stream/processor' @@ -22,8 +23,8 @@ const user: UIMessage = { parts: [{ type: 'text', content: 'Clean up' }], } -/** A child that asks approval to delete a file, then reports. */ -function cleanerAgent({ needsApproval = true } = {}) { +/** A child model that asks to delete a file, then reports. */ +function cleanerChild({ needsApproval = true } = {}) { const execute = vi.fn().mockReturnValue({ deleted: true }) const { adapter } = createMockAdapter({ iterations: [ @@ -51,27 +52,60 @@ function cleanerAgent({ needsApproval = true } = {}) { ], ], }) + const tools = [{ ...serverTool('deleteFile', execute), needsApproval }] + const run = (ctx: SubagentRunContext) => + chat({ + adapter, + messages: ctx.messages, + threadId: ctx.threadId, + runId: ctx.runId, + parentRunId: ctx.parentRunId, + subagentRunId: ctx.subagentRunId, + resume: ctx.resume, + tools, + }) + return { run, execute } +} + +/** A child that asks approval to delete a file, then reports. */ +function cleanerAgent({ needsApproval = true } = {}) { + const { run, execute } = cleanerChild({ needsApproval }) const contexts: Array = [] const agent = defineAgent({ name: 'cleaner', description: 'Deletes files', run: (ctx) => { contexts.push(ctx) - return chat({ - adapter, - messages: ctx.messages, - threadId: ctx.threadId, - runId: ctx.runId, - parentRunId: ctx.parentRunId, - subagentRunId: ctx.subagentRunId, - resume: ctx.resume, - tools: [{ ...serverTool('deleteFile', execute), needsApproval }], - }) + return run(ctx) }, }) return { agent, execute, contexts } } +/** A parent model that calls `cleaner` with `args`, then wraps up. */ +function toolParent(args: string) { + return createMockAdapter({ + iterations: [ + [ + ev.runStarted('p1'), + ev.textStart('parent-note'), + ev.textContent('Checking the disk first', 'parent-note'), + ev.textEnd('parent-note'), + ev.toolStart('call_p', 'cleaner'), + ev.toolArgs('call_p', args), + ev.runFinished('tool_calls', 'p1'), + ], + [ + ev.runStarted('p2'), + ev.textStart('parent-done'), + ev.textContent('All clean', 'parent-done'), + ev.textEnd('parent-done'), + ev.runFinished('stop', 'p2'), + ], + ], + }).adapter +} + function cardOf(messages: ReadonlyArray): SubagentPart { for (const message of messages) { for (const part of message.parts) { @@ -386,26 +420,7 @@ describe('subagent interrupts', () => { it('suspends a child that a tool call started, then resumes it', async () => { const { agent, execute, contexts } = cleanerAgent() - const { adapter: parent } = createMockAdapter({ - iterations: [ - [ - ev.runStarted('p1'), - ev.textStart('parent-note'), - ev.textContent('Checking the disk first', 'parent-note'), - ev.textEnd('parent-note'), - ev.toolStart('call_p', 'cleaner'), - ev.toolArgs('call_p', '{}'), - ev.runFinished('tool_calls', 'p1'), - ], - [ - ev.runStarted('p2'), - ev.textStart('parent-done'), - ev.textContent('All clean', 'parent-done'), - ev.textEnd('parent-done'), - ev.runFinished('stop', 'p2'), - ], - ], - }) + const parent = toolParent('{}') const subagents = { agents: [agent] } const first = await collectChunks( @@ -477,6 +492,62 @@ describe('subagent interrupts', () => { expect(second.at(-1)).toMatchObject({ type: EventType.RUN_FINISHED }) expect(second.at(-1)).not.toHaveProperty('outcome.type', 'interrupt') }) + + it('gives a resumed child the same input the model wrote', async () => { + const { run, execute } = cleanerChild() + const inputs: Array<{ folder: string }> = [] + const agent = defineAgent({ + name: 'cleaner', + description: 'Deletes files', + inputSchema: z.object({ folder: z.string() }), + run: (ctx) => { + inputs.push(ctx.input) + return run(ctx) + }, + }) + const parent = toolParent('{"folder":"old"}') + const subagents = { agents: [agent] } + + const first = await collectChunks( + chat({ + adapter: parent, + threadId: 't', + runId: 'run-1', + messages: [user], + subagents, + }) as AsyncIterable, + ) + const processor = new StreamProcessor({ initialMessages: [user] }) + replay(processor, first) + // The input stays on the parent's tool call, next to the card. + expect( + processor.getMessages().flatMap((message) => message.parts), + ).toContainEqual( + expect.objectContaining({ + type: 'tool-call', + id: 'call_p', + arguments: '{"folder":"old"}', + }), + ) + + processor.addToolApprovalResponse('approval_call_c', true) + await collectChunks( + chat({ + adapter: parent, + threadId: 't', + runId: 'run-2', + parentRunId: 'run-1', + messages: requestMessages(processor.getMessages()), + resume: [ + { interruptId: 'approval_call_c', status: 'resolved', payload: true }, + ], + subagents, + }) as AsyncIterable, + ) + + expect(inputs).toEqual([{ folder: 'old' }, { folder: 'old' }]) + expect(execute).toHaveBeenCalledTimes(1) + }) }) describe('subagent cards', () => { From 92268fae1d9fba7d39570cc97c6f1760767d8e9e Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 25 Sep 2026 17:44:13 +0200 Subject: [PATCH 002/286] test(e2e): cover a subagent brief written by the parent model The child fixture matches only the brief text, so the test fails if the child gets the parent transcript instead. --- testing/e2e/fixtures/subagents/basic.json | 31 ++++++++++++++++++++ testing/e2e/src/lib/subagents-test.ts | 10 +++++-- testing/e2e/src/routes/api.subagents-test.ts | 25 +++++++++++++++- testing/e2e/tests/subagents.spec.ts | 23 ++++++++++++++- 4 files changed, 85 insertions(+), 4 deletions(-) diff --git a/testing/e2e/fixtures/subagents/basic.json b/testing/e2e/fixtures/subagents/basic.json index 610e6dd093..fd9d5b682f 100644 --- a/testing/e2e/fixtures/subagents/basic.json +++ b/testing/e2e/fixtures/subagents/basic.json @@ -78,6 +78,37 @@ "response": { "content": "Research is done." } + }, + { + "match": { + "userMessage": "[subagent-brief] compare squid facts", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "researcher", + "arguments": "{\"task\":\"[subagent-brief-task] List two squid facts.\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[subagent-brief-task] List two squid facts." + }, + "response": { + "content": "Squids have three hearts and blue blood." + } + }, + { + "match": { + "userMessage": "[subagent-brief] compare squid facts", + "sequenceIndex": 1 + }, + "response": { + "content": "The comparison is ready." + } } ] } diff --git a/testing/e2e/src/lib/subagents-test.ts b/testing/e2e/src/lib/subagents-test.ts index d6ff957740..b7dfad6c23 100644 --- a/testing/e2e/src/lib/subagents-test.ts +++ b/testing/e2e/src/lib/subagents-test.ts @@ -18,15 +18,21 @@ export const lookupFacts = toolDefinition({ inputSchema: z.object({ topic: z.string() }), }) -export type SubagentScenario = 'route' | 'approval' | 'tool' +export type SubagentScenario = 'route' | 'approval' | 'tool' | 'brief' /** The user message for each scenario. It matches `fixtures/subagents`. */ export const SUBAGENT_PROMPTS: Record = { route: '[subagent-route] research squids', approval: '[subagent-approval] clean up the old logs', tool: '[subagent-tool] research squids', + brief: '[subagent-brief] compare squid facts', } export function isSubagentScenario(value: unknown): value is SubagentScenario { - return value === 'route' || value === 'approval' || value === 'tool' + return ( + value === 'route' || + value === 'approval' || + value === 'tool' || + value === 'brief' + ) } diff --git a/testing/e2e/src/routes/api.subagents-test.ts b/testing/e2e/src/routes/api.subagents-test.ts index 5abcfe4e2e..561fa9185f 100644 --- a/testing/e2e/src/routes/api.subagents-test.ts +++ b/testing/e2e/src/routes/api.subagents-test.ts @@ -6,7 +6,8 @@ import { maxIterations, toServerSentEventsResponse, } from '@tanstack/ai' -import type { Tool } from '@tanstack/ai' +import { z } from 'zod' +import type { DefinedAgent, Tool } from '@tanstack/ai' import { createTextAdapter } from '@/lib/providers' import { deleteLogs, @@ -27,6 +28,8 @@ import type { SubagentScenario } from '@/lib/subagents-test' * the same child. * - `tool`: no router. The parent model calls the `researcher` tool, the * 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. */ function subagentsFor( scenario: SubagentScenario, @@ -72,6 +75,26 @@ function subagentsFor( ]) return { agents: [cleaner], router: () => 'cleaner' } } + if (scenario === 'brief') { + const researcher = defineAgent({ + name: 'researcher', + description: 'Researches one focused task', + inputSchema: z.object({ task: z.string() }), + run: (ctx) => + chat({ + ...createTextAdapter('openai', undefined, aimockPort, testId), + messages: [{ role: 'user', content: ctx.input.task }], + threadId: ctx.threadId, + runId: ctx.runId, + parentRunId: ctx.parentRunId, + subagentRunId: ctx.subagentRunId, + resume: ctx.resume, + }), + }) + // Each branch returns its own bag, so give them one agent type. + const agents: Array = [researcher] + return { agents } + } return { agents: [child('researcher', [])] } } diff --git a/testing/e2e/tests/subagents.spec.ts b/testing/e2e/tests/subagents.spec.ts index 8faf7d8968..20ae0d3f44 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', + scenario: 'route' | 'approval' | 'tool' | 'brief', testId: string, aimockPort: number, ) { @@ -111,4 +111,25 @@ test.describe('subagents', () => { 'text:Squids have three hearts.', ) }) + + test('the parent model writes the brief the child runs on', async ({ + page, + testId, + aimockPort, + }) => { + await open(page, 'brief', testId, aimockPort) + await run(page) + + await expect(page.getByTestId('parent-text')).toHaveText( + 'The comparison is ready.', + ) + await expect(page.getByTestId('card-status-researcher')).toHaveText( + 'finished', + ) + // The child's fixture matches only the brief. A child that got the + // transcript instead would get no reply. + await expect(page.getByTestId('card-researcher')).toContainText( + 'text:Squids have three hearts and blue blood.', + ) + }) }) From 650a5f871d7f3a66e73fd884bd7b672d24d8970c Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 25 Sep 2026 17:44:13 +0200 Subject: [PATCH 003/286] docs: document subagent briefs with inputSchema Add 'Let the model write the brief' and 'Show the brief on a card' to the subagents page. --- docs/chat/subagents.md | 85 ++++++++++++++++++++++++++++++++++++++++++ docs/config.json | 2 +- 2 files changed, 86 insertions(+), 1 deletion(-) diff --git a/docs/chat/subagents.md b/docs/chat/subagents.md index a0c02f973f..de2f4698a3 100644 --- a/docs/chat/subagents.md +++ b/docs/chat/subagents.md @@ -156,6 +156,56 @@ 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. +## Let the model write the brief + +In a long thread, the child gets the full conversation at the tool call, with every earlier tool result. Each parallel child gets its own copy. The child must also find its task in that conversation. + +Add `inputSchema` to the agent. The main model then writes the child's input when it calls the tool. `run` reads that input as `ctx.input`: + +```ts +import { chat, defineAgent } from '@tanstack/ai' +import { openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' + +const messages = [ + { role: 'user' as const, content: 'Compare pricing for the vendors we picked' }, +] + +const researcher = defineAgent({ + name: 'researcher', + description: 'Researches one focused question and returns sourced findings', + inputSchema: z.object({ + task: z + .string() + .describe('What to find and what to return. The researcher sees only this text.'), + }), + run: (ctx) => + chat({ + adapter: openaiText('gpt-5.6'), + messages: [{ role: 'user', content: ctx.input.task }], + threadId: ctx.threadId, + runId: ctx.runId, + parentRunId: ctx.parentRunId, + subagentRunId: ctx.subagentRunId, + resume: ctx.resume, + }), +}) + +const stream = chat({ + adapter: openaiText('gpt-5.6'), + messages, + subagents: { agents: [researcher] }, +}) +``` + +- `ctx.input` has the type of the schema. Here it is `{ task: string }`. +- If the input does not match the schema, the model gets a tool error and can call the tool again. The child does not start. +- `ctx.messages` still holds the parent conversation. Add parts of it when the child needs more context. +- A resumed child gets the same `ctx.input` as the first run. +- `inputSchema` needs tool mode. If `subagents.router` is set and an agent has `inputSchema`, `chat()` throws. + +To show the brief in the UI, see [Show the brief on a card](#show-the-brief-on-a-card). + ## Strategy - `exclusive` (default): the chosen child owns the turn. Main does not answer after it. @@ -497,6 +547,41 @@ export function ChatScreen() { When `run` imports server code, do not import the agent into the browser. Declare the agent once in a shared file, with its `name`, `description`, and tool definitions from `toolDefinition`. The server passes `defineAgent({ ...researcher, run })` to `chat()`. The client passes the declaration. +### 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: + +```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 +} + +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)}

+ }) +} +``` + +A child that a router started has no `parentToolCallId`, so it has no brief. + ### Style one child's parts The researcher's reasoning and tool calls use the root widgets by default. To make them look different on the researcher card only, pass widgets to that card's `Parts`: diff --git a/docs/config.json b/docs/config.json index 0d4770f15b..3c6e1f6f68 100644 --- a/docs/config.json +++ b/docs/config.json @@ -155,7 +155,7 @@ "label": "Subagents", "to": "chat/subagents", "addedAt": "2026-09-21", - "updatedAt": "2026-09-24" + "updatedAt": "2026-09-25" }, { "label": "Agentic Cycle", From d2c6910ce7730511bbe4c0e6bed191e14739d850 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 25 Sep 2026 17:44:13 +0200 Subject: [PATCH 004/286] feat(examples): add a subagent brief demo to ts-react-chat The /subagent-brief page shows each researcher card with the brief the parent model wrote and the child's answer. --- .../ts-react-chat/src/components/Header.tsx | 14 ++ examples/ts-react-chat/src/routeTree.gen.ts | 42 ++++++ .../src/routes/api.subagent-brief.ts | 68 +++++++++ .../src/routes/subagent-brief.tsx | 135 ++++++++++++++++++ 4 files changed, 259 insertions(+) create mode 100644 examples/ts-react-chat/src/routes/api.subagent-brief.ts create mode 100644 examples/ts-react-chat/src/routes/subagent-brief.tsx diff --git a/examples/ts-react-chat/src/components/Header.tsx b/examples/ts-react-chat/src/components/Header.tsx index fc6d50cbb5..20d344a520 100644 --- a/examples/ts-react-chat/src/components/Header.tsx +++ b/examples/ts-react-chat/src/components/Header.tsx @@ -18,6 +18,7 @@ import { MessageSquare, Mic, Music, + Network, PauseCircle, Plug, RefreshCw, @@ -381,6 +382,19 @@ export default function Header() { WebSocket Chat + setIsOpen(false)} + className="flex items-center gap-3 p-3 rounded-lg hover:bg-gray-800 transition-colors mb-2" + activeProps={{ + className: + 'flex items-center gap-3 p-3 rounded-lg bg-cyan-600 hover:bg-cyan-700 transition-colors mb-2', + }} + > + + Subagent Brief + + setIsOpen(false)} diff --git a/examples/ts-react-chat/src/routeTree.gen.ts b/examples/ts-react-chat/src/routeTree.gen.ts index b0edbb8e74..4e7cafac3c 100644 --- a/examples/ts-react-chat/src/routeTree.gen.ts +++ b/examples/ts-react-chat/src/routeTree.gen.ts @@ -12,6 +12,7 @@ import { Route as rootRouteImport } from './routes/__root' import { Route as WebsocketChatRouteImport } from './routes/websocket-chat' import { Route as TypesafeToolsRouteImport } from './routes/typesafe-tools' import { Route as ThreadsRouteImport } from './routes/threads' +import { Route as SubagentBriefRouteImport } from './routes/subagent-brief' import { Route as ServerFnChatRouteImport } from './routes/server-fn-chat' import { Route as SandboxesDurableRouteImport } from './routes/sandboxes-durable' import { Route as SandboxesRouteImport } from './routes/sandboxes' @@ -46,6 +47,7 @@ import { Route as ExampleRuntimeContextRouteImport } from './routes/example.runt import { Route as ApiTranscribeRouteImport } from './routes/api.transcribe' import { Route as ApiTanchatRouteImport } from './routes/api.tanchat' import { Route as ApiSummarizeRouteImport } from './routes/api.summarize' +import { Route as ApiSubagentBriefRouteImport } from './routes/api.subagent-brief' import { Route as ApiStructuredOutputRouteImport } from './routes/api.structured-output' import { Route as ApiStructuredChatRouteImport } from './routes/api.structured-chat' import { Route as ApiSandboxTriageDurableRouteImport } from './routes/api.sandbox-triage-durable' @@ -96,6 +98,11 @@ const ThreadsRoute = ThreadsRouteImport.update({ path: '/threads', getParentRoute: () => rootRouteImport, } as any) +const SubagentBriefRoute = SubagentBriefRouteImport.update({ + id: '/subagent-brief', + path: '/subagent-brief', + getParentRoute: () => rootRouteImport, +} as any) const ServerFnChatRoute = ServerFnChatRouteImport.update({ id: '/server-fn-chat', path: '/server-fn-chat', @@ -271,6 +278,11 @@ const ApiSummarizeRoute = ApiSummarizeRouteImport.update({ path: '/api/summarize', getParentRoute: () => rootRouteImport, } as any) +const ApiSubagentBriefRoute = ApiSubagentBriefRouteImport.update({ + id: '/api/subagent-brief', + path: '/api/subagent-brief', + getParentRoute: () => rootRouteImport, +} as any) const ApiStructuredOutputRoute = ApiStructuredOutputRouteImport.update({ id: '/api/structured-output', path: '/api/structured-output', @@ -464,6 +476,7 @@ export interface FileRoutesByFullPath { '/sandboxes': typeof SandboxesRoute '/sandboxes-durable': typeof SandboxesDurableRoute '/server-fn-chat': typeof ServerFnChatRoute + '/subagent-brief': typeof SubagentBriefRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute '/websocket-chat': typeof WebsocketChatRoute @@ -494,6 +507,7 @@ export interface FileRoutesByFullPath { '/api/sandbox-triage-durable': typeof ApiSandboxTriageDurableRoute '/api/structured-chat': typeof ApiStructuredChatRoute '/api/structured-output': typeof ApiStructuredOutputRoute + '/api/subagent-brief': typeof ApiSubagentBriefRoute '/api/summarize': typeof ApiSummarizeRoute '/api/tanchat': typeof ApiTanchatRoute '/api/transcribe': typeof ApiTranscribeRoute @@ -537,6 +551,7 @@ export interface FileRoutesByTo { '/sandboxes': typeof SandboxesRoute '/sandboxes-durable': typeof SandboxesDurableRoute '/server-fn-chat': typeof ServerFnChatRoute + '/subagent-brief': typeof SubagentBriefRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute '/websocket-chat': typeof WebsocketChatRoute @@ -567,6 +582,7 @@ export interface FileRoutesByTo { '/api/sandbox-triage-durable': typeof ApiSandboxTriageDurableRoute '/api/structured-chat': typeof ApiStructuredChatRoute '/api/structured-output': typeof ApiStructuredOutputRoute + '/api/subagent-brief': typeof ApiSubagentBriefRoute '/api/summarize': typeof ApiSummarizeRoute '/api/tanchat': typeof ApiTanchatRoute '/api/transcribe': typeof ApiTranscribeRoute @@ -611,6 +627,7 @@ export interface FileRoutesById { '/sandboxes': typeof SandboxesRoute '/sandboxes-durable': typeof SandboxesDurableRoute '/server-fn-chat': typeof ServerFnChatRoute + '/subagent-brief': typeof SubagentBriefRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute '/websocket-chat': typeof WebsocketChatRoute @@ -641,6 +658,7 @@ export interface FileRoutesById { '/api/sandbox-triage-durable': typeof ApiSandboxTriageDurableRoute '/api/structured-chat': typeof ApiStructuredChatRoute '/api/structured-output': typeof ApiStructuredOutputRoute + '/api/subagent-brief': typeof ApiSubagentBriefRoute '/api/summarize': typeof ApiSummarizeRoute '/api/tanchat': typeof ApiTanchatRoute '/api/transcribe': typeof ApiTranscribeRoute @@ -686,6 +704,7 @@ export interface FileRouteTypes { | '/sandboxes' | '/sandboxes-durable' | '/server-fn-chat' + | '/subagent-brief' | '/threads' | '/typesafe-tools' | '/websocket-chat' @@ -716,6 +735,7 @@ export interface FileRouteTypes { | '/api/sandbox-triage-durable' | '/api/structured-chat' | '/api/structured-output' + | '/api/subagent-brief' | '/api/summarize' | '/api/tanchat' | '/api/transcribe' @@ -759,6 +779,7 @@ export interface FileRouteTypes { | '/sandboxes' | '/sandboxes-durable' | '/server-fn-chat' + | '/subagent-brief' | '/threads' | '/typesafe-tools' | '/websocket-chat' @@ -789,6 +810,7 @@ export interface FileRouteTypes { | '/api/sandbox-triage-durable' | '/api/structured-chat' | '/api/structured-output' + | '/api/subagent-brief' | '/api/summarize' | '/api/tanchat' | '/api/transcribe' @@ -832,6 +854,7 @@ export interface FileRouteTypes { | '/sandboxes' | '/sandboxes-durable' | '/server-fn-chat' + | '/subagent-brief' | '/threads' | '/typesafe-tools' | '/websocket-chat' @@ -862,6 +885,7 @@ export interface FileRouteTypes { | '/api/sandbox-triage-durable' | '/api/structured-chat' | '/api/structured-output' + | '/api/subagent-brief' | '/api/summarize' | '/api/tanchat' | '/api/transcribe' @@ -906,6 +930,7 @@ export interface RootRouteChildren { SandboxesRoute: typeof SandboxesRoute SandboxesDurableRoute: typeof SandboxesDurableRoute ServerFnChatRoute: typeof ServerFnChatRoute + SubagentBriefRoute: typeof SubagentBriefRoute ThreadsRoute: typeof ThreadsRoute TypesafeToolsRoute: typeof TypesafeToolsRoute WebsocketChatRoute: typeof WebsocketChatRoute @@ -936,6 +961,7 @@ export interface RootRouteChildren { ApiSandboxTriageDurableRoute: typeof ApiSandboxTriageDurableRoute ApiStructuredChatRoute: typeof ApiStructuredChatRoute ApiStructuredOutputRoute: typeof ApiStructuredOutputRoute + ApiSubagentBriefRoute: typeof ApiSubagentBriefRoute ApiSummarizeRoute: typeof ApiSummarizeRoute ApiTanchatRoute: typeof ApiTanchatRoute ApiTranscribeRoute: typeof ApiTranscribeRoute @@ -981,6 +1007,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ThreadsRouteImport parentRoute: typeof rootRouteImport } + '/subagent-brief': { + id: '/subagent-brief' + path: '/subagent-brief' + fullPath: '/subagent-brief' + preLoaderRoute: typeof SubagentBriefRouteImport + parentRoute: typeof rootRouteImport + } '/server-fn-chat': { id: '/server-fn-chat' path: '/server-fn-chat' @@ -1219,6 +1252,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiSummarizeRouteImport parentRoute: typeof rootRouteImport } + '/api/subagent-brief': { + id: '/api/subagent-brief' + path: '/api/subagent-brief' + fullPath: '/api/subagent-brief' + preLoaderRoute: typeof ApiSubagentBriefRouteImport + parentRoute: typeof rootRouteImport + } '/api/structured-output': { id: '/api/structured-output' path: '/api/structured-output' @@ -1492,6 +1532,7 @@ const rootRouteChildren: RootRouteChildren = { SandboxesRoute: SandboxesRoute, SandboxesDurableRoute: SandboxesDurableRoute, ServerFnChatRoute: ServerFnChatRoute, + SubagentBriefRoute: SubagentBriefRoute, ThreadsRoute: ThreadsRoute, TypesafeToolsRoute: TypesafeToolsRoute, WebsocketChatRoute: WebsocketChatRoute, @@ -1522,6 +1563,7 @@ const rootRouteChildren: RootRouteChildren = { ApiSandboxTriageDurableRoute: ApiSandboxTriageDurableRoute, ApiStructuredChatRoute: ApiStructuredChatRoute, ApiStructuredOutputRoute: ApiStructuredOutputRoute, + ApiSubagentBriefRoute: ApiSubagentBriefRoute, ApiSummarizeRoute: ApiSummarizeRoute, ApiTanchatRoute: ApiTanchatRoute, ApiTranscribeRoute: ApiTranscribeRoute, diff --git a/examples/ts-react-chat/src/routes/api.subagent-brief.ts b/examples/ts-react-chat/src/routes/api.subagent-brief.ts new file mode 100644 index 0000000000..0f229eb84c --- /dev/null +++ b/examples/ts-react-chat/src/routes/api.subagent-brief.ts @@ -0,0 +1,68 @@ +import { createFileRoute } from '@tanstack/react-router' +import { + chat, + chatParamsFromRequestBody, + defineAgent, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' + +/** + * 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.', + ), + }), + run: (ctx) => { + console.log( + '[subagent-brief] child input:', + JSON.stringify(ctx.input), + '| parent messages available:', + ctx.messages.length, + ) + return chat({ + adapter: openaiText('gpt-5.6'), + messages: [{ role: 'user', content: ctx.input.task }], + threadId: ctx.threadId, + runId: ctx.runId, + parentRunId: ctx.parentRunId, + subagentRunId: ctx.subagentRunId, + resume: ctx.resume, + }) + }, +}) + +export const Route = createFileRoute('/api/subagent-brief')({ + server: { + handlers: { + POST: async ({ request }) => { + const abortController = new AbortController() + const params = await chatParamsFromRequestBody(await request.json()) + const stream = chat({ + adapter: openaiText('gpt-5.6'), + messages: params.messages, + threadId: params.threadId, + runId: params.runId, + ...(params.parentRunId ? { parentRunId: params.parentRunId } : {}), + ...(params.resume ? { resume: params.resume } : {}), + systemPrompts: [ + 'For any research question, call the researcher tool. Then answer the user from its result.', + ], + subagents: { agents: [researcher] }, + abortController, + }) + return toServerSentEventsResponse(stream, { abortController }) + }, + }, + }, +}) diff --git a/examples/ts-react-chat/src/routes/subagent-brief.tsx b/examples/ts-react-chat/src/routes/subagent-brief.tsx new file mode 100644 index 0000000000..704e83a75b --- /dev/null +++ b/examples/ts-react-chat/src/routes/subagent-brief.tsx @@ -0,0 +1,135 @@ +import { createFileRoute } from '@tanstack/react-router' +import { useState } from 'react' +import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' +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) => + message.parts.flatMap((part) => + part.type === 'text' ? [part.content] : [], + ), + ) + .join('') +} + +function SubagentBriefPage() { + const { messages, sendMessage, isLoading } = useChat({ + connection: fetchServerSentEvents('/api/subagent-brief'), + }) + const [input, setInput] = useState( + 'We picked three sea animals earlier: squid, octopus, and cuttlefish. Find how many hearts each one has.', + ) + + return ( +
+

Subagent brief

+

+ The parent model writes a short task for the researcher{' '} + subagent. The researcher reads only that task. +

+ +
+ {messages.map((message) => ( +
+
{message.role}
+ {message.parts.map((part, index) => { + if (part.type === 'text' && part.content) { + return ( +

+ {part.content} +

+ ) + } + if (part.type !== 'subagent') return null + const toolCallId = part.subagent.parentToolCallId + return ( +
+
+ {part.subagent.name}{' '} + + {part.subagent.status} + +
+
+ Brief:{' '} + {toolCallId === undefined + ? '(none)' + : (briefOf(message, toolCallId) ?? '(none)')} +
+

+ {childText(part.subagent.messages)} +

+
+ ) + })} +
+ ))} +
+ +
{ + event.preventDefault() + const text = input.trim() + if (!text || isLoading) return + setInput('') + void sendMessage(text) + }} + style={{ marginTop: 16, display: 'flex', gap: 8 }} + > + setInput(event.target.value)} + style={{ flex: 1, padding: 8 }} + /> + +
+
+ ) +} + +const page: React.CSSProperties = { + maxWidth: 720, + margin: '0 auto', + padding: 24, + fontFamily: 'system-ui, sans-serif', +} + +const card: React.CSSProperties = { + border: '1px solid #ddd', + borderRadius: 8, + padding: 12, + display: 'flex', + flexDirection: 'column', + gap: 6, +} + +const roleLabel: React.CSSProperties = { + fontSize: 11, + textTransform: 'uppercase', + color: '#999', +} From e82d79c36e9b2e2f17870949795788217c2d98a2 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 11:40:12 +0200 Subject: [PATCH 005/286] feat(ai): let a subagent return a value and call any activity defineAgent run can now resolve to a plain value (an image result, a string). The value lands on SUBAGENT_FINISHED.result and goes to the parent model as the tool result, with very long strings shortened for the model. run also gets bound activities on ctx (ctx.chat, ctx.generateImage, ...) that fill in the child ids and abort signal, plus ctx.forward and an optional produces field. SubagentsBag.binding lets a host add middleware to every bound call. RunRecord gets optional harness-session fields (kind, activity, agent, result, artifacts, principal, leaseOwner, leaseExpiresAt, checkpoint), and both memory run stores keep them. Also fixes the stale adapter comment about capabilities. --- .changeset/harness-p0-agent-results.md | 6 + docs/chat/subagents.md | 51 +++ docs/config.json | 2 +- packages/ai-persistence/src/memory.ts | 30 +- packages/ai/src/activities/chat/adapter.ts | 6 +- .../ai/src/activities/chat/agents/bound.ts | 151 +++++++++ .../activities/chat/agents/define-agent.ts | 62 +++- .../ai/src/activities/chat/agents/spawn.ts | 114 ++++++- .../src/activities/chat/middleware/index.ts | 3 + .../activities/chat/middleware/run-store.ts | 80 +++-- .../src/activities/chat/tools/tool-calls.ts | 41 ++- packages/ai/src/activities/index.ts | 7 + packages/ai/src/client.ts | 7 + packages/ai/src/index.ts | 10 + .../ai/tests/define-agent-results.test.ts | 298 ++++++++++++++++++ packages/ai/tests/run-record-fields.test.ts | 57 ++++ testing/e2e/fixtures/subagents/basic.json | 31 ++ testing/e2e/src/lib/subagents-test.ts | 11 +- testing/e2e/src/routes/api.subagents-test.ts | 18 ++ testing/e2e/tests/subagents.spec.ts | 19 +- 20 files changed, 944 insertions(+), 60 deletions(-) create mode 100644 .changeset/harness-p0-agent-results.md create mode 100644 packages/ai/src/activities/chat/agents/bound.ts create mode 100644 packages/ai/tests/define-agent-results.test.ts create mode 100644 packages/ai/tests/run-record-fields.test.ts diff --git a/.changeset/harness-p0-agent-results.md b/.changeset/harness-p0-agent-results.md new file mode 100644 index 0000000000..af26a929b6 --- /dev/null +++ b/.changeset/harness-p0-agent-results.md @@ -0,0 +1,6 @@ +--- +'@tanstack/ai': minor +'@tanstack/ai-persistence': patch +--- + +A subagent can now return a plain value and call any activity. `defineAgent`'s `run` can return a promise of any value, such as an image result. The value arrives on `SUBAGENT_FINISHED.result`, and the parent model gets it as the tool result. A very long string in that result (for example a base64 image) reaches the parent model as a short note. `run` also gets the activities on `ctx` (`ctx.chat`, `ctx.generateImage`, `ctx.generateVideo`, and the rest), which fill in the thread id, a run id, and the abort signal, plus `ctx.forward` for a plain `chat()` call. The optional `produces` field says what an agent makes. `RunRecord` gets optional fields for harness sessions (`kind`, `activity`, `agent`, `result`, `artifacts`, `principal`, `leaseOwner`, `leaseExpiresAt`, `checkpoint`), and the memory run stores keep them. diff --git a/docs/chat/subagents.md b/docs/chat/subagents.md index de2f4698a3..25a05c6ad2 100644 --- a/docs/chat/subagents.md +++ b/docs/chat/subagents.md @@ -206,6 +206,57 @@ const stream = chat({ To show the brief in the UI, see [Show the brief on a card](#show-the-brief-on-a-card). +## 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: + +```ts group=subagent-values +import { chat, defineAgent } from '@tanstack/ai' +import { openaiImage, openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' + +const heroImage = defineAgent({ + name: 'heroImage', + description: 'Makes a hero image for a post', + produces: 'image', + inputSchema: z.object({ prompt: z.string() }), + run: (ctx) => + ctx.generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: `${ctx.input.prompt}. Brand colors: blue and white.`, + size: '1536x1024', + }), +}) + +const stream = chat({ + adapter: openaiText('gpt-5.6'), + messages: [{ role: 'user', content: 'Make a hero image about squids' }], + subagents: { agents: [heroImage] }, +}) +``` + +- `ctx.chat`, `ctx.generateImage`, `ctx.generateVideo`, `ctx.generateSpeech`, and the other activities on `ctx` take the same options as the plain functions. They fill in the thread id, a run id, and the abort signal. +- `ctx.chat` also uses the parent conversation (`ctx.messages`) when you do not pass `messages`. +- `run` can return a promise of any value. The value arrives on `SUBAGENT_FINISHED.result`, and the parent model gets it as the tool result. +- A string result also streams as the child's text, so the card shows it. +- A very long string in the result (for example a base64 image) reaches the parent model as a short note, `[omitted 5000 characters]`. The full value stays on `SUBAGENT_FINISHED.result`. +- `produces` says what the agent makes (`'image'`, `'text'`, and so on). It does not change how the agent runs. + +When you call the plain `chat()` instead, spread `ctx.forward` to pass the child's ids and abort controller in one line: + +```ts group=subagent-values +const researcher = defineAgent({ + name: 'researcher', + description: 'Looks up facts', + run: (ctx) => + chat({ + adapter: openaiText('gpt-5.6'), + messages: ctx.messages, + ...ctx.forward, + }), +}) +``` + ## 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 3c6e1f6f68..420a15ec03 100644 --- a/docs/config.json +++ b/docs/config.json @@ -155,7 +155,7 @@ "label": "Subagents", "to": "chat/subagents", "addedAt": "2026-09-21", - "updatedAt": "2026-09-25" + "updatedAt": "2026-09-26" }, { "label": "Agentic Cycle", diff --git a/packages/ai-persistence/src/memory.ts b/packages/ai-persistence/src/memory.ts index 5ab3c85ff7..849e30acf1 100644 --- a/packages/ai-persistence/src/memory.ts +++ b/packages/ai-persistence/src/memory.ts @@ -55,15 +55,9 @@ class MemoryMessageStore implements MessageStore { class MemoryRunStore implements RunStore { private readonly runs = new Map() - createOrResume(input: { - runId: string - threadId: string - status?: RunRecord['status'] - startedAt: number - parentRunId?: string - subagentRunId?: string - name?: string - }): Promise { + createOrResume( + input: Parameters[0], + ): Promise { const existing = this.runs.get(input.runId) if (existing) return Promise.resolve(existing) const record: RunRecord = { @@ -78,25 +72,17 @@ class MemoryRunStore implements RunStore { ? { subagentRunId: input.subagentRunId } : {}), ...(input.name !== undefined ? { name: input.name } : {}), + ...(input.kind !== undefined ? { kind: input.kind } : {}), + ...(input.activity !== undefined ? { activity: input.activity } : {}), + ...(input.agent !== undefined ? { agent: input.agent } : {}), + ...(input.principal !== undefined ? { principal: input.principal } : {}), } this.runs.set(record.runId, record) return Promise.resolve(record) } update( runId: string, - patch: Partial< - Pick< - RunRecord, - | 'status' - | 'finishedAt' - | 'error' - | 'usage' - | 'sandboxKey' - | 'detachedSince' - | 'cancelRequested' - | 'driverEpoch' - > - >, + patch: Parameters[1], ): Promise { const existing = this.runs.get(runId) if (existing) this.runs.set(runId, { ...existing, ...patch }) diff --git a/packages/ai/src/activities/chat/adapter.ts b/packages/ai/src/activities/chat/adapter.ts index 102d297827..dd1b1d2cce 100644 --- a/packages/ai/src/activities/chat/adapter.ts +++ b/packages/ai/src/activities/chat/adapter.ts @@ -83,9 +83,9 @@ export interface TextAdapter< /** * Capabilities this adapter requires at runtime. `chat()` validates that the * configured middleware provides each one. Model adapters omit this; harness - * adapters (e.g. a future `claudeCode()`) declare e.g. `[sandboxCapability]`. - * Runtime access to capabilities from inside the adapter is not yet wired — - * this is the declaration/validation surface only. + * adapters (for example `claudeCode()`) declare `[sandboxCapability]`. + * At runtime the engine passes the provided capabilities to `chatStream` as + * `options.capabilities`. */ readonly requires?: ReadonlyArray diff --git a/packages/ai/src/activities/chat/agents/bound.ts b/packages/ai/src/activities/chat/agents/bound.ts new file mode 100644 index 0000000000..bc73311b1a --- /dev/null +++ b/packages/ai/src/activities/chat/agents/bound.ts @@ -0,0 +1,151 @@ +import { chat } from '../index' +import { summarize } from '../../summarize/index' +import { generateImage } from '../../generateImage/index' +import { generateVideo } from '../../generateVideo/index' +import { generateAudio } from '../../generateAudio/index' +import { generateSpeech } from '../../generateSpeech/index' +import { generateVoice } from '../../generateVoice/index' +import { generateTranscription } from '../../generateTranscription/index' +import { generateWorld } from '../../generateWorld/index' +import { generateLiveVideo } from '../../generateLiveVideo/index' +import { embed } from '../../embed/index' +import { rerank } from '../../rerank/index' +import { decide } from '../../evaluate/index' +import type { GenerationMiddleware } from '../../middleware/types' +import type { AnyChatMiddleware } from '../middleware/types' +import type { RunAgentResumeItem } from '../../../types' +import type { SubagentRunInput } from './define-agent' + +/** + * The fields a child `chat()` needs, in one spread: + * `chat({ adapter, messages, ...ctx.forward })`. + */ +export interface SubagentForward { + threadId: string + runId: string + parentRunId: string + subagentRunId: string + resume?: Array + /** Aborts when the parent run or the subagent group stops. */ + abortController: AbortController +} + +/** + * What a host (for example a harness session) adds to every activity call an + * agent makes through `ctx`. Apps using plain `chat({ subagents })` do not set + * it. + */ +export interface SubagentBinding { + /** Added before the call's own middleware on every `ctx.chat` call. */ + chatMiddleware?: ReadonlyArray + /** Added before the call's own middleware on every generation call. */ + generationMiddleware?: ReadonlyArray +} + +/** + * Activity functions bound to one child run. Each takes the same options as + * the plain function. Options you pass win over the bound defaults. + */ +export interface BoundActivities { + chat: typeof chat + summarize: typeof summarize + generateImage: typeof generateImage + generateVideo: typeof generateVideo + generateAudio: typeof generateAudio + generateSpeech: typeof generateSpeech + generateVoice: typeof generateVoice + generateTranscription: typeof generateTranscription + generateWorld: typeof generateWorld + generateLiveVideo: typeof generateLiveVideo + embed: typeof embed + rerank: typeof rerank + decide: typeof decide +} + +/** Options every bound call can receive. Typed loosely: each wrapper forwards + * to a function whose own signature checks the caller. */ +type LooseOptions = Record & { + middleware?: Array +} + +export function createBoundActivities( + input: SubagentRunInput, + abortController: AbortController, + binding?: SubagentBinding, +): BoundActivities { + let calls = 0 + const signal = abortController.signal + const chatMiddleware = binding?.chatMiddleware ?? [] + const generationMiddleware = binding?.generationMiddleware ?? [] + + // Every generation call gets its own run id under the child run. + const ids = (activity: string) => { + calls += 1 + return { + threadId: input.threadId, + runId: `${input.runId}:${activity}-${calls}`, + } + } + const withGenerationMiddleware = (options: LooseOptions) => ({ + ...options, + middleware: [...generationMiddleware, ...(options.middleware ?? [])], + }) + // Activities that accept thread and run ids plus an abort signal. + const full = (activity: string, options: LooseOptions) => + withGenerationMiddleware({ + ...ids(activity), + abortSignal: signal, + ...options, + }) + + return { + chat: ((options: LooseOptions) => + chat({ + messages: input.messages, + threadId: input.threadId, + runId: input.runId, + parentRunId: input.parentRunId, + subagentRunId: input.subagentRunId, + ...(input.resume ? { resume: input.resume } : {}), + abortController, + ...options, + middleware: [...chatMiddleware, ...(options.middleware ?? [])], + } as never)) as typeof chat, + summarize: ((options: LooseOptions) => + summarize(full('summarize', options) as never)) as typeof summarize, + generateImage: ((options: LooseOptions) => + generateImage(full('image', options) as never)) as typeof generateImage, + generateVideo: ((options: LooseOptions) => + generateVideo(full('video', options) as never)) as typeof generateVideo, + generateAudio: ((options: LooseOptions) => + generateAudio(full('audio', options) as never)) as typeof generateAudio, + generateSpeech: ((options: LooseOptions) => + generateSpeech( + full('speech', options) as never, + )) as typeof generateSpeech, + generateVoice: ((options: LooseOptions) => + generateVoice(full('voice', options) as never)) as typeof generateVoice, + generateTranscription: ((options: LooseOptions) => + generateTranscription( + full('transcription', options) as never, + )) as typeof generateTranscription, + generateWorld: ((options: LooseOptions) => + generateWorld(full('world', options) as never)) as typeof generateWorld, + generateLiveVideo: ((options: LooseOptions) => + generateLiveVideo( + full('liveVideo', options) as never, + )) as typeof generateLiveVideo, + // `embed` takes neither ids nor an abort signal. + embed: ((options: LooseOptions) => + embed(withGenerationMiddleware(options) as never)) as typeof embed, + // `rerank` and `decide` take an abort signal but no ids. + rerank: ((options: LooseOptions) => + rerank( + withGenerationMiddleware({ abortSignal: signal, ...options }) as never, + )) as typeof rerank, + decide: ((options: LooseOptions) => + decide( + withGenerationMiddleware({ abortSignal: signal, ...options }) as never, + )) as typeof decide, + } +} diff --git a/packages/ai/src/activities/chat/agents/define-agent.ts b/packages/ai/src/activities/chat/agents/define-agent.ts index e9822f64b8..0c6ea78b38 100644 --- a/packages/ai/src/activities/chat/agents/define-agent.ts +++ b/packages/ai/src/activities/chat/agents/define-agent.ts @@ -10,12 +10,13 @@ import type { UIMessage, } from '../../../types' import type { AnyClientTool } from '../tools/tool-definition' +import type { BoundActivities, SubagentForward } from './bound' /** - * Context the library passes into {@link defineAgent} `run`. + * What the library knows about a child run before `run` starts. * `TInput` is the agent's `inputSchema`. */ -export interface SubagentRunContext< +export interface SubagentRunInput< TInput extends SchemaInput | undefined = any, > { /** @@ -45,6 +46,35 @@ export interface SubagentRunContext< parentSubagentRunId?: string } +/** + * Context the library passes into {@link defineAgent} `run`. + * + * - `forward` holds the fields a child `chat()` needs, in one spread: + * `chat({ adapter, messages, ...ctx.forward })`. + * - 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. + */ +export type SubagentRunContext = + SubagentRunInput & BoundActivities & { forward: SubagentForward } + +/** + * What an agent makes. Informative: plugins use it to find an agent by the + * kind of output it produces. + */ +export type AgentProduces = + | 'text' + | 'image' + | 'video' + | 'audio' + | 'speech' + | 'voice' + | 'transcription' + | 'embedding' + | 'world' + | 'liveVideo' + | (string & {}) + /** * A tool a child agent can carry into client part types. * Server tools and client tools both qualify. @@ -52,9 +82,10 @@ export interface SubagentRunContext< export type SubagentTool = AnyTool | AnyClientTool /** - * A named child agent. `run` is a `chat()` call (or any stream of AG-UI chunks). + * A named child agent. `run` is a `chat()` call (or any stream of AG-UI + * chunks), or a promise of a plain value such as an image result. * `TTools` and `TSchema` stay on the object so `useChat({ subagents })` can - * type that child's parts. + * type that child's parts. `TResult` is the value a promise `run` resolves to. */ export interface DefinedAgent< TName extends string = string, @@ -63,13 +94,20 @@ export interface DefinedAgent< TInterrupts extends ReadonlyArray> = ReadonlyArray>, TInput extends SchemaInput | undefined = any, + TResult = unknown, + TProduces extends AgentProduces | undefined = AgentProduces | undefined, > extends AGUISubagentInfo { name: TName /** Required here: the router and the synthetic tool both read it. */ description: string run: ( ctx: SubagentRunContext, - ) => AsyncIterable | Promise> + ) => + | AsyncIterable + | Promise> + | Promise + /** What this agent makes. See {@link AgentProduces}. */ + 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 @@ -121,7 +159,19 @@ export function defineAgent< InterruptDefinition > = readonly [], TInput extends SchemaInput | undefined = undefined, ->(agent: DefinedAgent) { + TResult = unknown, + const TProduces extends AgentProduces | undefined = undefined, +>( + agent: DefinedAgent< + TName, + TTools, + TSchema, + TInterrupts, + TInput, + TResult, + TProduces + >, +) { if (agent.name.trim() === '') { throw new Error('defineAgent requires a non-empty name') } diff --git a/packages/ai/src/activities/chat/agents/spawn.ts b/packages/ai/src/activities/chat/agents/spawn.ts index 18e712007a..0b3733de42 100644 --- a/packages/ai/src/activities/chat/agents/spawn.ts +++ b/packages/ai/src/activities/chat/agents/spawn.ts @@ -25,7 +25,13 @@ import type { Tool, UIMessage, } from '../../../types' -import type { DefinedAgent, SubagentRunContext } from './define-agent' +import { createBoundActivities } from './bound' +import type { SubagentBinding } from './bound' +import type { + DefinedAgent, + SubagentRunContext, + SubagentRunInput, +} from './define-agent' import type { ChatMiddleware } from '../middleware/types' import type { SubagentTurn } from './turn' @@ -75,6 +81,12 @@ export interface SubagentsBag< */ order?: 'parallel' | 'sequence' sandbox?: 'own' | 'inherit' + /** + * What a host adds to every activity call a child makes through `ctx` + * (`ctx.chat`, `ctx.generateImage`, and the rest). A harness session sets + * it. Apps do not. + */ + binding?: SubagentBinding } /** What the children of one parent run left behind for the parent terminal. */ @@ -261,6 +273,7 @@ function openAgentStream( }, sink, parentToolCallId, + bag.binding, ) } @@ -402,12 +415,87 @@ export function withChildUsage( return withTanstackMetadata(next, { usage: leftover }) as ParentTerminal } +function isAsyncIterable(value: unknown): value is AsyncIterable { + return ( + typeof value === 'object' && value !== null && Symbol.asyncIterator in value + ) +} + +/** + * Chunks for an agent whose `run` resolved to a plain value instead of a + * stream. A string also streams as the child's text, so the parent model, a + * `sequence` router, and the UI all read it like chat output. + */ +function* valueResultChunks( + value: unknown, + id: string, +): Generator { + if (typeof value === 'string' && value !== '') { + const messageId = `${id}-text` + const timestamp = Date.now() + yield attributeChunk( + { + type: EventType.TEXT_MESSAGE_START, + messageId, + role: 'assistant', + timestamp, + }, + id, + ) + yield attributeChunk( + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId, + delta: value, + timestamp, + }, + id, + ) + yield attributeChunk( + { type: EventType.TEXT_MESSAGE_END, messageId, timestamp }, + id, + ) + } + yield { + type: SUBAGENT_FINISHED, + subagentRunId: id, + ...(value !== undefined ? { result: value } : {}), + timestamp: Date.now(), + } satisfies SubagentFinishedEvent +} + +/** + * Build the context `run` receives: the spawn input plus `forward` and the + * bound activity functions. + */ +function runContext( + input: SubagentRunInput, + abortController: AbortController, + binding?: SubagentBinding, +): SubagentRunContext { + return { + ...input, + forward: { + threadId: input.threadId, + runId: input.runId, + parentRunId: input.parentRunId, + subagentRunId: input.subagentRunId, + ...(input.resume ? { resume: input.resume } : {}), + abortController, + }, + ...createBoundActivities(input, abortController, binding), + } +} + export async function* spawnAgentStream( agent: DefinedAgent, - ctx: SubagentRunContext, + input: SubagentRunInput, sink?: SubagentSink, parentToolCallId?: string, + binding?: SubagentBinding, ): AsyncIterable { + const link = linkAbort(input.abortSignal) + const ctx = runContext(input, link.controller, binding) const id = ctx.subagentRunId yield { type: SUBAGENT_STARTED, @@ -428,11 +516,19 @@ export async function* spawnAgentStream( yield stoppedEvent(id) return } - const stream = await orAbort( + const produced: unknown = await orAbort( Promise.resolve(agent.run(ctx)), ctx.abortSignal, ) - iterator = stream[Symbol.asyncIterator]() + if (!isAsyncIterable(produced)) { + if (ctx.abortSignal?.aborted) { + yield stoppedEvent(id) + return + } + yield* valueResultChunks(produced, id) + return + } + iterator = produced[Symbol.asyncIterator]() while (true) { if (ctx.abortSignal?.aborted) { yield stoppedEvent(id) @@ -522,6 +618,7 @@ export async function* spawnAgentStream( } catch { // Child stream may already be closed or aborted. } + link.dispose() } } @@ -760,6 +857,7 @@ export function createSyntheticSubagentTools( let subagentRunId = suspended?.subagentRunId ?? '' let text = suspended?.text ?? '' let error: string | undefined + let result: unknown try { for await (const chunk of openAgentStream( entry, @@ -792,6 +890,13 @@ export function createSyntheticSubagentTools( ) { error = chunk.message } + if ( + chunk.type === SUBAGENT_FINISHED && + chunk.subagentRunId === subagentRunId && + chunk.result !== undefined + ) { + result = chunk.result + } toolContext?.[EMIT_STREAM_CHUNK]?.(chunk) } } finally { @@ -806,6 +911,7 @@ export function createSyntheticSubagentTools( return { subagentRunId, text, + ...(result !== undefined ? { result } : {}), ...(error !== undefined ? { error } : {}), ...(sink.interrupts.length > 0 ? { interrupts: sink.interrupts } : {}), } satisfies SubagentToolOutcome diff --git a/packages/ai/src/activities/chat/middleware/index.ts b/packages/ai/src/activities/chat/middleware/index.ts index 53da815b24..2c4b8a7fbd 100644 --- a/packages/ai/src/activities/chat/middleware/index.ts +++ b/packages/ai/src/activities/chat/middleware/index.ts @@ -80,6 +80,9 @@ export type { RunStatus, TerminalRunStatus, RunRecord, + RunKind, + RunArtifactRef, + RunCheckpoint, RunError, RunStore, } from './run-store' diff --git a/packages/ai/src/activities/chat/middleware/run-store.ts b/packages/ai/src/activities/chat/middleware/run-store.ts index 93d2361e8a..77847a74d9 100644 --- a/packages/ai/src/activities/chat/middleware/run-store.ts +++ b/packages/ai/src/activities/chat/middleware/run-store.ts @@ -176,6 +176,54 @@ export interface RunRecord { * loser nothing to read. Absent on a run that was never claimed. */ driverEpoch?: number + /** + * What kind of work this run is, when a harness session recorded it: + * a chat turn, an agent run, one activity call inside an agent, a user + * command, or a compaction. Absent on runs recorded by plain `chat()`. + */ + kind?: RunKind + /** The activity of a `kind: 'activity'` run (`'image'`, `'chat'`, ...). */ + activity?: string + /** The agent name of a `kind: 'agent'` run. */ + agent?: string + /** Terminal result metadata (ids, urls, text). Never media bytes. */ + result?: unknown + /** Durable artifact references produced by this run. */ + artifacts?: Array + /** Who started this run, from the host's `authorize`. */ + principal?: { id: string } + /** + * The host that currently drives this run, and when its claim expires. The + * driver renews `leaseExpiresAt`. Another host may resume the run after it + * passes. Pairs with `driverEpoch` for fencing. + */ + leaseOwner?: string + leaseExpiresAt?: number + /** Crash-resume data saved at the last model or tool boundary. */ + checkpoint?: RunCheckpoint +} + +/** Kinds of work a harness session records. */ +export type RunKind = 'chat' | 'agent' | 'activity' | 'command' | 'compact' + +/** A durable artifact a run produced. */ +export interface RunArtifactRef { + artifactId: string + name?: string + mimeType?: string + url?: string +} + +/** What a run saved at its last model or tool boundary. */ +export interface RunCheckpoint { + /** Epoch ms of the checkpoint. */ + at: number + /** Tools that started but had no result yet at the checkpoint. */ + pendingTools?: Array<{ + toolCallId: string + name: string + replay: 'safe' | 'never' + }> } /** @@ -211,7 +259,7 @@ export interface RunStore { parentRunId?: string subagentRunId?: string name?: string - }, + } & Partial>, ) => Promise /** * Patch a record's mutable fields. @@ -232,6 +280,11 @@ export interface RunStore { | 'detachedSince' | 'cancelRequested' | 'driverEpoch' + | 'result' + | 'artifacts' + | 'leaseOwner' + | 'leaseExpiresAt' + | 'checkpoint' > >, ) => Promise @@ -371,12 +424,7 @@ export class InMemoryRunStore implements RunStore { private readonly runs = new Map() createOrResume( - input: Pick & { - status?: RunStatus - parentRunId?: string - subagentRunId?: string - name?: string - }, + input: Parameters[0], ): Promise { const existing = this.runs.get(input.runId) if (existing) return Promise.resolve(existing) @@ -392,6 +440,10 @@ export class InMemoryRunStore implements RunStore { ? { subagentRunId: input.subagentRunId } : {}), ...(input.name !== undefined ? { name: input.name } : {}), + ...(input.kind !== undefined ? { kind: input.kind } : {}), + ...(input.activity !== undefined ? { activity: input.activity } : {}), + ...(input.agent !== undefined ? { agent: input.agent } : {}), + ...(input.principal !== undefined ? { principal: input.principal } : {}), } this.runs.set(record.runId, record) return Promise.resolve(record) @@ -399,19 +451,7 @@ export class InMemoryRunStore implements RunStore { update( runId: string, - patch: Partial< - Pick< - RunRecord, - | 'status' - | 'finishedAt' - | 'error' - | 'usage' - | 'sandboxKey' - | 'detachedSince' - | 'cancelRequested' - | 'driverEpoch' - > - >, + patch: Parameters[1], ): Promise { const existing = this.runs.get(runId) if (existing) this.runs.set(runId, { ...existing, ...patch }) diff --git a/packages/ai/src/activities/chat/tools/tool-calls.ts b/packages/ai/src/activities/chat/tools/tool-calls.ts index 2315ae5225..5293c2df7c 100644 --- a/packages/ai/src/activities/chat/tools/tool-calls.ts +++ b/packages/ai/src/activities/chat/tools/tool-calls.ts @@ -56,6 +56,12 @@ export const EMIT_STREAM_CHUNK = Symbol.for('tanstack.ai.emitStreamChunk') export interface SubagentToolOutcome { subagentRunId: string text: string + /** + * The child's typed result: the value a promise `run` resolved to, or the + * `RUN_FINISHED.result` of a child stream. When set, the parent model gets + * this instead of the child's text. + */ + result?: unknown error?: string /** Set when the child stopped for outside input. The tool call stays open. */ interrupts?: Array @@ -65,6 +71,33 @@ function isSubagentTool(tool: AnyTool): boolean { return (tool as { [SUBAGENT_TOOL]?: true })[SUBAGENT_TOOL] === true } +/** Longest string a subagent result sends to the parent model unchanged. */ +const MODEL_RESULT_MAX_STRING = 2048 + +/** + * A copy of a subagent result that is safe to send to the parent model. + * A string longer than {@link MODEL_RESULT_MAX_STRING} (for example a base64 + * image) becomes a short note, so one image result cannot fill the context. + * The full result still travels on `SUBAGENT_FINISHED` for the UI. + */ +export function compactForModel(value: unknown, depth = 0): unknown { + if (typeof value === 'string') { + return value.length > MODEL_RESULT_MAX_STRING + ? `[omitted ${value.length} characters]` + : value + } + if (depth > 20 || typeof value !== 'object' || value === null) return value + if (Array.isArray(value)) { + return value.map((item) => compactForModel(item, depth + 1)) + } + return Object.fromEntries( + Object.entries(value).map(([key, item]) => [ + key, + compactForModel(item, depth + 1), + ]), + ) +} + /** * MCP Apps metadata attached to a server tool at discovery (see * `@tanstack/ai-mcp` discovery + `MCPManager.discover()`). @@ -672,7 +705,13 @@ export async function* executeServerTool( } const modelResult = outcome.error ? { subagentRunId: outcome.subagentRunId, error: outcome.error } - : { subagentRunId: outcome.subagentRunId, result: outcome.text } + : { + subagentRunId: outcome.subagentRunId, + result: + outcome.result !== undefined + ? compactForModel(outcome.result) + : outcome.text, + } results.push({ toolCallId: toolCall.id, toolName, diff --git a/packages/ai/src/activities/index.ts b/packages/ai/src/activities/index.ts index 2e7f298748..c2d9013816 100644 --- a/packages/ai/src/activities/index.ts +++ b/packages/ai/src/activities/index.ts @@ -45,7 +45,14 @@ export { type DefinedAgent, type SubagentChoiceOptions, type SubagentRunContext, + type SubagentRunInput, + type AgentProduces, } from './chat/agents/define-agent' +export type { + BoundActivities, + SubagentBinding, + SubagentForward, +} from './chat/agents/bound' export { subagentRoute, type SubagentRouteOptions } from './chat/agents/route' export type { SubagentOrder, diff --git a/packages/ai/src/client.ts b/packages/ai/src/client.ts index be1d03f72d..e6ab3963da 100644 --- a/packages/ai/src/client.ts +++ b/packages/ai/src/client.ts @@ -231,7 +231,14 @@ export { type DefinedAgent, type SubagentChoiceOptions, type SubagentRunContext, + type SubagentRunInput, + type AgentProduces, } from './activities/chat/agents/define-agent' +export type { + BoundActivities, + SubagentBinding, + SubagentForward, +} from './activities/chat/agents/bound' export { subagentRoute, type SubagentRouteOptions, diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 30f7001403..bb2055aa0c 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -91,7 +91,14 @@ export { type DefinedAgent, type SubagentChoiceOptions, type SubagentRunContext, + type SubagentRunInput, + type AgentProduces, } from './activities/chat/agents/define-agent' +export type { + BoundActivities, + SubagentBinding, + SubagentForward, +} from './activities/chat/agents/bound' export { subagentRoute, type SubagentRouteOptions, @@ -367,6 +374,9 @@ export type { RunStatus, TerminalRunStatus, RunRecord, + RunKind, + RunArtifactRef, + RunCheckpoint, RunError, RunStore, } from './activities/chat/middleware/index' diff --git a/packages/ai/tests/define-agent-results.test.ts b/packages/ai/tests/define-agent-results.test.ts new file mode 100644 index 0000000000..647f1bbdfb --- /dev/null +++ b/packages/ai/tests/define-agent-results.test.ts @@ -0,0 +1,298 @@ +import { describe, expect, expectTypeOf, it, vi } from 'vitest' +import { z } from 'zod' +import { defineAgent } from '../src/activities/chat/agents/define-agent' +import { chat } from '../src/activities/chat' +import { compactForModel } from '../src/activities/chat/tools/tool-calls' +import { collectChunks, createMockAdapter, ev } from './test-utils' +import type { ImageAdapter } from '../src/activities/generateImage/adapter' +import type { GenerationMiddleware } from '../src/activities/middleware/types' +import type { ChatMiddleware } from '../src/activities/chat/middleware/types' +import type { DefinedAgent } from '../src/activities/chat/agents/define-agent' +import type { StreamChunk } from '../src/types' + +/** A parent model that calls `agentName` once with `args`, then answers. */ +function parentCalling(agentName: string, args = '{}') { + return createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('call_1', agentName), + ev.toolArgs('call_1', args), + ev.runFinished('tool_calls'), + ], + [ev.runStarted(), ev.runFinished('stop')], + ], + }) +} + +async function runParent( + agent: DefinedAgent, + binding?: { + chatMiddleware?: Array + generationMiddleware?: Array + }, +) { + const { adapter } = parentCalling(agent.name) + return collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + threadId: 'thread-p', + runId: 'run-p', + subagents: { agents: [agent], ...(binding ? { binding } : {}) }, + }) as AsyncIterable, + ) +} + +function toolResult(chunks: Array): unknown { + const chunk = chunks.find( + (entry) => + entry.type === 'TOOL_CALL_RESULT' && entry.toolCallId === 'call_1', + ) + if (chunk?.type !== 'TOOL_CALL_RESULT' || typeof chunk.content !== 'string') { + throw new Error('no tool result') + } + return JSON.parse(chunk.content) +} + +function finished(chunks: Array) { + const chunk = chunks.find((entry) => entry.type === 'SUBAGENT_FINISHED') + if (chunk?.type !== 'SUBAGENT_FINISHED') throw new Error('no finish') + return chunk +} + +function imageAdapter(generateImages = vi.fn()): ImageAdapter { + generateImages.mockResolvedValue({ + id: 'img-1', + model: 'test-model', + images: [{ url: 'https://example.com/a.png' }], + }) + return { + kind: 'image' as const, + name: 'test-image', + model: 'test-model', + '~types': {} as any, + generateImages, + } +} + +describe('defineAgent promise results', () => { + it('sends a promise result to the parent model and on SUBAGENT_FINISHED', async () => { + const agent = defineAgent({ + name: 'pricer', + description: 'Returns a price', + run: () => Promise.resolve({ price: 42, currency: 'EUR' }), + }) + + const chunks = await runParent(agent) + + expect(finished(chunks).result).toEqual({ price: 42, currency: 'EUR' }) + expect(toolResult(chunks)).toMatchObject({ + result: { price: 42, currency: 'EUR' }, + }) + }) + + it('streams a string result as the child text', async () => { + const agent = defineAgent({ + name: 'titler', + description: 'Writes a title', + run: () => Promise.resolve('A short title'), + }) + + const chunks = await runParent(agent) + const started = chunks.find((entry) => entry.type === 'SUBAGENT_STARTED') + const text = chunks.find( + (entry) => + entry.type === 'TEXT_MESSAGE_CONTENT' && 'subagentRunId' in entry, + ) + + expect(text).toMatchObject({ + delta: 'A short title', + subagentRunId: + started?.type === 'SUBAGENT_STARTED' ? started.subagentRunId : 'none', + }) + expect(toolResult(chunks)).toMatchObject({ result: 'A short title' }) + }) + + it('shortens long strings for the model but keeps them on SUBAGENT_FINISHED', async () => { + const b64 = 'a'.repeat(5000) + const agent = defineAgent({ + name: 'painter', + description: 'Paints', + run: () => Promise.resolve({ images: [{ b64Json: b64 }] }), + }) + + const chunks = await runParent(agent) + + expect(finished(chunks).result).toEqual({ images: [{ b64Json: b64 }] }) + expect(toolResult(chunks)).toMatchObject({ + result: { images: [{ b64Json: '[omitted 5000 characters]' }] }, + }) + }) + + it('reports a rejected promise as SUBAGENT_ERROR', async () => { + const agent = defineAgent({ + name: 'broken', + description: 'Fails', + run: () => Promise.reject(new Error('provider down')), + }) + + const chunks = await runParent(agent) + const error = chunks.find((entry) => entry.type === 'SUBAGENT_ERROR') + + expect(error).toMatchObject({ message: 'provider down' }) + expect(toolResult(chunks)).toMatchObject({ error: 'provider down' }) + }) +}) + +describe('agent run context', () => { + it('gives run ctx.forward with the child ids and an abort controller', async () => { + const seen: Array = [] + const agent = defineAgent({ + name: 'probe', + description: 'Records ctx.forward', + run: (ctx) => { + seen.push({ + threadId: ctx.forward.threadId, + runId: ctx.forward.runId, + parentRunId: ctx.forward.parentRunId, + subagentRunId: ctx.forward.subagentRunId, + aborted: ctx.forward.abortController.signal.aborted, + }) + return Promise.resolve('ok') + }, + }) + + await runParent(agent) + + expect(seen).toEqual([ + expect.objectContaining({ + threadId: 'thread-p:probe', + parentRunId: 'run-p', + aborted: false, + }), + ]) + const forward = seen[0] as { runId: string; subagentRunId: string } + expect(forward.runId).toBe(`run-p:${forward.subagentRunId}`) + }) + + it('runs ctx.chat with the child ids, the parent messages, and host middleware', async () => { + const child = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.textStart('m'), + ev.textContent('child says hi', 'm'), + ev.textEnd('m'), + ev.runFinished('stop'), + ], + ], + }) + const onStart = vi.fn() + const agent = defineAgent({ + name: 'helper', + description: 'Uses ctx.chat', + run: (ctx) => ctx.chat({ adapter: child.adapter }), + }) + + const chunks = await runParent(agent, { + chatMiddleware: [{ name: 'host', onStart }], + }) + + expect(onStart).toHaveBeenCalledTimes(1) + const hostCtx = onStart.mock.calls[0]?.[0] as { + threadId: string + subagentRunId?: string + } + expect(hostCtx.threadId).toBe('thread-p:helper') + expect(hostCtx.subagentRunId).toBeDefined() + expect(child.calls[0]?.messages).toEqual( + expect.arrayContaining([expect.objectContaining({ content: 'Go' })]), + ) + expect(toolResult(chunks)).toMatchObject({ result: 'child says hi' }) + }) + + it('runs ctx.generateImage with a child run id, the abort signal, and host middleware', async () => { + const generateImages = vi.fn() + const adapter = imageAdapter(generateImages) + const onStart = vi.fn() + const agent = defineAgent({ + name: 'illustrator', + description: 'Makes an image', + produces: 'image', + inputSchema: z.object({ prompt: z.string() }), + run: (ctx) => + ctx.generateImage({ + adapter, + prompt: ctx.input.prompt, + size: '1024x1024', + }), + }) + + const { adapter: parent } = parentCalling( + 'illustrator', + '{"prompt":"a cat"}', + ) + const chunks = await collectChunks( + chat({ + adapter: parent, + messages: [{ role: 'user', content: 'Draw' }], + threadId: 'thread-p', + runId: 'run-p', + subagents: { + agents: [agent], + binding: { generationMiddleware: [{ name: 'host', onStart }] }, + }, + }) as AsyncIterable, + ) + + expect(generateImages).toHaveBeenCalledWith( + expect.objectContaining({ prompt: 'a cat', size: '1024x1024' }), + ) + const request = generateImages.mock.calls[0]?.[0] as { + abortSignal?: AbortSignal + threadId?: unknown + } + expect(request.abortSignal).toBeInstanceOf(AbortSignal) + // Ids and middleware stay in the activity; the adapter never sees them. + expect(request.threadId).toBeUndefined() + const genCtx = onStart.mock.calls[0]?.[0] as { + threadId: string + runId: string + } + expect(genCtx.threadId).toBe('thread-p:illustrator') + expect(genCtx.runId).toMatch(/^run-p:subagent-.+:image-1$/) + expect(finished(chunks).result).toMatchObject({ + images: [{ url: 'https://example.com/a.png' }], + }) + }) +}) + +describe('agent types', () => { + it('keeps produces as a literal and infers a promise result', () => { + const agent = defineAgent({ + name: 'counter', + description: 'Counts', + produces: 'text', + run: () => Promise.resolve({ count: 1 }), + }) + + expectTypeOf(agent.produces).toEqualTypeOf<'text' | undefined>() + expectTypeOf(agent.run).returns.toMatchTypeOf< + AsyncIterable | Promise + >() + }) +}) + +describe('compactForModel', () => { + it('keeps short values and shortens long strings at any depth', () => { + expect( + compactForModel({ a: 'x', b: [{ c: 'y'.repeat(3000) }], n: 1, z: null }), + ).toEqual({ + a: 'x', + b: [{ c: '[omitted 3000 characters]' }], + n: 1, + z: null, + }) + }) +}) diff --git a/packages/ai/tests/run-record-fields.test.ts b/packages/ai/tests/run-record-fields.test.ts new file mode 100644 index 0000000000..4276de9a2d --- /dev/null +++ b/packages/ai/tests/run-record-fields.test.ts @@ -0,0 +1,57 @@ +import { describe, expect, it } from 'vitest' +import { InMemoryRunStore } from '../src/activities/chat/middleware/run-store' + +describe('RunRecord harness fields', () => { + it('stores kind, activity, agent, and principal on first create only', async () => { + const store = new InMemoryRunStore() + + await store.createOrResume({ + runId: 'op-1', + threadId: 't-1', + startedAt: 1, + kind: 'activity', + activity: 'image', + agent: 'heroImage', + principal: { id: 'user-1' }, + }) + // A resume must not change identity fields. + const again = await store.createOrResume({ + runId: 'op-1', + threadId: 't-1', + startedAt: 2, + kind: 'chat', + }) + + expect(again).toMatchObject({ + kind: 'activity', + activity: 'image', + agent: 'heroImage', + principal: { id: 'user-1' }, + startedAt: 1, + }) + }) + + it('patches result, artifacts, lease, and checkpoint', async () => { + const store = new InMemoryRunStore() + await store.createOrResume({ runId: 'op-2', threadId: 't-1', startedAt: 1 }) + + await store.update('op-2', { + result: { url: 'https://example.com/a.png' }, + artifacts: [{ artifactId: 'art-1', mimeType: 'image/png' }], + leaseOwner: 'host-a', + leaseExpiresAt: 30_000, + checkpoint: { + at: 10, + pendingTools: [{ toolCallId: 'c1', name: 'deploy', replay: 'never' }], + }, + }) + + expect(await store.get('op-2')).toMatchObject({ + result: { url: 'https://example.com/a.png' }, + artifacts: [{ artifactId: 'art-1', mimeType: 'image/png' }], + leaseOwner: 'host-a', + leaseExpiresAt: 30_000, + checkpoint: { at: 10, pendingTools: [{ replay: 'never' }] }, + }) + }) +}) diff --git a/testing/e2e/fixtures/subagents/basic.json b/testing/e2e/fixtures/subagents/basic.json index fd9d5b682f..284943524d 100644 --- a/testing/e2e/fixtures/subagents/basic.json +++ b/testing/e2e/fixtures/subagents/basic.json @@ -109,6 +109,37 @@ "response": { "content": "The comparison is ready." } + }, + { + "match": { + "userMessage": "[subagent-result] price vendor a", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "pricer", + "arguments": "{\"task\":\"[subagent-result-task] Price Vendor A.\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[subagent-result-task] Price Vendor A." + }, + "response": { + "content": "Vendor A costs 42 EUR." + } + }, + { + "match": { + "userMessage": "[subagent-result] price vendor a", + "sequenceIndex": 1 + }, + "response": { + "content": "Pricing is ready." + } } ] } diff --git a/testing/e2e/src/lib/subagents-test.ts b/testing/e2e/src/lib/subagents-test.ts index b7dfad6c23..1b50135a93 100644 --- a/testing/e2e/src/lib/subagents-test.ts +++ b/testing/e2e/src/lib/subagents-test.ts @@ -18,7 +18,12 @@ export const lookupFacts = toolDefinition({ inputSchema: z.object({ topic: z.string() }), }) -export type SubagentScenario = 'route' | 'approval' | 'tool' | 'brief' +export type SubagentScenario = + | 'route' + | 'approval' + | 'tool' + | 'brief' + | 'result' /** The user message for each scenario. It matches `fixtures/subagents`. */ export const SUBAGENT_PROMPTS: Record = { @@ -26,6 +31,7 @@ export const SUBAGENT_PROMPTS: Record = { approval: '[subagent-approval] clean up the old logs', tool: '[subagent-tool] research squids', brief: '[subagent-brief] compare squid facts', + result: '[subagent-result] price vendor a', } export function isSubagentScenario(value: unknown): value is SubagentScenario { @@ -33,6 +39,7 @@ export function isSubagentScenario(value: unknown): value is SubagentScenario { value === 'route' || value === 'approval' || value === 'tool' || - value === 'brief' + value === 'brief' || + value === 'result' ) } diff --git a/testing/e2e/src/routes/api.subagents-test.ts b/testing/e2e/src/routes/api.subagents-test.ts index 561fa9185f..5f03bb403f 100644 --- a/testing/e2e/src/routes/api.subagents-test.ts +++ b/testing/e2e/src/routes/api.subagents-test.ts @@ -30,6 +30,9 @@ 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. + * - `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. */ function subagentsFor( scenario: SubagentScenario, @@ -95,6 +98,21 @@ function subagentsFor( const agents: Array = [researcher] return { agents } } + if (scenario === 'result') { + const pricer = defineAgent({ + name: 'pricer', + description: 'Prices one vendor', + inputSchema: z.object({ task: z.string() }), + run: (ctx) => + ctx.chat({ + ...createTextAdapter('openai', undefined, aimockPort, testId), + messages: [{ role: 'user', content: ctx.input.task }], + stream: false, + }), + }) + const agents: Array = [pricer] + return { agents } + } return { agents: [child('researcher', [])] } } diff --git a/testing/e2e/tests/subagents.spec.ts b/testing/e2e/tests/subagents.spec.ts index 20ae0d3f44..c596063fe6 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', + scenario: 'route' | 'approval' | 'tool' | 'brief' | 'result', testId: string, aimockPort: number, ) { @@ -132,4 +132,21 @@ test.describe('subagents', () => { 'text:Squids have three hearts and blue blood.', ) }) + + test('a child whose run resolves to a string reports it to the parent', async ({ + page, + testId, + aimockPort, + }) => { + await open(page, 'result', testId, aimockPort) + await run(page) + + await expect(page.getByTestId('parent-text')).toHaveText( + 'Pricing is ready.', + ) + await expect(page.getByTestId('card-status-pricer')).toHaveText('finished') + await expect(page.getByTestId('card-pricer')).toContainText( + 'text:Vendor A costs 42 EUR.', + ) + }) }) From 71bef105a672efc1c31ea7cf5e7b7ab8d611f0aa Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 12:09:40 +0200 Subject: [PATCH 006/286] feat(ai-harness): add sessions, typed agents, and plugins New package @tanstack/ai-harness. defineHarness takes the chat() options plus typed agents, plugins, and a busy policy. createHarnessHost().open() opens a long-lived session per thread. A session runs chat turns (prompt), takes messages during a turn (steer, followUp, or busy: queue/steer/reject), answers approvals (resolve), and cancels work (cancel). session.agents..run(input) runs a typed agent from code and notes the result in the transcript; start(input, { wake }) runs it in the background and starts a turn when it is done. Every operation streams AG-UI events with cursors, plus harness.* CUSTOM events for input receipts and operation lifecycle. definePlugin adds tools, prompts, chat middleware, generation middleware, and agents. Plugins provide capabilities to each other, own resources with rollback, and live for the session or for one turn. Collisions name both owners. @tanstack/ai-persistence adds the optional inbox store, so accepted inputs survive a restart. @tanstack/ai exports the helpers a host needs. --- .changeset/harness-p1-core.md | 9 + docs/config.json | 11 + docs/harness/overview.md | 182 ++++ packages/ai-harness/README.md | 71 ++ packages/ai-harness/package.json | 77 ++ packages/ai-harness/src/agents.ts | 93 ++ packages/ai-harness/src/define.ts | 137 +++ packages/ai-harness/src/feed.ts | 102 ++ packages/ai-harness/src/host.ts | 106 +++ packages/ai-harness/src/index.ts | 57 ++ packages/ai-harness/src/operation.ts | 109 +++ packages/ai-harness/src/plugins.ts | 357 +++++++ packages/ai-harness/src/resources.ts | 84 ++ packages/ai-harness/src/session.ts | 870 ++++++++++++++++++ packages/ai-harness/src/types.ts | 99 ++ packages/ai-harness/tests/define.test.ts | 64 ++ packages/ai-harness/tests/helpers.ts | 154 ++++ packages/ai-harness/tests/plugins.test.ts | 167 ++++ packages/ai-harness/tests/session.test.ts | 498 ++++++++++ packages/ai-harness/tests/types.test-d.ts | 39 + packages/ai-harness/tsconfig.json | 8 + packages/ai-harness/vite.config.ts | 35 + packages/ai-persistence/src/index.ts | 5 + packages/ai-persistence/src/memory.ts | 48 +- .../ai-persistence/src/testkit/conformance.ts | 63 ++ packages/ai-persistence/src/types.ts | 50 + packages/ai-persistence/tests/inbox.test.ts | 37 + packages/ai-persistence/tests/memory.test.ts | 1 + .../tests/persistence-types.test-d.ts | 2 + packages/ai/src/index.ts | 8 + pnpm-lock.yaml | 18 + testing/e2e/fixtures/harness/basic.json | 20 + testing/e2e/package.json | 1 + testing/e2e/src/routeTree.gen.ts | 21 + testing/e2e/src/routes/api.harness-test.ts | 80 ++ testing/e2e/tests/harness.spec.ts | 31 + 36 files changed, 3712 insertions(+), 2 deletions(-) create mode 100644 .changeset/harness-p1-core.md create mode 100644 docs/harness/overview.md create mode 100644 packages/ai-harness/README.md create mode 100644 packages/ai-harness/package.json create mode 100644 packages/ai-harness/src/agents.ts create mode 100644 packages/ai-harness/src/define.ts create mode 100644 packages/ai-harness/src/feed.ts create mode 100644 packages/ai-harness/src/host.ts create mode 100644 packages/ai-harness/src/index.ts create mode 100644 packages/ai-harness/src/operation.ts create mode 100644 packages/ai-harness/src/plugins.ts create mode 100644 packages/ai-harness/src/resources.ts create mode 100644 packages/ai-harness/src/session.ts create mode 100644 packages/ai-harness/src/types.ts create mode 100644 packages/ai-harness/tests/define.test.ts create mode 100644 packages/ai-harness/tests/helpers.ts create mode 100644 packages/ai-harness/tests/plugins.test.ts create mode 100644 packages/ai-harness/tests/session.test.ts create mode 100644 packages/ai-harness/tests/types.test-d.ts create mode 100644 packages/ai-harness/tsconfig.json create mode 100644 packages/ai-harness/vite.config.ts create mode 100644 packages/ai-persistence/tests/inbox.test.ts create mode 100644 testing/e2e/fixtures/harness/basic.json create mode 100644 testing/e2e/src/routes/api.harness-test.ts create mode 100644 testing/e2e/tests/harness.spec.ts diff --git a/.changeset/harness-p1-core.md b/.changeset/harness-p1-core.md new file mode 100644 index 0000000000..28c75291c4 --- /dev/null +++ b/.changeset/harness-p1-core.md @@ -0,0 +1,9 @@ +--- +'@tanstack/ai-harness': minor +'@tanstack/ai': minor +'@tanstack/ai-persistence': minor +--- + +New package `@tanstack/ai-harness`. `defineHarness` takes the same options as `chat()`, plus typed `agents`, `plugins`, and a `busy` policy. `createHarnessHost().open(harness, { threadId })` opens a long-lived session. A session runs chat turns with `prompt`, takes messages during a turn with `steer` and `followUp`, answers approvals with `resolve`, and cancels work with `cancel`. `session.agents..run(input)` runs a typed agent from code, and `start(input, { wake: true })` runs it in the background. Every operation streams AG-UI events with cursors. `definePlugin` adds tools, prompts, chat middleware, generation middleware, and agents, with capabilities that plugins provide to each other and resources that close with the session or the turn. + +`@tanstack/ai` exports `CapabilityRegistry`, `runAgentStream`, `createSubagentId`, `compactForModel`, and the `SubagentsBag` type for harness hosts. `@tanstack/ai-persistence` adds the optional `inbox` store (`InboxStore`, `defineInboxStore`) that keeps accepted session inputs across restarts. The memory backend and the conformance testkit cover it. diff --git a/docs/config.json b/docs/config.json index 420a15ec03..e08f6892b5 100644 --- a/docs/config.json +++ b/docs/config.json @@ -800,6 +800,17 @@ ], "tab": "guides" }, + { + "label": "Harness", + "tab": "guides", + "children": [ + { + "label": "Build your first harness", + "to": "harness/overview", + "addedAt": "2026-09-26" + } + ] + }, { "label": "Sandboxes", "tab": "guides", diff --git a/docs/harness/overview.md b/docs/harness/overview.md new file mode 100644 index 0000000000..5364f1e554 --- /dev/null +++ b/docs/harness/overview.md @@ -0,0 +1,182 @@ +--- +title: Build your first harness +id: harness-overview +order: 1 +description: "Keep one agent conversation alive across many turns. Queue or steer messages while it works, run typed agents from code, and add plugins." +keywords: + - tanstack ai + - harness + - defineHarness + - session + - plugins + - agents +--- + +`chat()` answers one request. A coding agent, a support bot, or a media studio needs more: one conversation that stays open, takes new messages while it works, and runs other agents from code. A harness gives you that. You define it once with the same options as `chat()`, then open a session for each conversation. + +At the end of this page you have a session that answers prompts, queues and steers messages, runs a typed agent, and loads a plugin. + +## Install + + + +react: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai +vue: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai +solid: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai +svelte: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai +preact: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai +angular: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai +vanilla: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai +octane: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai + + + +## 1. Define the harness + +`defineHarness` takes the `chat()` options you know (`adapter`, `systemPrompts`, `tools`, `middleware`, `subagents`), plus a `name`. Importing it starts nothing. + +```ts group=harness-first +import { defineAgent } from '@tanstack/ai' +import { defineHarness } from '@tanstack/ai-harness' +import { openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' + +const pricer = defineAgent({ + name: 'pricer', + description: 'Looks up the price of a vendor plan', + inputSchema: z.object({ vendor: z.string() }), + run: async (ctx) => ({ vendor: ctx.input.vendor, cents: 1200 }), +}) + +export const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-5.6'), + systemPrompts: ['You are a helpful assistant.'], + agents: [pricer], +}) +``` + +`agents` are agents you run from code. Put an agent in `subagents` instead when the model must call it as a tool, the same as in `chat()`. + +## 2. Open a session + +A host runs sessions. Give it the stores from `@tanstack/ai-persistence`, then open a session for a conversation id. + +```ts group=harness-first +import { createHarnessHost } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' + +const host = createHarnessHost({ persistence: memoryPersistence() }) +const session = await host.open(assistant, { threadId: 'thread-1' }) +``` + +Opening the same `threadId` again returns the live session. + +## 3. Send a prompt + +`prompt` returns an operation. Await it for the text of the turn. + +```ts group=harness-first +const turn = await session.prompt('Write a haiku about the sea.') +console.log(turn.text) +``` + +The session saves the transcript, so the next prompt continues the same conversation. + +To show the answer while it streams, read the events of the operation: + +```ts group=harness-first +const next = session.prompt('Now one about the mountains.') +for await (const entry of next.events()) { + if (entry.event.type === 'TEXT_MESSAGE_CONTENT') { + process.stdout.write(entry.event.delta) + } +} +``` + +The events are standard AG-UI events. `session.events({ from: cursor })` gives the events of every operation, from a cursor you saved. + +## 4. Send messages while it works + +Users type while the agent works. Pick what a new message does: + +- `session.prompt(text)`: runs after the current turn. This is the default. +- `session.prompt(text, { busy: 'reject' })`: refuses the message while a turn runs. +- `session.steer(text)`: adds the message to the running turn, before its next model call. +- `session.followUp(text)`: runs after all current work. + +```ts group=harness-first +const long = session.prompt('Plan a three-day trip to Lisbon.') +await session.steer('Keep it under 500 euros.') +await long +``` + +If the turn makes no further model call, the steer message runs as the next turn. + +## 5. Run an agent from code + +Every agent in `agents` and `subagents` is on `session.agents`. The input and the result are typed from the definition. + +```ts group=harness-first +const price = await session.agents.pricer.run({ vendor: 'acme' }) +console.log(price.cents) +``` + +The session adds a short note about the result to the transcript. The model sees it on the next turn. To run the agent in the background and start a new turn when it is done, use `start`: + +```ts group=harness-first +session.agents.pricer.start({ vendor: 'globex' }, { wake: true }) +``` + +## 6. Add a plugin + +A plugin adds tools, prompts, middleware, and agents to every session. `setup` runs once when the session opens. Resources you acquire close when the session closes. + +```ts group=harness-first +import { toolDefinition } from '@tanstack/ai' +import { definePlugin } from '@tanstack/ai-harness' + +const clock = definePlugin({ + name: 'acme/clock', + setup: () => ({ + prompts: [`Today is ${new Date().toDateString()}.`], + tools: [ + toolDefinition({ + name: 'current_time', + description: 'Get the current time', + }).server(async () => new Date().toISOString()), + ], + }), +}) + +export const withClock = defineHarness({ + name: 'acme/assistant-with-clock', + adapter: openaiText('gpt-5.6'), + plugins: () => [clock], +}) +``` + +`plugins` is a function, so every session gets its own plugin instances. Set `lifetime: 'run'` on a plugin to set it up again for each turn. + +If two plugins add a tool with the same name, `host.open` fails and names both plugins. + +## Keep sessions across restarts + +The host takes the same stores as `withPersistence`. Add an `inbox` store and the session also keeps the messages it accepted but did not start yet. When the session opens again, it runs them. + +- `messages`: the transcript. Required. +- `runs`: the record of every turn and agent run. +- `interrupts`: approvals that wait for a user. +- `inbox`: accepted messages that did not run yet. + +`memoryPersistence()` has every store, but it keeps them in memory only. Write your own stores to keep data in your database. See [Build your own adapter](../persistence/build-your-own-adapter). + +## What you have now + +- A harness defined once, with the same options as `chat()`. +- A session that keeps the conversation and streams AG-UI events. +- Queue, steer, and follow-up for messages that arrive during a turn. +- Typed agents you run from code, in the foreground or the background. +- Plugins that add tools and prompts to each session. + +Next: open the [subagents guide](../chat/subagents) to let the model call your agents as tools. diff --git a/packages/ai-harness/README.md b/packages/ai-harness/README.md new file mode 100644 index 0000000000..c6a4806cf6 --- /dev/null +++ b/packages/ai-harness/README.md @@ -0,0 +1,71 @@ +
+ + + + TanStack AI + +
+ +
+ +# @tanstack/ai-harness + +Build agent harnesses on TanStack AI: long-lived sessions, typed agents, and plugins. + +`chat()` answers one request. A harness keeps one conversation open across many turns. It takes new messages while it works, runs typed agents from code, and loads plugins that add tools, prompts, and middleware. + +## Installation + +```bash +npm install @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence +``` + +## Usage + +```ts +import { defineAgent } from '@tanstack/ai' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' + +const pricer = defineAgent({ + name: 'pricer', + description: 'Looks up the price of a vendor plan', + inputSchema: z.object({ vendor: z.string() }), + run: async (ctx) => ({ vendor: ctx.input.vendor, cents: 1200 }), +}) + +const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-5.6'), + agents: [pricer], +}) + +const host = createHarnessHost({ persistence: memoryPersistence() }) +const session = await host.open(assistant, { threadId: 'thread-1' }) + +const turn = await session.prompt('Hello!') +console.log(turn.text) + +const price = await session.agents.pricer.run({ vendor: 'acme' }) +console.log(price.cents) +``` + +## Documentation + +Read [Build your first harness](https://tanstack.com/ai/latest/docs/harness/overview). + +## License + +MIT diff --git a/packages/ai-harness/package.json b/packages/ai-harness/package.json new file mode 100644 index 0000000000..5d45fa9b9d --- /dev/null +++ b/packages/ai-harness/package.json @@ -0,0 +1,77 @@ +{ + "name": "@tanstack/ai-harness", + "version": "0.0.0", + "description": "Build agent harnesses on TanStack AI: long-lived sessions, typed agents, and plugins", + "keywords": [ + "agent", + "ai", + "harness", + "plugins", + "session", + "tanstack" + ], + "license": "MIT", + "author": "", + "repository": { + "type": "git", + "url": "git+https://github.com/TanStack/ai.git", + "directory": "packages/ai-harness" + }, + "files": [ + "dist", + "src" + ], + "type": "module", + "sideEffects": false, + "module": "./dist/esm/index.js", + "types": "./dist/esm/index.d.ts", + "exports": { + ".": { + "types": "./dist/esm/index.d.ts", + "import": "./dist/esm/index.js" + } + }, + "scripts": { + "build": "vite build", + "clean": "premove ./build ./dist", + "lint:fix": "oxlint src --type-aware --fix", + "test:build": "publint --strict", + "test:oxlint": "oxlint src --type-aware", + "test:coverage": "vitest run --passWithNoTests --coverage --coverage.include='src/**' --coverage.reporter=text-summary --coverage.reporter=json-summary", + "test:lib": "vitest --passWithNoTests", + "test:lib:dev": "pnpm test:lib --watch", + "test:types": "tsc" + }, + "devDependencies": { + "@tanstack/ai": "workspace:*", + "@tanstack/ai-persistence": "workspace:*", + "@vitest/coverage-v8": "4.1.10", + "zod": "^4.2.0" + }, + "peerDependencies": { + "@tanstack/ai": "workspace:^", + "@tanstack/ai-persistence": "workspace:^" + }, + "nx": { + "targets": { + "test:lib": { + "dependsOn": [ + "build", + "^build" + ] + }, + "test:coverage": { + "dependsOn": [ + "build", + "^build" + ] + }, + "test:types": { + "dependsOn": [ + "build", + "^build" + ] + } + } + } +} diff --git a/packages/ai-harness/src/agents.ts b/packages/ai-harness/src/agents.ts new file mode 100644 index 0000000000..1562c6e4ef --- /dev/null +++ b/packages/ai-harness/src/agents.ts @@ -0,0 +1,93 @@ +import type { + AgentProduces, + DefinedAgent, + InferSchemaType, + SchemaInput, +} from '@tanstack/ai' + +/** Any agent made with `defineAgent`, with its types kept. */ +export type AnyAgent = DefinedAgent + +/** The input an agent takes: its `inputSchema` type, or `undefined`. */ +export type AgentInputOf = TAgent extends { + inputSchema?: infer TSchema +} + ? TSchema extends SchemaInput + ? InferSchemaType + : undefined + : undefined + +/** + * What running an agent resolves to: + * - its `outputSchema` type, when it has one; + * - else the value its promise `run` resolves to; + * - else the child's text. + */ +export type AgentResultOf = + TAgent extends DefinedAgent< + any, + any, + infer TSchema, + any, + any, + infer TResult, + any + > + ? TSchema extends SchemaInput + ? InferSchemaType + : unknown extends TResult + ? string + : TResult + : unknown + +/** A read-only view of the agents a session can run. */ +export interface AgentRegistryView { + /** Every agent, in registration order. */ + list: () => ReadonlyArray + /** The agent named `name`, or `undefined`. */ + get: (name: string) => AnyAgent | undefined + /** The first agent whose `produces` matches, or `undefined`. */ + find: (query: { produces: AgentProduces }) => AnyAgent | undefined +} + +/** + * The agents of one session, with the owner of each (the harness, or a + * plugin name). Adding a name twice with two different agents is an error that + * names both owners. Adding the same agent object twice is allowed, so an + * agent can sit in both `agents` and `subagents.agents`. + */ +export class AgentRegistry implements AgentRegistryView { + private readonly agents = new Map< + string, + { agent: AnyAgent; owner: string } + >() + + add(agent: AnyAgent, owner: string): void { + const existing = this.agents.get(agent.name) + if (existing && existing.agent !== agent) { + throw new Error( + `Duplicate agent "${agent.name}": first owner ${existing.owner}, second owner ${owner}.`, + ) + } + if (!existing) this.agents.set(agent.name, { agent, owner }) + } + + list(): ReadonlyArray { + return [...this.agents.values()].map((entry) => entry.agent) + } + + get(name: string): AnyAgent | undefined { + return this.agents.get(name)?.agent + } + + find(query: { produces: AgentProduces }): AnyAgent | undefined { + return this.list().find((agent) => agent.produces === query.produces) + } + + /** A copy for one chat turn, so run plugins can add agents for that turn only. */ + fork(): AgentRegistry { + const copy = new AgentRegistry() + for (const [name, entry] of this.agents) copy.agents.set(name, entry) + return copy + } +} diff --git a/packages/ai-harness/src/define.ts b/packages/ai-harness/src/define.ts new file mode 100644 index 0000000000..ca22eafa31 --- /dev/null +++ b/packages/ai-harness/src/define.ts @@ -0,0 +1,137 @@ +import type { + AgentLoopStrategy, + AnyChatMiddleware, + AnyTextAdapter, + AnyTool, + InterruptDefinition, + SubagentsBag, + SystemPrompt, +} from '@tanstack/ai' +import type { AnyAgent } from './agents' +import type { HarnessPlugin } from './plugins' +import type { BusyPolicy } from './types' + +/** The `subagents` option: the same as `chat({ subagents })`. */ +export type HarnessSubagents> = Omit< + SubagentsBag, + 'binding' +> + +/** + * What `defineHarness` takes. Where it overlaps with `chat()`, the option + * names and types are the same. + */ +export interface HarnessConfig< + TAdapter extends AnyTextAdapter = AnyTextAdapter, + TAgents extends ReadonlyArray = ReadonlyArray, + TSubagents extends ReadonlyArray = ReadonlyArray, +> { + /** A stable name, for example `'acme/studio'`. */ + name: string + /** The main agent-loop model, the same as `chat({ adapter })`. */ + adapter: TAdapter + systemPrompts?: Array + tools?: ReadonlyArray + middleware?: ReadonlyArray + agentLoopStrategy?: AgentLoopStrategy + modelOptions?: TAdapter['~types']['providerOptions'] + interrupts?: ReadonlyArray> + /** Runtime context passed to middleware hooks and server tools. */ + context?: unknown + /** + * Typed agents the session can run from code, commands, plugins, and + * exposed clients: `session.agents..run(input)`. + */ + agents?: TAgents + /** + * Agents the main model can call as tools, the same as `chat({ subagents })`. + * They are also registered in `session.agents`. + */ + subagents?: HarnessSubagents + /** Called once per session, so every session gets fresh plugin instances. */ + plugins?: () => ReadonlyArray + /** What a `prompt` does while a chat turn runs. Default `'queue'`. */ + busy?: BusyPolicy + /** What clients may call. Nothing is exposed by default. */ + expose?: { + agents?: ReadonlyArray + } +} + +const HARNESS_KIND = 'tanstack-ai-harness' as const + +/** An immutable harness definition. Importing it starts nothing. */ +export type HarnessDefinition< + TAdapter extends AnyTextAdapter = AnyTextAdapter, + TAgents extends ReadonlyArray = ReadonlyArray, + TSubagents extends ReadonlyArray = ReadonlyArray, +> = Readonly> & { + readonly kind: typeof HARNESS_KIND + readonly version: 1 +} + +/** A harness definition with any type parameters. */ +export type AnyHarness = HarnessDefinition + +/** Every agent a harness registers, as a union. */ +export type HarnessAgentsOf = + THarness extends HarnessDefinition + ? TAgents[number] | TSubagents[number] + : never + +/** True for a value made with `defineHarness`. */ +export function isHarnessDefinition(value: unknown): value is AnyHarness { + return ( + typeof value === 'object' && + value !== null && + (value as { kind?: unknown }).kind === HARNESS_KIND && + (value as { version?: unknown }).version === 1 + ) +} + +/** + * Define a harness: a reusable, typed agent configuration. Use the same option + * names as `chat()`, plus `agents`, `plugins`, `busy`, and `expose`. + * + * @example + * ```ts + * const studio = defineHarness({ + * name: 'acme/studio', + * adapter: anthropicText('claude-sonnet-4-5'), + * agents: [heroImage], + * subagents: { agents: [researcher] }, + * }) + * ``` + */ +export function defineHarness< + TAdapter extends AnyTextAdapter, + const TAgents extends ReadonlyArray = readonly [], + const TSubagents extends ReadonlyArray = readonly [], +>( + config: HarnessConfig, +): HarnessDefinition { + if (config.name.trim() === '') { + throw new Error('defineHarness requires a non-empty name') + } + const byName = new Map() + for (const agent of [ + ...(config.agents ?? []), + ...(config.subagents?.agents ?? []), + ]) { + const existing = byName.get(agent.name) + if (existing && existing !== agent) { + throw new Error( + `defineHarness "${config.name}": two different agents are named "${agent.name}".`, + ) + } + byName.set(agent.name, agent) + } + for (const name of config.expose?.agents ?? []) { + if (!byName.has(name)) { + throw new Error( + `defineHarness "${config.name}": expose.agents names "${name}", which is not a registered agent.`, + ) + } + } + return Object.freeze({ ...config, kind: HARNESS_KIND, version: 1 as const }) +} diff --git a/packages/ai-harness/src/feed.ts b/packages/ai-harness/src/feed.ts new file mode 100644 index 0000000000..0d9e031cfb --- /dev/null +++ b/packages/ai-harness/src/feed.ts @@ -0,0 +1,102 @@ +import type { StreamChunk } from '@tanstack/ai' +import type { Cursor, SessionEvent } from './types' + +// ponytail: in-memory log capped at MAX_EVENTS per session. A reader whose +// cursor fell out of the window resumes from the oldest kept event. Durable +// cross-process delivery goes through StreamDurability in the protocol layer. +const MAX_EVENTS = 10_000 + +/** + * The ordered event stream of one session. Every operation publishes here. + * Cursors are opaque to callers and increase with each event. + */ +export class SessionFeed { + private readonly entries: Array = [] + private sequence = 0 + private readonly waiters = new Set<() => void>() + private closed = false + + publish(operationId: string, event: StreamChunk): SessionEvent { + this.sequence += 1 + const entry: SessionEvent = { + cursor: String(this.sequence), + operationId, + event, + } + this.entries.push(entry) + if (this.entries.length > MAX_EVENTS) this.entries.shift() + this.wake() + return entry + } + + /** The cursor of the newest event, or `'0'` for an empty feed. */ + head(): Cursor { + return String(this.sequence) + } + + /** + * Events after `from` (exclusive), then live events until `signal` aborts + * or the feed closes. Pass a `filter` to read one operation's events. + */ + async *read(options: { + from?: Cursor + signal?: AbortSignal + filter?: (entry: SessionEvent) => boolean + /** Stop after `until()` returns true and no buffered event is left. */ + until?: () => boolean + }): AsyncIterable { + let after = Number(options.from ?? '0') + if (!Number.isFinite(after)) after = 0 + while (true) { + if (options.signal?.aborted) return + const next = this.entries.filter( + (entry) => + Number(entry.cursor) > after && + (options.filter ? options.filter(entry) : true), + ) + for (const entry of next) { + if (options.signal?.aborted) return + after = Number(entry.cursor) + yield entry + } + if (next.length === 0) { + // Skip past events the filter rejected, so they are not scanned again. + after = Math.max(after, this.sequence) + } + if (this.closed || options.until?.()) { + const rest = this.entries.filter( + (entry) => + Number(entry.cursor) > after && + (options.filter ? options.filter(entry) : true), + ) + for (const entry of rest) yield entry + return + } + await this.waitForNext(options.signal) + } + } + + close(): void { + this.closed = true + this.wake() + } + + private wake(): void { + const waiters = [...this.waiters] + this.waiters.clear() + for (const wake of waiters) wake() + } + + private waitForNext(signal?: AbortSignal): Promise { + return new Promise((resolve) => { + if (signal?.aborted) return resolve() + const done = () => { + signal?.removeEventListener('abort', done) + this.waiters.delete(done) + resolve() + } + this.waiters.add(done) + signal?.addEventListener('abort', done, { once: true }) + }) + } +} diff --git a/packages/ai-harness/src/host.ts b/packages/ai-harness/src/host.ts new file mode 100644 index 0000000000..dda6dd628e --- /dev/null +++ b/packages/ai-harness/src/host.ts @@ -0,0 +1,106 @@ +import { memoryPersistence } from '@tanstack/ai-persistence' +import { HarnessSession } from './session' +import type { + AIPersistence, + ChatTranscriptStores, + InboxStore, +} from '@tanstack/ai-persistence' +import type { AnyHarness } from './define' +import type { Principal } from './types' + +/** + * The stores a host needs: a message store, plus any of runs, interrupts, + * metadata, and inbox. Without an inbox, inputs are kept in memory and a + * restart loses the ones not yet applied. + */ +export type HarnessPersistence = AIPersistence< + ChatTranscriptStores & { inbox?: InboxStore } +> + +export interface HarnessHostOptions { + /** Where sessions keep their state. Default: in memory, lost on restart. */ + persistence?: HarnessPersistence +} + +export interface OpenSessionOptions { + /** The conversation id. Opening the same id again returns the live session. */ + threadId: string + /** Who opened the session. Stored on inputs and run records. */ + principal?: Principal +} + +/** Runs sessions for one or more harnesses in this process. */ +export interface HarnessHost { + /** + * Open a session, or return the live one for this harness and thread. + * Session plugins are set up here, so a plugin error rejects the promise. + */ + open: ( + harness: THarness, + options: OpenSessionOptions, + ) => Promise> + /** Close every live session. */ + close: () => Promise +} + +let warned = false + +/** + * Create a host for harness sessions. + * + * @example + * ```ts + * const host = createHarnessHost({ persistence }) + * const session = await host.open(studio, { threadId: 'thread-1' }) + * const turn = await session.prompt('Write a haiku about the sea.') + * ``` + */ +export function createHarnessHost( + options: HarnessHostOptions = {}, +): HarnessHost { + if (!options.persistence && !warned) { + warned = true + console.warn( + '[@tanstack/ai-harness] No persistence given: sessions live in memory and are lost on restart.', + ) + } + const persistence = options.persistence ?? memoryPersistence() + // ponytail: a memory inbox when the stores have none. Pass `stores.inbox` + // to keep accepted inputs across restarts. + const inbox = persistence.stores.inbox ?? memoryPersistence().stores.inbox + const sessions = new Map>() + + return { + open(harness, { threadId, principal }) { + const key = `${harness.name}\u0000${threadId}` + let session = sessions.get(key) + if (!session) { + const created = new HarnessSession({ + harness, + threadId, + persistence, + inbox, + ...(principal ? { principal } : {}), + onClose: () => sessions.delete(key), + }) + session = created.open().then( + () => created, + (error: unknown) => { + sessions.delete(key) + throw error + }, + ) + sessions.set(key, session) + } + return session as Promise> + }, + async close() { + const live = await Promise.allSettled(sessions.values()) + await Promise.all( + live + .filter((entry) => entry.status === 'fulfilled') + .map((entry) => entry.value.close()), + ) + }, + } +} diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts new file mode 100644 index 0000000000..ac889765b3 --- /dev/null +++ b/packages/ai-harness/src/index.ts @@ -0,0 +1,57 @@ +export { defineHarness, isHarnessDefinition } from './define' +export type { + AnyHarness, + HarnessAgentsOf, + HarnessConfig, + HarnessDefinition, + HarnessSubagents, +} from './define' + +export { definePlugin } from './plugins' +export type { + HarnessPlugin, + PluginContributions, + PluginDefinition, + PluginLifetime, + PluginPrompt, + PluginSetupContext, +} from './plugins' + +export type { + AgentInputOf, + AgentRegistryView, + AgentResultOf, + AnyAgent, +} from './agents' + +export { createHarnessHost } from './host' +export type { + HarnessHost, + HarnessHostOptions, + HarnessPersistence, + OpenSessionOptions, +} from './host' + +export { HarnessSession } from './session' +export type { + AgentHandle, + AgentHandles, + AgentRunOptions, + AgentStartOptions, + SessionSnapshot, +} from './session' + +export { HARNESS_EVENTS } from './types' +export type { + BusyPolicy, + ChatTurnResult, + Cursor, + HarnessInput, + Operation, + OperationKind, + OperationStatus, + Principal, + Receipt, + SessionEvent, + UserInput, +} from './types' diff --git a/packages/ai-harness/src/operation.ts b/packages/ai-harness/src/operation.ts new file mode 100644 index 0000000000..33f80cd2e6 --- /dev/null +++ b/packages/ai-harness/src/operation.ts @@ -0,0 +1,109 @@ +import type { StreamChunk } from '@tanstack/ai' +import type { SessionFeed } from './feed' +import type { + Cursor, + Operation, + OperationKind, + OperationStatus, + Receipt, + SessionEvent, +} from './types' + +let counter = 0 + +/** A new operation id. It is also the AG-UI `runId` of the operation. */ +export function createOperationId(kind: OperationKind): string { + counter += 1 + return `op-${kind}-${Date.now().toString(36)}-${counter}` +} + +const TERMINAL: ReadonlySet = new Set([ + 'completed', + 'failed', + 'cancelled', + 'interrupted', +]) + +/** The session-side implementation of {@link Operation}. */ +export class OperationImpl implements Operation { + readonly id: string + readonly abortController = new AbortController() + private current: OperationStatus = 'accepted' + private readonly settled: Promise + private resolveResult!: (value: TResult) => void + private rejectResult!: (error: unknown) => void + + constructor( + readonly kind: OperationKind, + private readonly feed: SessionFeed, + private readonly onCancel: ( + operation: OperationImpl, + ) => Promise, + readonly agent?: string, + ) { + this.id = createOperationId(kind) + this.settled = new Promise((resolve, reject) => { + this.resolveResult = resolve + this.rejectResult = reject + }) + // Nobody may await a failed background operation. Do not report that as an + // unhandled rejection. `then` still rejects for callers that await. + this.settled.catch(() => {}) + } + + then( + onfulfilled?: ((value: TResult) => TResult1 | PromiseLike) | null, + onrejected?: ((reason: unknown) => TResult2 | PromiseLike) | null, + ): PromiseLike { + return this.settled.then(onfulfilled, onrejected) + } + + status(): OperationStatus { + return this.current + } + + isSettled(): boolean { + return TERMINAL.has(this.current) + } + + setStatus(status: OperationStatus): void { + this.current = status + } + + publish(event: StreamChunk): SessionEvent { + return this.feed.publish(this.id, event) + } + + finish(status: 'completed' | 'interrupted', result: TResult): void { + this.current = status + this.resolveResult(result) + } + + fail(status: 'failed' | 'cancelled', error: unknown): void { + this.current = status + this.rejectResult(error) + } + + events(options?: { + from?: Cursor + signal?: AbortSignal + }): AsyncIterable { + return this.feed.read({ + ...options, + filter: (entry) => entry.operationId === this.id, + until: () => this.isSettled(), + }) + } + + async *stream(options?: { + signal?: AbortSignal + }): AsyncIterable { + for await (const entry of this.events({ from: '0', ...options })) { + yield entry.event + } + } + + cancel(_reason?: string): Promise { + return this.onCancel(this) + } +} diff --git a/packages/ai-harness/src/plugins.ts b/packages/ai-harness/src/plugins.ts new file mode 100644 index 0000000000..550ec092fa --- /dev/null +++ b/packages/ai-harness/src/plugins.ts @@ -0,0 +1,357 @@ +import { CapabilityRegistry } from '@tanstack/ai' +import { ResourceScope, disposeAll } from './resources' +import type { + AgentProduces, + AnyChatMiddleware, + AnyGenerationMiddleware, + AnyTool, + Capability, + CapabilityHandle, + SystemPrompt, +} from '@tanstack/ai' +import type { AgentRegistry, AgentRegistryView, AnyAgent } from './agents' + +/** + * How long a plugin's `setup` result and resources live: + * - `session` (default): from session open to session close. Resources such + * as an index, a watcher, or an LSP server stay up across turns. + * - `run`: set up for each chat turn, disposed when that turn ends. + */ +export type PluginLifetime = 'session' | 'run' + +/** A prompt a plugin contributes. `id` must be unique in the session. */ +export interface PluginPrompt { + id: string + text: string +} + +/** What a plugin's `setup` returns. Every field is optional. */ +export interface PluginContributions { + /** Tools for the main model. */ + tools?: ReadonlyArray + /** System prompt text for every chat turn. */ + prompts?: ReadonlyArray + /** Chat middleware, the same type as `chat({ middleware })`. */ + middleware?: ReadonlyArray + /** Middleware for the activities agents call (`ctx.generateImage`, ...). */ + generationMiddleware?: ReadonlyArray + /** Agents added to `session.agents`. */ + agents?: ReadonlyArray +} + +/** What a plugin's `setup` receives. */ +export interface PluginSetupContext { + /** Resources this plugin owns, with rollback and reverse-order cleanup. */ + resources: { + acquire: ( + open: () => T | Promise, + close: (resource: T) => unknown, + ) => Promise + signal: AbortSignal + } + /** Read a capability another plugin or middleware provided. Throws if absent. */ + get: (capability: Capability) => T + /** Read a capability, or `undefined` when nobody provided it. */ + getOptional: (capability: Capability) => T | undefined + /** Provide a capability this plugin declared in `provides`. */ + provide: (capability: Capability, value: T) => void + /** The agents of this session (harness agents plus earlier plugins' agents). */ + agents: AgentRegistryView + session: { threadId: string } +} + +export interface PluginDefinition { + /** A stable, unique name, for example `'acme/todos'`. */ + name: string + lifetime?: PluginLifetime + /** Capabilities that an earlier plugin or harness middleware must provide. */ + requires?: ReadonlyArray + /** Capabilities this plugin provides in `setup`. */ + provides?: ReadonlyArray + /** Capabilities used when present. Never an error when missing. */ + optionalRequires?: ReadonlyArray + /** Agent outputs this plugin needs the session to have. */ + needs?: { produces?: ReadonlyArray } + setup?: ( + ctx: PluginSetupContext, + ) => PluginContributions | void | Promise +} + +const PLUGIN_BRAND = Symbol.for('tanstack.ai.harnessPlugin') + +/** A plugin made with {@link definePlugin}. */ +export type HarnessPlugin< + TDefinition extends PluginDefinition = PluginDefinition, +> = Readonly & { readonly [PLUGIN_BRAND]: true } + +/** + * Define a harness plugin. `setup` returns what the plugin contributes, and + * closures in it can use the resources it acquired. + * + * @example + * ```ts + * const today = definePlugin({ + * name: 'acme/today', + * setup: () => ({ prompts: [`Today is ${new Date().toDateString()}.`] }), + * }) + * ``` + */ +export function definePlugin( + definition: TDefinition, +): HarnessPlugin { + if (definition.name.trim() === '') { + throw new Error('definePlugin requires a non-empty name') + } + return Object.freeze({ ...definition, [PLUGIN_BRAND]: true as const }) +} + +/** A tool, prompt, or middleware with the plugin that contributed it. */ +interface Owned { + value: T + owner: string +} + +/** Everything a set of mounted plugins contributes, plus its cleanup. */ +export interface MountedPlugins { + tools: Array + prompts: Array + middleware: Array + generationMiddleware: Array + /** + * Chat middleware that provides every plugin capability to each chat run, + * so any chat middleware can read it with `getX(ctx)`. Goes first. + */ + capabilityBridge: AnyChatMiddleware | undefined + /** The capability values, so a run mount can read session capabilities. */ + values: CapabilityValues + /** Dispose every plugin scope, newest first. Safe to call twice. */ + dispose: () => Promise +} + +export interface MountEnvironment { + threadId: string + registry: AgentRegistry + /** Names already taken by the harness itself. */ + harnessTools: ReadonlyArray + /** Capabilities provided by harness-level middleware. */ + harnessProvides: ReadonlyArray + /** Capability values from an outer mount (session plugins, for a run mount). */ + inherited?: CapabilityValues +} + +/** Capability values provided by plugins, keyed by handle. */ +export class CapabilityValues { + readonly context = { capabilities: new CapabilityRegistry() } + readonly handles: Array = [] + + constructor(private readonly parent?: CapabilityValues) {} + + provide(handle: Capability, value: T): void { + handle[1](this.context, value) + this.handles.push(handle) + } + + get(handle: Capability): T | undefined { + if (handle.has(this.context)) return handle[0](this.context) + return this.parent?.get(handle) + } + + /** Every handle provided here and in the parent chain. */ + all(): Array { + return [...(this.parent?.all() ?? []), ...this.handles] + } +} + +function checkCapabilityOrder( + plugins: ReadonlyArray, + available: Set, +): void { + const byName = new Map() + for (const handle of available) byName.set(handle.capabilityName, handle) + const names = new Set() + for (const plugin of plugins) { + if (names.has(plugin.name)) { + throw new Error(`Duplicate plugin: ${plugin.name}`) + } + names.add(plugin.name) + for (const handle of [ + ...(plugin.requires ?? []), + ...(plugin.provides ?? []), + ...(plugin.optionalRequires ?? []), + ]) { + const other = byName.get(handle.capabilityName) + if (other && other !== handle) { + throw new Error( + `Two different capabilities are named "${handle.capabilityName}" (plugin ${plugin.name}). Capability names must be unique.`, + ) + } + byName.set(handle.capabilityName, handle) + } + for (const handle of plugin.requires ?? []) { + if (!available.has(handle)) { + throw new Error( + `Plugin ${plugin.name} requires capability "${handle.capabilityName}", but no earlier plugin or harness middleware provides it. Add a provider before ${plugin.name}.`, + ) + } + } + for (const handle of plugin.provides ?? []) available.add(handle) + } +} + +function promptOf(prompt: string | PluginPrompt, owner: string, index: number) { + return typeof prompt === 'string' + ? { id: `${owner}#${index}`, text: prompt } + : prompt +} + +/** + * Set up plugins in order and collect what they contribute. + * + * Before any `setup` runs: duplicate plugin names, capability name clashes, + * and missing providers are errors. If a `setup` throws, or a contribution + * clashes (the same tool, prompt id, or agent name from two owners), every + * plugin set up so far is disposed newest first, and no chat turn starts. + */ +export async function mountPlugins( + plugins: ReadonlyArray, + env: MountEnvironment, +): Promise { + const available = new Set([ + ...env.harnessProvides, + ...(env.inherited?.all() ?? []), + ]) + checkCapabilityOrder(plugins, available) + + const values = new CapabilityValues(env.inherited) + const scopes: Array = [] + const tools: Array> = env.harnessTools.map((tool) => ({ + value: tool, + owner: 'the harness', + })) + const prompts: Array> = [] + const middleware: Array = [] + const generationMiddleware: Array = [] + + try { + for (const plugin of plugins) { + const scope = new ResourceScope() + scopes.push(scope) + const provided = new Set() + const declared = new Set([ + ...(plugin.requires ?? []), + ...(plugin.optionalRequires ?? []), + ]) + const contributions = await plugin.setup?.({ + resources: { + acquire: (open, close) => scope.acquire(open, close), + signal: scope.signal, + }, + get: (handle) => { + const value = values.get(handle) + if (value === undefined && !declared.has(handle)) { + throw new Error( + `Plugin ${plugin.name} reads capability "${handle.capabilityName}" without declaring it in requires.`, + ) + } + if (value === undefined) { + throw new Error( + `Capability "${handle.capabilityName}" was requested by ${plugin.name} but never provided.`, + ) + } + return value + }, + getOptional: (handle) => values.get(handle), + provide: (handle, value) => { + if (!(plugin.provides ?? []).includes(handle)) { + throw new Error( + `Plugin ${plugin.name} provides "${handle.capabilityName}" without declaring it in provides.`, + ) + } + values.provide(handle, value) + provided.add(handle) + }, + agents: env.registry, + session: { threadId: env.threadId }, + }) + for (const handle of plugin.provides ?? []) { + if (!provided.has(handle)) { + throw new Error( + `Plugin ${plugin.name} declares capability "${handle.capabilityName}" in provides but did not provide it in setup.`, + ) + } + } + if (!contributions) continue + for (const tool of contributions.tools ?? []) { + const clash = tools.find((entry) => entry.value.name === tool.name) + if (clash) { + throw new Error( + `Duplicate tool "${tool.name}": first owner ${clash.owner}, second owner ${plugin.name}. No run started.`, + ) + } + tools.push({ value: tool, owner: plugin.name }) + } + ;(contributions.prompts ?? []).forEach((prompt, index) => { + const section = promptOf(prompt, plugin.name, index) + const clash = prompts.find((entry) => entry.value.id === section.id) + if (clash) { + throw new Error( + `Duplicate prompt id "${section.id}": first owner ${clash.owner}, second owner ${plugin.name}.`, + ) + } + prompts.push({ value: section, owner: plugin.name }) + }) + middleware.push(...(contributions.middleware ?? [])) + generationMiddleware.push(...(contributions.generationMiddleware ?? [])) + for (const agent of contributions.agents ?? []) { + env.registry.add(agent, plugin.name) + } + } + for (const plugin of plugins) { + for (const produces of plugin.needs?.produces ?? []) { + if (!env.registry.find({ produces })) { + throw new Error( + `Plugin ${plugin.name} needs an agent that produces "${produces}", but the session has none. Add one to agents or subagents.`, + ) + } + } + } + } catch (error) { + try { + await disposeAll(scopes) + } catch (cleanupError) { + throw new AggregateError( + [error, cleanupError], + 'Plugin setup failed, and cleanup also failed', + { cause: error }, + ) + } + throw error + } + + const bridgeHandles = values.handles + const capabilityBridge: AnyChatMiddleware | undefined = + bridgeHandles.length > 0 + ? { + name: 'harness:plugin-capabilities', + provides: bridgeHandles, + setup(ctx) { + for (const handle of bridgeHandles) { + handle[1](ctx, values.get(handle)) + } + }, + } + : undefined + + let disposed: Promise | undefined + return { + tools: tools + .filter((entry) => entry.owner !== 'the harness') + .map((entry) => entry.value), + prompts: prompts.map((entry) => entry.value.text), + middleware, + generationMiddleware, + capabilityBridge, + values, + dispose: () => (disposed ??= disposeAll(scopes)), + } +} diff --git a/packages/ai-harness/src/resources.ts b/packages/ai-harness/src/resources.ts new file mode 100644 index 0000000000..41c093a521 --- /dev/null +++ b/packages/ai-harness/src/resources.ts @@ -0,0 +1,84 @@ +/** + * Resources one plugin owns for one lifetime (a session or a run). + * + * - `acquire(open, close)` registers `close` only after `open` succeeds. + * - If the scope closes while `open` is still running, the late resource is + * closed as soon as it arrives. + * - `dispose()` closes resources newest first, runs every closer even when + * one throws, and returns the same promise to every caller. + */ +export class ResourceScope { + private readonly controller = new AbortController() + readonly signal: AbortSignal = this.controller.signal + private readonly closers: Array<() => unknown> = [] + private readonly pending = new Set>() + private closing?: Promise + private closed = false + + async acquire( + open: () => T | Promise, + close: (resource: T) => unknown, + ): Promise { + this.signal.throwIfAborted() + let release!: () => void + const pending = new Promise((resolve) => { + release = resolve + }) + this.pending.add(pending) + try { + const resource = await open() + if (this.closed) { + await close(resource) + throw this.signal.reason + } + this.closers.push(() => close(resource)) + return resource + } finally { + this.pending.delete(pending) + release() + } + } + + dispose(): Promise { + if (this.closing) return this.closing + this.closed = true + this.controller.abort(new Error('Resource scope disposed')) + // Defer, so a closer that calls dispose() again gets the same promise. + this.closing = Promise.resolve().then(async () => { + await Promise.all(this.pending) + const errors: Array = [] + for (const close of this.closers.reverse()) { + try { + await close() + } catch (error) { + errors.push(error) + } + } + this.closers.length = 0 + if (errors.length > 0) { + throw new AggregateError(errors, 'Resource cleanup failed') + } + }) + return this.closing + } +} + +/** + * Dispose scopes newest first. Every scope is disposed even when one fails. + * Throws an `AggregateError` with every failure, after all scopes ran. + */ +export async function disposeAll( + scopes: ReadonlyArray, +): Promise { + const errors: Array = [] + for (const scope of [...scopes].reverse()) { + try { + await scope.dispose() + } catch (error) { + errors.push(error) + } + } + if (errors.length > 0) { + throw new AggregateError(errors, 'Plugin cleanup failed') + } +} diff --git a/packages/ai-harness/src/session.ts b/packages/ai-harness/src/session.ts new file mode 100644 index 0000000000..9e5a74b5c3 --- /dev/null +++ b/packages/ai-harness/src/session.ts @@ -0,0 +1,870 @@ +import { + EventType, + RUN_CANCEL_REASON, + chat, + compactForModel, + createSubagentId, + runAgentStream, + validateWithStandardSchema, +} from '@tanstack/ai' +import { withPersistence } from '@tanstack/ai-persistence' +import { AgentRegistry } from './agents' +import { SessionFeed } from './feed' +import { OperationImpl } from './operation' +import { mountPlugins } from './plugins' +import { HARNESS_EVENTS } from './types' +import type { + AnyChatMiddleware, + Interrupt, + ModelMessage, + RunAgentResumeItem, + StreamChunk, + SubagentBinding, +} from '@tanstack/ai' +import type { InboxEntry, InboxStore } from '@tanstack/ai-persistence' +import type { + AgentInputOf, + AgentRegistryView, + AgentResultOf, + AnyAgent, +} from './agents' +import type { AnyHarness, HarnessAgentsOf } from './define' +import type { HarnessPersistence } from './host' +import type { HarnessPlugin, MountedPlugins } from './plugins' +import type { + BusyPolicy, + ChatTurnResult, + Cursor, + HarnessInput, + Operation, + Principal, + Receipt, + SessionEvent, + UserInput, +} from './types' + +/** Options for running an agent from code. */ +export interface AgentRunOptions { + /** + * How the main model learns about the result on its next turn: + * `'reference'` (default) adds a short note to the transcript, `'none'` adds + * nothing. + */ + attach?: 'reference' | 'none' +} + +/** Options for starting an agent in the background. */ +export interface AgentStartOptions extends AgentRunOptions { + /** When the agent finishes, start a new chat turn with its result. */ + wake?: boolean +} + +type RunArgs = + AgentInputOf extends undefined + ? [input?: undefined, options?: TOptions] + : [input: AgentInputOf, options?: TOptions] + +/** Run one agent of a session, typed from its definition. */ +export interface AgentHandle { + run: ( + ...args: RunArgs + ) => Operation> + start: ( + ...args: RunArgs + ) => Operation> +} + +/** `session.agents`: one typed handle per registered agent name. */ +export type AgentHandles = { + [TAgent in HarnessAgentsOf as TAgent['name']]: AgentHandle +} + +/** What a session looks like right now. */ +export interface SessionSnapshot { + threadId: string + /** + * `running`: a chat turn runs. `requires_action`: the last turn stopped for + * outside input. `idle`: ready for a prompt (agents may still run). + */ + status: 'idle' | 'running' | 'requires_action' + activeOperations: Array<{ id: string; kind: string; agent?: string }> + queuedTurns: number + pendingInterrupts: Array + /** The cursor of the newest event. */ + cursor: Cursor +} + +/** What the host hands a new session. */ +export interface SessionDependencies { + harness: AnyHarness + threadId: string + persistence: HarnessPersistence + inbox: InboxStore + principal?: Principal + onClose: () => void +} + +interface QueuedTurn { + operation: OperationImpl + message?: UserInput + resume?: Array + parentRunId?: string + inputId?: string +} + +function createInputId(): string { + return `in-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 9)}` +} + +function createMessageId(): string { + return `msg-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 9)}` +} + +function customEvent( + name: string, + value: Record, +): StreamChunk { + return { type: EventType.CUSTOM, name, value, timestamp: Date.now() } +} + +/** A short transcript note about an agent result, for the next turn. */ +function referenceNote(agent: string, result: unknown): string { + const body = + typeof result === 'string' + ? result + : JSON.stringify(compactForModel(result)) + const clipped = body.length > 2000 ? `${body.slice(0, 2000)}...` : body + return `[${agent} finished] ${clipped}` +} + +/** + * A live harness session: one conversation (`threadId`) with its plugins, + * operations, inbox, and event stream. Open one with `host.open()`. + */ +export class HarnessSession { + readonly threadId: string + readonly agents: AgentHandles + /** The agents this session can run, for discovery. */ + readonly registry: AgentRegistryView + + private readonly harness: THarness + private readonly persistence: HarnessPersistence + private readonly inbox: InboxStore + private readonly principal: Principal | undefined + private readonly feed = new SessionFeed() + private readonly agentRegistry = new AgentRegistry() + private readonly operations = new Map>() + private readonly queue: Array = [] + /** + * Messages for the running turn. A steer that the turn never reached (no + * further model call) runs as the next turn instead. + */ + private readonly steerQueue: Array<{ + inputId: string + message: UserInput + operation?: OperationImpl + }> = [] + private readonly pendingNotes: Array = [] + private activeTurn: OperationImpl | undefined + private interrupted: + | { runId: string; interrupts: Array } + | undefined + private plugins: ReadonlyArray = [] + private sessionPlugins: MountedPlugins | undefined + private closing: Promise | undefined + private readonly onClose: () => void + + constructor(deps: SessionDependencies) { + this.harness = deps.harness as THarness + this.threadId = deps.threadId + this.persistence = deps.persistence + this.inbox = deps.inbox + this.principal = deps.principal + this.onClose = deps.onClose + this.registry = this.agentRegistry + for (const agent of this.harness.agents ?? []) { + this.agentRegistry.add(agent, 'the harness') + } + for (const agent of this.harness.subagents?.agents ?? []) { + this.agentRegistry.add(agent, 'the harness') + } + this.agents = new Proxy({} as AgentHandles, { + get: (_target, name) => { + if (typeof name !== 'string') return undefined + return { + run: (input?: unknown, options?: AgentRunOptions) => + this.runAgent(name, input, { ...options, wake: false }), + start: (input?: unknown, options?: AgentStartOptions) => + this.runAgent(name, input, options ?? {}), + } + }, + }) + } + + /** @internal Mount session plugins and replay inputs left in the inbox. */ + async open(): Promise { + this.plugins = this.harness.plugins?.() ?? [] + this.sessionPlugins = await mountPlugins( + this.plugins.filter( + (plugin) => (plugin.lifetime ?? 'session') === 'session', + ), + { + threadId: this.threadId, + registry: this.agentRegistry, + harnessTools: this.harness.tools ?? [], + harnessProvides: (this.harness.middleware ?? []).flatMap( + (middleware) => middleware.provides ?? [], + ), + }, + ) + await this.recoverInbox() + } + + // =========================== + // Inputs + // =========================== + + /** Start a chat turn, or queue it while one runs (see `busy`). */ + prompt( + message: UserInput, + options?: { busy?: BusyPolicy }, + ): Operation { + const busy = options?.busy ?? this.harness.busy ?? 'queue' + const inputId = createInputId() + const operation = this.createTurnOperation() + void this.accept(inputId, { op: 'prompt', message, busy }).then(() => { + if (this.activeTurn && busy === 'reject') { + this.reject(inputId, 'busy') + operation.fail('failed', new Error('A chat turn is already running.')) + return + } + if (this.activeTurn && busy === 'steer') { + this.steerQueue.push({ inputId, message, operation }) + return + } + this.enqueueTurn({ operation, message, inputId }) + }) + return operation + } + + /** Add a message to the running turn at its next model call. */ + async steer(message: UserInput): Promise { + const inputId = createInputId() + await this.accept(inputId, { op: 'steer', message }) + if (!this.activeTurn) { + const operation = this.createTurnOperation() + this.enqueueTurn({ operation, message, inputId }) + return { inputId, status: 'accepted', operationId: operation.id } + } + this.steerQueue.push({ inputId, message }) + return { inputId, status: 'accepted', operationId: this.activeTurn.id } + } + + /** Run a turn after the current work settles. */ + async followUp(message: UserInput): Promise { + const inputId = createInputId() + await this.accept(inputId, { op: 'followUp', message }) + const operation = this.createTurnOperation() + const status = + this.activeTurn || this.queue.length > 0 ? 'queued' : 'accepted' + this.enqueueTurn({ operation, message, inputId }) + return { inputId, status, operationId: operation.id } + } + + /** + * Answer the interrupts of the last turn. One resume must answer every open + * interrupt of that turn (the AG-UI rule). + */ + async resolve(resume: Array): Promise { + const inputId = createInputId() + await this.accept(inputId, { op: 'resolve', resume }) + if (!this.interrupted) { + this.reject(inputId, 'no_pending_interrupts') + return { inputId, status: 'rejected', reason: 'no_pending_interrupts' } + } + if (this.activeTurn) { + this.reject(inputId, 'busy') + return { inputId, status: 'rejected', reason: 'busy' } + } + const parentRunId = this.interrupted.runId + this.interrupted = undefined + const operation = this.createTurnOperation() + this.enqueueTurn({ operation, resume, parentRunId, inputId }) + return { inputId, status: 'accepted', operationId: operation.id } + } + + /** Cancel one operation, or the running chat turn. */ + async cancel(operationId?: string): Promise { + const inputId = createInputId() + await this.accept(inputId, { op: 'cancel', operationId }) + const target = operationId + ? this.operations.get(operationId) + : this.activeTurn + if (!target || target.isSettled()) { + this.reject(inputId, 'not_running') + return { inputId, status: 'rejected', reason: 'not_running' } + } + const queued = this.queue.findIndex((turn) => turn.operation === target) + if (queued >= 0) { + this.queue.splice(queued, 1) + target.fail('cancelled', new Error('Cancelled before it started.')) + this.publishFinished(target) + } else { + target.abortController.abort(RUN_CANCEL_REASON) + } + await this.inbox.markApplied(inputId, target.id) + return { inputId, status: 'accepted', operationId: target.id } + } + + // =========================== + // Events and state + // =========================== + + /** The ordered events of every operation, from `from` (exclusive). */ + events(options?: { + from?: Cursor + signal?: AbortSignal + }): AsyncIterable { + return this.feed.read(options ?? {}) + } + + snapshot(): SessionSnapshot { + const active = [...this.operations.values()].filter( + (operation) => !operation.isSettled(), + ) + return { + threadId: this.threadId, + status: this.activeTurn + ? 'running' + : this.interrupted + ? 'requires_action' + : 'idle', + activeOperations: active.map((operation) => ({ + id: operation.id, + kind: operation.kind, + ...(operation.agent ? { agent: operation.agent } : {}), + })), + queuedTurns: this.queue.length, + pendingInterrupts: this.interrupted?.interrupts ?? [], + cursor: this.feed.head(), + } + } + + /** + * Stop every running operation, wait for them, then dispose session plugins. + * Safe to call twice. + */ + close(): Promise { + this.closing ??= (async () => { + for (const turn of this.queue.splice(0)) { + turn.operation.fail('cancelled', new Error('Session closed.')) + } + const running = [...this.operations.values()].filter( + (operation) => !operation.isSettled(), + ) + for (const operation of running) { + operation.abortController.abort(RUN_CANCEL_REASON) + } + await Promise.allSettled( + running.map((operation) => Promise.resolve(operation)), + ) + try { + await this.sessionPlugins?.dispose() + } finally { + this.feed.close() + this.onClose() + } + })() + return this.closing + } + + // =========================== + // Chat turns + // =========================== + + private createTurnOperation(): OperationImpl { + const operation = new OperationImpl( + 'chat', + this.feed, + (target) => this.cancel(target.id), + ) + this.operations.set(operation.id, operation as OperationImpl) + return operation + } + + private enqueueTurn(turn: QueuedTurn): void { + this.queue.push(turn) + this.drain() + } + + private drain(): void { + if (this.activeTurn || this.closing) return + const next = this.queue.shift() + if (!next) return + this.activeTurn = next.operation + void this.runTurn(next).finally(() => { + this.activeTurn = undefined + this.drain() + }) + } + + /** Middleware that adds queued steer messages before each model call. */ + private steering(): AnyChatMiddleware { + return { + name: 'harness:steering', + onConfig: (ctx, config) => { + if (ctx.phase !== 'beforeModel' || this.steerQueue.length === 0) { + return undefined + } + const steers = this.steerQueue.splice(0) + const running = this.activeTurn + for (const steer of steers) { + // A prompt that joined the running turn settles with that turn. + if (steer.operation && running) { + const joined = steer.operation + this.operations.delete(joined.id) + running.then( + (result) => joined.finish('completed', result), + (error: unknown) => joined.fail('failed', error), + ) + } + void this.inbox.markApplied(steer.inputId, ctx.runId) + this.feed.publish( + ctx.runId, + customEvent(HARNESS_EVENTS.inputApplied, { + inputId: steer.inputId, + operationId: ctx.runId, + }), + ) + } + return { + messages: [ + ...config.messages, + ...steers.map( + (steer): ModelMessage => ({ + id: createMessageId(), + role: 'user', + content: steer.message, + }), + ), + ], + } + }, + } + } + + private binding(): SubagentBinding { + return { + generationMiddleware: this.sessionPlugins?.generationMiddleware ?? [], + } + } + + private async runTurn(turn: QueuedTurn): Promise { + const { operation } = turn + operation.setStatus('running') + if (turn.inputId) await this.applied(turn.inputId, operation.id) + this.publishStarted(operation) + + let runPlugins: MountedPlugins | undefined + let text = '' + let interrupts: Array | undefined + let failure: string | undefined + try { + await this.flushNotes() + const perRun = this.plugins.filter((plugin) => plugin.lifetime === 'run') + if (perRun.length > 0) { + runPlugins = await mountPlugins(perRun, { + threadId: this.threadId, + registry: this.agentRegistry.fork(), + harnessTools: [ + ...(this.harness.tools ?? []), + ...(this.sessionPlugins?.tools ?? []), + ], + harnessProvides: (this.harness.middleware ?? []).flatMap( + (middleware) => middleware.provides ?? [], + ), + ...(this.sessionPlugins + ? { inherited: this.sessionPlugins.values } + : {}), + }) + } + const session = this.sessionPlugins + const bridges = [ + session?.capabilityBridge, + runPlugins?.capabilityBridge, + ].filter((bridge): bridge is AnyChatMiddleware => bridge !== undefined) + const subagents = this.harness.subagents + const stream = chat({ + adapter: this.harness.adapter, + messages: + turn.message !== undefined + ? [{ id: createMessageId(), role: 'user', content: turn.message }] + : [], + systemPrompts: [ + ...(this.harness.systemPrompts ?? []), + ...(session?.prompts ?? []), + ...(runPlugins?.prompts ?? []), + ], + tools: [ + ...(this.harness.tools ?? []), + ...(session?.tools ?? []), + ...(runPlugins?.tools ?? []), + ], + middleware: [ + ...bridges, + withPersistence(this.persistence), + ...(this.harness.middleware ?? []), + ...(session?.middleware ?? []), + ...(runPlugins?.middleware ?? []), + this.steering(), + ], + ...(subagents + ? { subagents: { ...subagents, binding: this.binding() } } + : {}), + ...(this.harness.agentLoopStrategy + ? { agentLoopStrategy: this.harness.agentLoopStrategy } + : {}), + ...(this.harness.modelOptions !== undefined + ? { modelOptions: this.harness.modelOptions } + : {}), + ...(this.harness.interrupts + ? { interrupts: this.harness.interrupts } + : {}), + ...(this.harness.context !== undefined + ? { context: this.harness.context } + : {}), + threadId: this.threadId, + runId: operation.id, + ...(turn.parentRunId ? { parentRunId: turn.parentRunId } : {}), + ...(turn.resume ? { resume: turn.resume } : {}), + abortController: operation.abortController, + stream: true, + } as never) as AsyncIterable + + for await (const chunk of stream) { + operation.publish(chunk) + if ( + chunk.type === EventType.TEXT_MESSAGE_CONTENT && + !('subagentRunId' in chunk && chunk.subagentRunId) + ) { + text += chunk.delta + } + if ( + chunk.type === EventType.RUN_FINISHED && + chunk.outcome?.type === 'interrupt' + ) { + interrupts = chunk.outcome.interrupts + } + if (chunk.type === EventType.RUN_ERROR) failure = chunk.message + } + } catch (error) { + failure = error instanceof Error ? error.message : String(error) + if (!operation.abortController.signal.aborted) { + operation.publish({ + type: EventType.RUN_ERROR, + message: failure, + timestamp: Date.now(), + }) + } + } finally { + await runPlugins?.dispose().catch(() => {}) + } + + if (operation.abortController.signal.aborted) { + operation.fail('cancelled', new Error('Cancelled.')) + } else if (failure !== undefined) { + operation.fail('failed', new Error(failure)) + } else if (interrupts && interrupts.length > 0) { + this.interrupted = { runId: operation.id, interrupts } + operation.finish('interrupted', { text, interrupts }) + } else { + operation.finish('completed', { text }) + } + this.publishFinished(operation) + // Steers the turn never reached run next, before other queued turns. + this.queue.unshift( + ...this.steerQueue.splice(0).map((steer) => ({ + operation: steer.operation ?? this.createTurnOperation(), + message: steer.message, + inputId: steer.inputId, + })), + ) + } + + // =========================== + // Agents + // =========================== + + private runAgent( + name: string, + input: unknown, + options: AgentStartOptions, + ): OperationImpl { + const operation = new OperationImpl( + 'agent', + this.feed, + (target) => this.cancel(target.id), + name, + ) + this.operations.set(operation.id, operation) + void this.executeAgent(operation, name, input, options) + return operation + } + + private async executeAgent( + operation: OperationImpl, + name: string, + input: unknown, + options: AgentStartOptions, + ): Promise { + const inputId = createInputId() + await this.accept(inputId, { + op: 'agent', + agent: name, + input, + ...(options.wake ? { detached: true } : {}), + }) + const agent: AnyAgent | undefined = this.agentRegistry.get(name) + if (!agent) { + this.reject(inputId, 'unknown_agent') + operation.fail('failed', new Error(`Unknown agent: ${name}`)) + return + } + let checkedInput: unknown = input + if (agent.inputSchema !== undefined) { + const checked = await validateWithStandardSchema( + agent.inputSchema, + input ?? {}, + ) + if (!checked.success) { + const reason = `Input validation failed for agent ${name}: ${checked.issues + .map((issue) => issue.message) + .join(', ')}` + this.reject(inputId, 'invalid_input') + operation.fail('failed', new Error(reason)) + return + } + checkedInput = checked.data + } + + operation.setStatus('running') + await this.applied(inputId, operation.id) + const runs = this.persistence.stores.runs + await runs?.createOrResume({ + runId: operation.id, + threadId: this.threadId, + startedAt: Date.now(), + kind: 'agent', + agent: name, + ...(this.principal ? { principal: { id: this.principal.id } } : {}), + }) + operation.publish({ + type: EventType.RUN_STARTED, + runId: operation.id, + threadId: this.threadId, + timestamp: Date.now(), + }) + this.publishStarted(operation) + + let text = '' + let result: unknown + let failure: string | undefined + let subagentRunId = '' + try { + const messages = await this.persistence.stores.messages?.loadThread( + this.threadId, + ) + subagentRunId = createSubagentId() + const stream = runAgentStream( + agent, + { + input: checkedInput, + messages: messages ?? [], + threadId: `${this.threadId}:${name}`, + runId: `${operation.id}:${subagentRunId}`, + parentRunId: operation.id, + subagentRunId, + abortSignal: operation.abortController.signal, + }, + undefined, + undefined, + this.binding(), + ) + for await (const chunk of stream) { + operation.publish(chunk) + if ( + chunk.type === EventType.TEXT_MESSAGE_CONTENT && + 'subagentRunId' in chunk && + chunk.subagentRunId === subagentRunId + ) { + text += chunk.delta + } + if ( + chunk.type === EventType.SUBAGENT_FINISHED && + chunk.subagentRunId === subagentRunId + ) { + result = chunk.result + } + if ( + chunk.type === EventType.SUBAGENT_ERROR && + chunk.subagentRunId === subagentRunId + ) { + failure = chunk.message + } + } + } catch (error) { + failure = error instanceof Error ? error.message : String(error) + } + + const value = result !== undefined ? result : text + if (operation.abortController.signal.aborted) { + operation.publish({ + type: EventType.RUN_FINISHED, + runId: operation.id, + threadId: this.threadId, + outcome: { type: 'cancelled' }, + timestamp: Date.now(), + } as StreamChunk) + await runs?.update(operation.id, { + status: 'aborted', + finishedAt: Date.now(), + }) + operation.fail('cancelled', new Error('Cancelled.')) + } else if (failure !== undefined) { + operation.publish({ + type: EventType.RUN_ERROR, + message: failure, + timestamp: Date.now(), + }) + await runs?.update(operation.id, { + status: 'failed', + finishedAt: Date.now(), + error: { message: failure }, + }) + operation.fail('failed', new Error(failure)) + } else { + operation.publish({ + type: EventType.RUN_FINISHED, + runId: operation.id, + threadId: this.threadId, + result: value, + timestamp: Date.now(), + } as StreamChunk) + await runs?.update(operation.id, { + status: 'completed', + finishedAt: Date.now(), + result: compactForModel(value), + }) + if ((options.attach ?? 'reference') === 'reference') { + this.pendingNotes.push(referenceNote(name, value)) + if (!this.activeTurn) await this.flushNotes() + } + operation.finish('completed', value) + if (options.wake) { + void this.followUp( + `Background agent ${name} finished: ${referenceNote(name, value)}`, + ) + } + } + this.publishFinished(operation) + } + + /** Write queued agent notes to the transcript while no turn is writing it. */ + private async flushNotes(): Promise { + const messages = this.persistence.stores.messages + if (!messages || this.pendingNotes.length === 0) return + const notes = this.pendingNotes.splice(0) + const history = await messages.loadThread(this.threadId) + await messages.saveThread(this.threadId, [ + ...history, + ...notes.map( + (note): ModelMessage => ({ + id: createMessageId(), + role: 'assistant', + content: note, + }), + ), + ]) + } + + // =========================== + // Inbox and lifecycle events + // =========================== + + private async accept(inputId: string, input: HarnessInput): Promise { + await this.inbox.append({ + inputId, + threadId: this.threadId, + input, + createdAt: Date.now(), + ...(this.principal ? { principal: { id: this.principal.id } } : {}), + }) + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.inputAccepted, { inputId, op: input.op }), + ) + } + + private async applied(inputId: string, operationId: string): Promise { + await this.inbox.markApplied(inputId, operationId) + this.feed.publish( + operationId, + customEvent(HARNESS_EVENTS.inputApplied, { inputId, operationId }), + ) + } + + private reject(inputId: string, reason: string): void { + void this.inbox.markRejected(inputId, reason) + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.inputRejected, { inputId, reason }), + ) + } + + private publishStarted( + operation: OperationImpl | OperationImpl, + ): void { + operation.publish( + customEvent(HARNESS_EVENTS.operationStarted, { + operationId: operation.id, + kind: operation.kind, + ...(operation.agent ? { agent: operation.agent } : {}), + }), + ) + } + + private publishFinished( + operation: OperationImpl | OperationImpl, + ): void { + operation.publish( + customEvent(HARNESS_EVENTS.operationFinished, { + operationId: operation.id, + status: operation.status(), + }), + ) + } + + /** Re-run turns that were accepted but never applied before a restart. */ + private async recoverInbox(): Promise { + const pending: Array = await this.inbox.listPending( + this.threadId, + ) + for (const entry of pending) { + const input = entry.input as HarnessInput + if ( + input.op === 'prompt' || + input.op === 'followUp' || + input.op === 'steer' + ) { + const operation = this.createTurnOperation() + this.enqueueTurn({ + operation, + message: input.message, + inputId: entry.inputId, + }) + } else { + this.reject(entry.inputId, 'expired_on_restart') + } + } + } +} diff --git a/packages/ai-harness/src/types.ts b/packages/ai-harness/src/types.ts new file mode 100644 index 0000000000..5d6b766ae9 --- /dev/null +++ b/packages/ai-harness/src/types.ts @@ -0,0 +1,99 @@ +import type { + ContentPart, + Interrupt, + RunAgentResumeItem, + StreamChunk, +} from '@tanstack/ai' + +/** A user message: plain text, or content parts (text, images, files). */ +export type UserInput = string | Array + +/** + * What a `prompt` does when a chat turn is already running: + * - `queue` (default): runs as a new turn after the current one settles. + * - `steer`: joins the running turn at its next model call. + * - `reject`: refuses the prompt. + */ +export type BusyPolicy = 'queue' | 'steer' | 'reject' + +/** Kinds of work a session runs. */ +export type OperationKind = 'chat' | 'agent' | 'command' | 'compact' + +export type OperationStatus = + | 'accepted' + | 'running' + | 'interrupted' + | 'completed' + | 'failed' + | 'cancelled' + +/** The immediate answer to an input. `applied` arrives later as an event. */ +export interface Receipt { + inputId: string + status: 'accepted' | 'queued' | 'rejected' + operationId?: string + reason?: string +} + +/** Opaque position in a session's event stream. */ +export type Cursor = string + +/** One event in a session's ordered stream. */ +export interface SessionEvent { + cursor: Cursor + /** The operation that produced the event. */ + operationId: string + event: StreamChunk +} + +/** An input a client sends to a session. Stored in the inbox. */ +export type HarnessInput = + | { op: 'prompt'; message: UserInput; busy?: BusyPolicy } + | { op: 'steer'; message: UserInput } + | { op: 'followUp'; message: UserInput } + | { op: 'resolve'; resume: Array } + | { op: 'agent'; agent: string; input?: unknown; detached?: boolean } + | { op: 'cancel'; operationId?: string } + +/** Who sent an input, from the host's `authorize`. */ +export interface Principal { + id: string + name?: string +} + +/** What a chat turn operation resolves to. */ +export interface ChatTurnResult { + /** The main model's text for this turn (child agent text excluded). */ + text: string + /** Set when the turn stopped for outside input. */ + interrupts?: Array +} + +/** + * One accepted unit of work in a session. Await it for its result. + */ +export interface Operation extends PromiseLike { + /** Also the AG-UI `runId` of the operation. */ + readonly id: string + readonly kind: OperationKind + /** The agent name, for `kind: 'agent'`. */ + readonly agent?: string + status: () => OperationStatus + /** This operation's events, from `from` (exclusive) onward. */ + events: (options?: { + from?: Cursor + signal?: AbortSignal + }) => AsyncIterable + /** This operation's raw AG-UI chunks, for existing transports. */ + stream: (options?: { signal?: AbortSignal }) => AsyncIterable + cancel: (reason?: string) => Promise +} + +/** Names of the `CUSTOM` events a harness session adds to the stream. */ +export const HARNESS_EVENTS = { + operationStarted: 'harness.operation.started', + operationFinished: 'harness.operation.finished', + inputAccepted: 'harness.input.accepted', + inputApplied: 'harness.input.applied', + inputRejected: 'harness.input.rejected', +} as const diff --git a/packages/ai-harness/tests/define.test.ts b/packages/ai-harness/tests/define.test.ts new file mode 100644 index 0000000000..7e46b4f603 --- /dev/null +++ b/packages/ai-harness/tests/define.test.ts @@ -0,0 +1,64 @@ +import { describe, expect, it } from 'vitest' +import { z } from 'zod' +import { defineAgent } from '@tanstack/ai' +import { defineHarness, isHarnessDefinition } from '../src' +import { mockAdapter } from './helpers' + +const agent = (name: string) => + defineAgent({ + name, + description: name, + inputSchema: z.object({ q: z.string() }), + run: async () => 'ok', + }) + +describe('defineHarness', () => { + const { adapter } = mockAdapter([]) + + it('returns a frozen, branded definition', () => { + const harness = defineHarness({ name: 'acme/test', adapter }) + expect(isHarnessDefinition(harness)).toBe(true) + expect(Object.isFrozen(harness)).toBe(true) + expect(isHarnessDefinition({ name: 'x' })).toBe(false) + }) + + it('rejects an empty name', () => { + expect(() => defineHarness({ name: ' ', adapter })).toThrow( + 'non-empty name', + ) + }) + + it('rejects two different agents with one name', () => { + expect(() => + defineHarness({ + name: 'acme/test', + adapter, + agents: [agent('same')], + subagents: { agents: [agent('same')] }, + }), + ).toThrow('two different agents are named "same"') + }) + + it('allows one agent in both agents and subagents', () => { + const shared = agent('shared') + expect(() => + defineHarness({ + name: 'acme/test', + adapter, + agents: [shared], + subagents: { agents: [shared] }, + }), + ).not.toThrow() + }) + + it('rejects expose names that are not agents', () => { + expect(() => + defineHarness({ + name: 'acme/test', + adapter, + agents: [agent('real')], + expose: { agents: ['missing' as 'real'] }, + }), + ).toThrow('expose.agents names "missing"') + }) +}) diff --git a/packages/ai-harness/tests/helpers.ts b/packages/ai-harness/tests/helpers.ts new file mode 100644 index 0000000000..e8e3876082 --- /dev/null +++ b/packages/ai-harness/tests/helpers.ts @@ -0,0 +1,154 @@ +import { EventType } from '@tanstack/ai' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' + +type Reply = (options: any) => AsyncIterable | Array + +/** A text adapter whose turns are scripted. Records every call. */ +export function mockAdapter(replies: Array | Reply) { + const calls: Array = [] + const adapter: AnyTextAdapter = { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + chatStream: (options: any) => { + calls.push(options) + const reply = Array.isArray(replies) ? replies[calls.length - 1] : replies + const result = reply ? reply(options) : text('') + return (async function* () { + yield* result + })() + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } + return { adapter, calls } +} + +const now = () => Date.now() + +/** One model call that answers with `content`. */ +export function text(content: string): Array { + return [ + { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: now(), + }, + { + type: EventType.TEXT_MESSAGE_START, + messageId: 'm', + role: 'assistant', + timestamp: now(), + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm', + delta: content, + timestamp: now(), + }, + { type: EventType.TEXT_MESSAGE_END, messageId: 'm', timestamp: now() }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'stop' } }, + }, + ] +} + +/** One model call that calls a tool. */ +export function toolCall( + name: string, + args: Record, + id = 'call-1', +): Array { + return [ + { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: now(), + }, + { + type: EventType.TOOL_CALL_START, + toolCallId: id, + toolCallName: name, + timestamp: now(), + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: id, + delta: JSON.stringify(args), + timestamp: now(), + }, + { type: EventType.TOOL_CALL_END, toolCallId: id, timestamp: now() }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'tool_calls' } }, + }, + ] +} + +/** A reply that waits until the call is aborted, then ends. */ +export function untilAborted(): Reply { + return (options: any) => + (async function* () { + yield { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: now(), + } + const signal: AbortSignal | undefined = + options.abortController?.signal ?? options.request?.signal + await new Promise((resolve) => { + if (!signal || signal.aborted) return resolve() + signal.addEventListener('abort', () => resolve(), { once: true }) + }) + })() +} + +/** A promise you resolve from outside. */ +export function gate() { + let open!: () => void + const opened = new Promise((resolve) => { + open = resolve + }) + return { open, opened } +} + +/** A reply that waits for `until` before answering with `content`. */ +export function after(until: Promise, content: string): Reply { + return () => + (async function* () { + await until + yield* text(content) + })() +} + +/** The text of every message the adapter got on call `index`. */ +export function messageTexts(call: any): Array { + return (call.messages as Array).map((message) => + typeof message.content === 'string' + ? message.content + : JSON.stringify(message.content), + ) +} diff --git a/packages/ai-harness/tests/plugins.test.ts b/packages/ai-harness/tests/plugins.test.ts new file mode 100644 index 0000000000..6cbfce4cf9 --- /dev/null +++ b/packages/ai-harness/tests/plugins.test.ts @@ -0,0 +1,167 @@ +import { describe, expect, it } from 'vitest' +import { createCapability, defineAgent } from '@tanstack/ai' +import { AgentRegistry } from '../src/agents' +import { definePlugin, mountPlugins } from '../src/plugins' +import { ResourceScope } from '../src/resources' + +const env = () => ({ + threadId: 't1', + registry: new AgentRegistry(), + harnessTools: [], + harnessProvides: [], +}) + +describe('ResourceScope', () => { + it('closes resources newest first and reports every close error', async () => { + const scope = new ResourceScope() + const log: Array = [] + await scope.acquire( + () => 'a', + () => log.push('close a'), + ) + await scope.acquire( + () => 'b', + () => { + log.push('close b') + throw new Error('b broke') + }, + ) + await scope.acquire( + () => 'c', + () => { + log.push('close c') + throw new Error('c broke') + }, + ) + await expect(scope.dispose()).rejects.toBeInstanceOf(AggregateError) + expect(log).toEqual(['close c', 'close b', 'close a']) + expect(scope.signal.aborted).toBe(true) + }) + + it('closes a resource that finishes opening after dispose', async () => { + const scope = new ResourceScope() + const log: Array = [] + let finishOpen!: () => void + const opening = scope.acquire( + () => + new Promise((resolve) => { + finishOpen = () => resolve('late') + }), + () => log.push('closed late'), + ) + const disposed = scope.dispose() + finishOpen() + await expect(opening).rejects.toThrow('Resource scope disposed') + await disposed + expect(log).toEqual(['closed late']) + await expect( + scope.acquire( + () => 'x', + () => {}, + ), + ).rejects.toThrow() + }) +}) + +describe('mountPlugins', () => { + const store = createCapability<{ items: Array }>()('test.store') + + it('fails before any setup when a required capability has no provider', async () => { + let ran = false + const reader = definePlugin({ + name: 'test/reader', + requires: [store], + setup: () => { + ran = true + }, + }) + await expect(mountPlugins([reader], env())).rejects.toThrow( + 'requires capability "test.store"', + ) + expect(ran).toBe(false) + }) + + it('lets a later plugin read what an earlier plugin provides', async () => { + const provider = definePlugin({ + name: 'test/provider', + provides: [store], + setup: ({ provide }) => provide(store, { items: ['x'] }), + }) + let read: Array = [] + const reader = definePlugin({ + name: 'test/reader', + requires: [store], + setup: ({ get }) => { + read = get(store).items + }, + }) + const mounted = await mountPlugins([provider, reader], env()) + expect(read).toEqual(['x']) + expect(mounted.capabilityBridge?.provides).toEqual([store]) + }) + + it('rejects a provide that the plugin did not declare', async () => { + const sneaky = definePlugin({ + name: 'test/sneaky', + setup: ({ provide }) => provide(store, { items: [] }), + }) + await expect(mountPlugins([sneaky], env())).rejects.toThrow( + 'without declaring it in provides', + ) + }) + + it('disposes earlier plugins newest first when a setup throws', async () => { + const log: Array = [] + const opener = (name: string) => + definePlugin({ + name, + setup: async ({ resources }) => { + await resources.acquire( + () => name, + () => log.push(`close ${name}`), + ) + }, + }) + const broken = definePlugin({ + name: 'test/broken', + setup: () => { + throw new Error('setup failed') + }, + }) + await expect( + mountPlugins([opener('test/one'), opener('test/two'), broken], env()), + ).rejects.toThrow('setup failed') + expect(log).toEqual(['close test/two', 'close test/one']) + }) + + it('checks needs.produces against the agents after setup', async () => { + const needsImages = definePlugin({ + name: 'test/gallery', + needs: { produces: ['image'] }, + }) + await expect(mountPlugins([needsImages], env())).rejects.toThrow( + 'needs an agent that produces "image"', + ) + + const painter = defineAgent({ + name: 'painter', + description: 'Paints', + produces: 'image', + run: async () => 'painted', + }) + const withPainter = definePlugin({ + name: 'test/painter', + setup: () => ({ agents: [painter] }), + }) + const environment = env() + await mountPlugins([withPainter, needsImages], environment) + expect(environment.registry.find({ produces: 'image' })).toBe(painter) + }) + + it('rejects duplicate plugin names', async () => { + const a = definePlugin({ name: 'test/same' }) + await expect(mountPlugins([a, a], env())).rejects.toThrow( + 'Duplicate plugin: test/same', + ) + }) +}) diff --git a/packages/ai-harness/tests/session.test.ts b/packages/ai-harness/tests/session.test.ts new file mode 100644 index 0000000000..a46b30a72c --- /dev/null +++ b/packages/ai-harness/tests/session.test.ts @@ -0,0 +1,498 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { + EventType, + createCapability, + defineAgent, + toolDefinition, +} from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + HARNESS_EVENTS, + createHarnessHost, + defineHarness, + definePlugin, +} from '../src' +import { + after, + gate, + messageTexts, + mockAdapter, + text, + toolCall, + untilAborted, +} from './helpers' +import type { StreamChunk } from '@tanstack/ai' +import type { SessionEvent } from '../src' + +function setup() { + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + return { persistence, host } +} + +function customNames(events: Array): Array { + return events + .map((entry) => entry.event) + .filter( + (event): event is Extract => + event.type === EventType.CUSTOM, + ) + .map((event) => event.name) +} + +async function collect( + iterable: AsyncIterable, +): Promise> { + const out: Array = [] + for await (const entry of iterable) out.push(entry) + return out +} + +describe('chat turns', () => { + it('runs a prompt, returns its text, and saves the transcript', async () => { + const { host, persistence } = setup() + const { adapter } = mockAdapter([() => text('Hello there')]) + const harness = defineHarness({ name: 'test/chat', adapter }) + const session = await host.open(harness, { threadId: 't1' }) + + const turn = session.prompt('Hi') + await expect(turn).resolves.toEqual({ text: 'Hello there' }) + expect(turn.status()).toBe('completed') + + const saved = await persistence.stores.messages.loadThread('t1') + expect(saved.map((message) => message.role)).toEqual(['user', 'assistant']) + const names = customNames(await collect(turn.events({ from: '0' }))) + expect(names).toContain(HARNESS_EVENTS.inputApplied) + expect(names).toContain(HARNESS_EVENTS.operationStarted) + expect(names.at(-1)).toBe(HARNESS_EVENTS.operationFinished) + await host.close() + }) + + it('queues a second prompt and gives it the first turn as history', async () => { + const { host } = setup() + const first = gate() + const { adapter, calls } = mockAdapter([ + after(first.opened, 'one'), + () => text('two'), + ]) + const session = await host.open( + defineHarness({ name: 'test/queue', adapter }), + { + threadId: 't1', + }, + ) + + const a = session.prompt('first') + const b = session.prompt('second') + await vi.waitFor(() => expect(calls).toHaveLength(1)) + expect(session.snapshot().queuedTurns).toBe(1) + first.open() + + await expect(a).resolves.toEqual({ text: 'one' }) + await expect(b).resolves.toEqual({ text: 'two' }) + expect(messageTexts(calls[1])).toEqual(['first', 'one', 'second']) + await host.close() + }) + + it('rejects a prompt with busy: reject while a turn runs', async () => { + const { host } = setup() + const first = gate() + const { adapter, calls } = mockAdapter([after(first.opened, 'one')]) + const session = await host.open( + defineHarness({ name: 'test/reject', adapter }), + { + threadId: 't1', + }, + ) + + const a = session.prompt('first') + await vi.waitFor(() => expect(calls).toHaveLength(1)) + await expect(session.prompt('second', { busy: 'reject' })).rejects.toThrow( + 'already running', + ) + first.open() + await a + await host.close() + }) + + it('adds a steer message before the next model call of the running turn', async () => { + const { host } = setup() + const toolStarted = gate() + const releaseTool = gate() + const lookup = toolDefinition({ + name: 'lookup', + description: 'Look something up', + inputSchema: z.object({ q: z.string() }), + }).server(async () => { + toolStarted.open() + await releaseTool.opened + return { found: true } + }) + const { adapter, calls } = mockAdapter([ + () => toolCall('lookup', { q: 'x' }), + () => text('done'), + ]) + const session = await host.open( + defineHarness({ name: 'test/steer', adapter, tools: [lookup] }), + { threadId: 't1' }, + ) + + const turn = session.prompt('find x') + await toolStarted.opened + const receipt = await session.steer('also check y') + expect(receipt).toMatchObject({ status: 'accepted', operationId: turn.id }) + releaseTool.open() + + await expect(turn).resolves.toEqual({ text: 'done' }) + expect(calls).toHaveLength(2) + expect(messageTexts(calls[1]).at(-1)).toBe('also check y') + await host.close() + }) + + it('runs a steer that arrived too late as the next turn', async () => { + const { host } = setup() + const first = gate() + const { adapter, calls } = mockAdapter([ + after(first.opened, 'one'), + () => text('two'), + ]) + const session = await host.open( + defineHarness({ name: 'test/late', adapter }), + { + threadId: 't1', + }, + ) + + const a = session.prompt('first') + await vi.waitFor(() => expect(calls).toHaveLength(1)) + const joined = session.prompt('and this', { busy: 'steer' }) + first.open() + + await expect(a).resolves.toEqual({ text: 'one' }) + await expect(joined).resolves.toEqual({ text: 'two' }) + expect(messageTexts(calls[1]).at(-1)).toBe('and this') + await host.close() + }) + + it('cancels the running turn', async () => { + const { host } = setup() + const { adapter, calls } = mockAdapter([untilAborted()]) + const session = await host.open( + defineHarness({ name: 'test/cancel', adapter }), + { + threadId: 't1', + }, + ) + + const turn = session.prompt('go') + await vi.waitFor(() => expect(calls).toHaveLength(1)) + const receipt = await turn.cancel() + expect(receipt.status).toBe('accepted') + await expect(turn).rejects.toThrow('Cancelled') + expect(turn.status()).toBe('cancelled') + expect(session.snapshot().status).toBe('idle') + await host.close() + }) + + it('stops for an approval and continues after resolve', async () => { + const { host } = setup() + const execute = vi.fn(async () => ({ ok: true })) + const remove = toolDefinition({ + name: 'remove', + description: 'Remove a file', + needsApproval: true, + inputSchema: z.object({ path: z.string() }), + }).server(execute) + const { adapter } = mockAdapter([ + () => toolCall('remove', { path: 'a.txt' }, 'call_1'), + () => text('removed'), + ]) + const session = await host.open( + defineHarness({ name: 'test/approve', adapter, tools: [remove] }), + { threadId: 't1' }, + ) + + const turn = await session.prompt('remove a.txt') + expect(turn.interrupts).toHaveLength(1) + expect(session.snapshot().status).toBe('requires_action') + + const receipt = await session.resolve([ + { + interruptId: turn.interrupts![0]!.id, + status: 'resolved', + payload: true, + }, + ]) + expect(receipt.status).toBe('accepted') + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + expect(execute).toHaveBeenCalledWith({ path: 'a.txt' }, expect.anything()) + await host.close() + }) + + it('rejects resolve when nothing is waiting', async () => { + const { host } = setup() + const { adapter } = mockAdapter([() => text('hi')]) + const session = await host.open( + defineHarness({ name: 'test/noresolve', adapter }), + { + threadId: 't1', + }, + ) + const receipt = await session.resolve([ + { interruptId: 'x', status: 'resolved', payload: true }, + ]) + expect(receipt).toMatchObject({ + status: 'rejected', + reason: 'no_pending_interrupts', + }) + await host.close() + }) +}) + +describe('events', () => { + it('replays events after a cursor', async () => { + const { host } = setup() + const { adapter } = mockAdapter([() => text('one'), () => text('two')]) + const session = await host.open( + defineHarness({ name: 'test/events', adapter }), + { + threadId: 't1', + }, + ) + await session.prompt('a') + const cursor = session.snapshot().cursor + const second = session.prompt('b') + await second + + const controller = new AbortController() + const seen: Array = [] + for await (const entry of session.events({ + from: cursor, + signal: controller.signal, + })) { + seen.push(entry) + if ( + entry.event.type === EventType.CUSTOM && + entry.event.name === HARNESS_EVENTS.operationFinished + ) { + controller.abort() + } + } + expect(seen.every((entry) => Number(entry.cursor) > Number(cursor))).toBe( + true, + ) + expect(seen.some((entry) => entry.operationId === second.id)).toBe(true) + await host.close() + }) +}) + +describe('agents', () => { + const pricer = defineAgent({ + name: 'pricer', + description: 'Prices a vendor', + inputSchema: z.object({ vendor: z.string() }), + run: async (ctx) => ({ vendor: ctx.input.vendor, cents: 1200 }), + }) + + it('runs a typed agent from code and notes the result in the transcript', async () => { + const { host, persistence } = setup() + const { adapter, calls } = mockAdapter([() => text('noted')]) + const session = await host.open( + defineHarness({ name: 'test/agents', adapter, agents: [pricer] }), + { threadId: 't1' }, + ) + + const run = session.agents.pricer.run({ vendor: 'acme' }) + const result = await run + expect(result).toEqual({ vendor: 'acme', cents: 1200 }) + expect(run.kind).toBe('agent') + + const record = await persistence.stores.runs.get(run.id) + expect(record).toMatchObject({ + kind: 'agent', + agent: 'pricer', + status: 'completed', + }) + + await session.prompt('what did it cost?') + expect( + messageTexts(calls[0]).some((line) => line.includes('[pricer finished]')), + ).toBe(true) + await host.close() + }) + + it('rejects input that fails the schema', async () => { + const { host } = setup() + const { adapter } = mockAdapter([]) + const session = await host.open( + defineHarness({ name: 'test/agents-bad', adapter, agents: [pricer] }), + { threadId: 't1' }, + ) + await expect( + session.agents.pricer.run({ vendor: 1 } as never), + ).rejects.toThrow('Input validation failed for agent pricer') + await host.close() + }) + + it('starts a background agent that wakes the session when done', async () => { + const { host } = setup() + const { adapter, calls } = mockAdapter([() => text('thanks')]) + const session = await host.open( + defineHarness({ name: 'test/wake', adapter, agents: [pricer] }), + { threadId: 't1' }, + ) + + await session.agents.pricer.start({ vendor: 'acme' }, { wake: true }) + await vi.waitFor(() => expect(calls).toHaveLength(1)) + expect(messageTexts(calls[0]).at(-1)).toContain( + 'Background agent pricer finished', + ) + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + await host.close() + }) +}) + +describe('inbox', () => { + it('runs a prompt that was accepted but never applied before a restart', async () => { + const persistence = memoryPersistence() + await persistence.stores.inbox.append({ + inputId: 'in-1', + threadId: 't1', + input: { op: 'prompt', message: 'left over' }, + createdAt: Date.now(), + }) + const { adapter, calls } = mockAdapter([() => text('recovered')]) + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ name: 'test/recover', adapter }), + { + threadId: 't1', + }, + ) + + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + expect((await persistence.stores.inbox.get('in-1'))?.status).toBe('applied') + expect(messageTexts(calls[0])).toEqual(['left over']) + await host.close() + }) +}) + +describe('plugins', () => { + it('keeps session plugins across turns and sets up run plugins per turn', async () => { + const { host } = setup() + const log: Array = [] + const sessionPlugin = definePlugin({ + name: 'test/session', + setup: async ({ resources }) => { + await resources.acquire( + () => log.push('session:open'), + () => log.push('session:close'), + ) + return { prompts: ['Session prompt.'] } + }, + }) + const runPlugin = definePlugin({ + name: 'test/run', + lifetime: 'run', + setup: async ({ resources }) => { + await resources.acquire( + () => log.push('run:open'), + () => log.push('run:close'), + ) + }, + }) + const { adapter, calls } = mockAdapter([() => text('a'), () => text('b')]) + const session = await host.open( + defineHarness({ + name: 'test/plugins', + adapter, + plugins: () => [sessionPlugin, runPlugin], + }), + { threadId: 't1' }, + ) + + await session.prompt('one') + await session.prompt('two') + await session.close() + + expect(log).toEqual([ + 'session:open', + 'run:open', + 'run:close', + 'run:open', + 'run:close', + 'session:close', + ]) + expect(JSON.stringify(calls[0].systemPrompts)).toContain('Session prompt.') + }) + + it('gives plugin tools to the model and plugin capabilities to middleware', async () => { + const { host } = setup() + const clockCapability = createCapability<{ now: () => string }>()( + 'test.clock', + ) + const [getClock] = clockCapability + const seen: Array = [] + const clockPlugin = definePlugin({ + name: 'test/clock', + provides: [clockCapability], + setup: ({ provide }) => { + provide(clockCapability, { now: () => 'noon' }) + return { + tools: [ + toolDefinition({ + name: 'clock', + description: 'Tell the time', + }).server(async () => 'noon'), + ], + middleware: [ + { + name: 'test/read-clock', + requires: [clockCapability], + onStart: (ctx) => { + seen.push(getClock(ctx).now()) + }, + }, + ], + } + }, + }) + const { adapter, calls } = mockAdapter([() => text('ok')]) + const session = await host.open( + defineHarness({ + name: 'test/tools', + adapter, + plugins: () => [clockPlugin], + }), + { threadId: 't1' }, + ) + + await session.prompt('time?') + expect(calls[0].tools.map((tool: { name: string }) => tool.name)).toContain( + 'clock', + ) + expect(seen).toEqual(['noon']) + await host.close() + }) + + it('fails open when two plugins add the same tool', async () => { + const { host } = setup() + const tool = toolDefinition({ name: 'dup', description: 'x' }).server( + async () => 1, + ) + const a = definePlugin({ name: 'test/a', setup: () => ({ tools: [tool] }) }) + const b = definePlugin({ name: 'test/b', setup: () => ({ tools: [tool] }) }) + const { adapter } = mockAdapter([]) + await expect( + host.open( + defineHarness({ name: 'test/dup', adapter, plugins: () => [a, b] }), + { + threadId: 't1', + }, + ), + ).rejects.toThrow( + 'Duplicate tool "dup": first owner test/a, second owner test/b', + ) + }) +}) diff --git a/packages/ai-harness/tests/types.test-d.ts b/packages/ai-harness/tests/types.test-d.ts new file mode 100644 index 0000000000..0a9ffe74ee --- /dev/null +++ b/packages/ai-harness/tests/types.test-d.ts @@ -0,0 +1,39 @@ +import { expectTypeOf, it } from 'vitest' +import { z } from 'zod' +import { defineAgent } from '@tanstack/ai' +import { defineHarness } from '../src' +import { mockAdapter } from './helpers' +import type { HarnessSession, Operation } from '../src' + +const pricer = defineAgent({ + name: 'pricer', + description: 'Prices a vendor', + inputSchema: z.object({ vendor: z.string() }), + run: async (ctx) => ({ vendor: ctx.input.vendor, cents: 1200 }), +}) + +const writer = defineAgent({ + name: 'writer', + description: 'Writes text', + run: () => (async function* () {})(), +}) + +const harness = defineHarness({ + name: 'test/types', + adapter: mockAdapter([]).adapter, + agents: [pricer], + subagents: { agents: [writer] }, +}) + +declare const session: HarnessSession + +it('types agent input and result from the definition', () => { + expectTypeOf(session.agents.pricer.run({ vendor: 'a' })).toEqualTypeOf< + Operation<{ vendor: string; cents: number }> + >() + expectTypeOf(session.agents.writer.run()).toEqualTypeOf>() + // @ts-expect-error unknown agent + void session.agents.missing + // @ts-expect-error wrong input + void session.agents.pricer.run({ vendor: 1 }) +}) diff --git a/packages/ai-harness/tsconfig.json b/packages/ai-harness/tsconfig.json new file mode 100644 index 0000000000..29112eff9f --- /dev/null +++ b/packages/ai-harness/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "dist" + }, + "include": ["vite.config.ts", "./src", "./tests"], + "exclude": ["node_modules", "dist"] +} diff --git a/packages/ai-harness/vite.config.ts b/packages/ai-harness/vite.config.ts new file mode 100644 index 0000000000..1f3542380f --- /dev/null +++ b/packages/ai-harness/vite.config.ts @@ -0,0 +1,35 @@ +import { defineConfig, mergeConfig } from 'vitest/config' +import { tanstackViteConfig } from '@tanstack/vite-config' +import packageJson from './package.json' + +const config = defineConfig({ + test: { + name: packageJson.name, + dir: './', + watch: false, + globals: true, + environment: 'node', + include: ['src/**/*.test.ts', 'tests/**/*.test.ts'], + coverage: { + provider: 'v8', + reporter: ['text', 'json', 'html', 'lcov'], + exclude: [ + 'node_modules/', + 'dist/', + 'tests/', + '**/*.test.ts', + '**/*.config.ts', + ], + include: ['src/**/*.ts'], + }, + }, +}) + +export default mergeConfig( + config, + tanstackViteConfig({ + entry: ['./src/index.ts'], + srcDir: './src', + cjs: false, + }), +) diff --git a/packages/ai-persistence/src/index.ts b/packages/ai-persistence/src/index.ts index 3100a67c4c..eac9972b34 100644 --- a/packages/ai-persistence/src/index.ts +++ b/packages/ai-persistence/src/index.ts @@ -9,6 +9,7 @@ export { defineGenerationRunStore, defineArtifactStore, defineBlobStore, + defineInboxStore, // Run lifecycle helpers owned by @tanstack/ai, transiting through ./types so // this package's public surface stays a single import for backend authors. isTerminalRunStatus, @@ -25,6 +26,10 @@ export type { InterruptStatus, InterruptStore, MetadataStore, + // Harness session inbox + InboxEntry, + InboxStatus, + InboxStore, // Named product shapes (prefer these over a sparse bag) ChatTranscriptStores, ChatPersistenceStores, diff --git a/packages/ai-persistence/src/memory.ts b/packages/ai-persistence/src/memory.ts index 849e30acf1..fd4df0b8b6 100644 --- a/packages/ai-persistence/src/memory.ts +++ b/packages/ai-persistence/src/memory.ts @@ -14,6 +14,8 @@ import type { BlobStore, GenerationRunRecord, GenerationRunStore, + InboxEntry, + InboxStore, InterruptCommitEntry, InterruptRecord, InterruptStore, @@ -542,6 +544,45 @@ class MemoryBlobStore implements BlobStore { } } +class MemoryInboxStore implements InboxStore { + private readonly entries = new Map() + append(entry: Omit): Promise { + const existing = this.entries.get(entry.inputId) + if (existing) return Promise.resolve({ ...existing }) + const stored: InboxEntry = { ...entry, status: 'pending' } + this.entries.set(entry.inputId, stored) + return Promise.resolve({ ...stored }) + } + listPending(threadId: string): Promise> { + // Map iteration keeps insertion order; sort by time in case clocks tie. + const pending = [...this.entries.values()] + .filter( + (entry) => entry.threadId === threadId && entry.status === 'pending', + ) + .sort((a, b) => a.createdAt - b.createdAt) + .map((entry) => ({ ...entry })) + return Promise.resolve(pending) + } + markApplied(inputId: string, operationId: string): Promise { + const existing = this.entries.get(inputId) + if (existing) { + this.entries.set(inputId, { ...existing, status: 'applied', operationId }) + } + return Promise.resolve() + } + markRejected(inputId: string, reason: string): Promise { + const existing = this.entries.get(inputId) + if (existing) { + this.entries.set(inputId, { ...existing, status: 'rejected', reason }) + } + return Promise.resolve() + } + get(inputId: string): Promise { + const existing = this.entries.get(inputId) + return Promise.resolve(existing ? { ...existing } : null) + } +} + interface MemoryPersistenceStores { messages: MessageStore runs: RunStore @@ -550,14 +591,16 @@ interface MemoryPersistenceStores { metadata: MetadataStore artifacts: ArtifactStore blobs: BlobStore + inbox: InboxStore } /** * In-process reference backend for the full state + generation store set. * * Returns messages + runs + generationRuns + interrupts + metadata + artifacts - * + blobs. Locks are not included — use `InMemoryLockStore` + `withLocks` from - * `@tanstack/ai` when a test or single-process app needs coordination. + * + blobs + inbox. Locks are not included — use `InMemoryLockStore` + + * `withLocks` from `@tanstack/ai` when a test or single-process app needs + * coordination. */ export function memoryPersistence() { const stores: MemoryPersistenceStores = { @@ -566,6 +609,7 @@ export function memoryPersistence() { generationRuns: new MemoryGenerationRunStore(), interrupts: new MemoryInterruptStore(), metadata: new MemoryMetadataStore(), + inbox: new MemoryInboxStore(), artifacts: new MemoryArtifactStore(), blobs: new MemoryBlobStore(), } diff --git a/packages/ai-persistence/src/testkit/conformance.ts b/packages/ai-persistence/src/testkit/conformance.ts index 253cc6e630..f621c7f6a2 100644 --- a/packages/ai-persistence/src/testkit/conformance.ts +++ b/packages/ai-persistence/src/testkit/conformance.ts @@ -1528,6 +1528,69 @@ export function runPersistenceConformance( }) }) + // The inbox is opt-in: only harness hosts read it. A backend without it + // skips these cases and needs no `skip` entry. + describe('inbox', () => { + it('appends idempotently, lists pending oldest first, and settles entries', async (ctx) => { + const store = persistence.stores.inbox + if (!store) return ctx.skip('inbox store not provided') + + const first = await store.append({ + inputId: 'in-1', + threadId: 'inbox-thread', + input: { op: 'prompt', message: 'hi' }, + createdAt: 1, + }) + expect(first.status).toBe('pending') + // Same id again: the stored entry wins, the new payload is ignored. + const again = await store.append({ + inputId: 'in-1', + threadId: 'inbox-thread', + input: { op: 'prompt', message: 'changed' }, + createdAt: 5, + }) + expect(again.input).toEqual({ op: 'prompt', message: 'hi' }) + + await store.append({ + inputId: 'in-2', + threadId: 'inbox-thread', + principal: { id: 'user-1' }, + input: { op: 'steer', message: 'shorter' }, + createdAt: 2, + }) + await store.append({ + inputId: 'in-other', + threadId: 'other-thread', + input: { op: 'prompt', message: 'x' }, + createdAt: 0, + }) + + expect( + (await store.listPending('inbox-thread')).map( + (entry) => entry.inputId, + ), + ).toEqual(['in-1', 'in-2']) + + await store.markApplied('in-1', 'op-1') + await store.markRejected('in-2', 'busy') + expect(await store.listPending('inbox-thread')).toEqual([]) + expect(await store.get('in-1')).toMatchObject({ + status: 'applied', + operationId: 'op-1', + }) + expect(await store.get('in-2')).toMatchObject({ + status: 'rejected', + reason: 'busy', + principal: { id: 'user-1' }, + }) + + // Unknown ids are no-ops. + await store.markApplied('missing', 'op-x') + await store.markRejected('missing', 'nope') + expect(await store.get('missing')).toBeNull() + }) + }) + describe('metadata', () => { it('sets, gets, namespaces, and deletes without composite-key collisions', async (ctx) => { const store = resolveStore('metadata') diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index 54c917d0b6..986c384bf0 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -368,6 +368,53 @@ export function defineInterruptStore(store: InterruptStore): InterruptStore { export function defineMetadataStore(store: MetadataStore): MetadataStore { return store } + +/** Lifecycle of one inbox entry. */ +export type InboxStatus = 'pending' | 'applied' | 'rejected' | 'expired' + +/** + * One input a client sent to a harness session (a prompt, a steer message, a + * follow-up, an interrupt answer, an agent run, a command). Written before the + * session answers with a receipt, so an accepted input survives a crash. + */ +export interface InboxEntry { + /** Idempotency key. A second append with the same id is a no-op. */ + inputId: string + threadId: string + /** Who sent it, from the host's `authorize`. */ + principal?: { id: string } + /** The input itself. Storage holds it as-is. The harness validates it. */ + input: unknown + status: InboxStatus + createdAt: number + expiresAt?: number + /** The operation that applied the input. */ + operationId?: string + /** Why the input was rejected. */ + reason?: string +} + +/** Durable store for harness session inputs. */ +export interface InboxStore { + /** + * Store a new entry as `'pending'`, or return the existing entry unchanged + * when `inputId` is already present. + */ + append: (entry: Omit) => Promise + /** Pending entries of a thread, oldest first. */ + listPending: (threadId: string) => Promise> + /** Mark an entry applied by `operationId`. A no-op for an unknown id. */ + markApplied: (inputId: string, operationId: string) => Promise + /** Mark an entry rejected with `reason`. A no-op for an unknown id. */ + markRejected: (inputId: string, reason: string) => Promise + /** The entry for `inputId`, or `null`. */ + get: (inputId: string) => Promise +} + +/** Type an {@link InboxStore} implementation inline. */ +export function defineInboxStore(store: InboxStore): InboxStore { + return store +} /** Type a {@link GenerationRunStore} implementation inline. */ export function defineGenerationRunStore( store: GenerationRunStore, @@ -620,6 +667,8 @@ export interface AIPersistenceStores { generationRuns?: GenerationRunStore artifacts?: ArtifactStore blobs?: BlobStore + /** Harness session inputs. Optional: only harness hosts read it. */ + inbox?: InboxStore } /** @@ -795,6 +844,7 @@ const storeKeys = [ 'metadata', 'artifacts', 'blobs', + 'inbox', ] satisfies Array const storeKeySet = new Set(storeKeys) diff --git a/packages/ai-persistence/tests/inbox.test.ts b/packages/ai-persistence/tests/inbox.test.ts new file mode 100644 index 0000000000..b88414483e --- /dev/null +++ b/packages/ai-persistence/tests/inbox.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from 'vitest' +import { + defineAIPersistence, + defineInboxStore, + memoryPersistence, +} from '../src' + +describe('inbox store', () => { + it('ships in memoryPersistence', () => { + expect(memoryPersistence().stores.inbox).toBeDefined() + }) + + it('returns a copy, so callers cannot change stored entries', async () => { + const inbox = memoryPersistence().stores.inbox + const entry = await inbox.append({ + inputId: 'in-1', + threadId: 't', + input: { op: 'prompt' }, + createdAt: 1, + }) + entry.status = 'applied' + + expect((await inbox.get('in-1'))?.status).toBe('pending') + }) + + it('accepts a user-implemented inbox in defineAIPersistence', () => { + const inbox = defineInboxStore({ + append: (entry) => Promise.resolve({ ...entry, status: 'pending' }), + listPending: () => Promise.resolve([]), + markApplied: () => Promise.resolve(), + markRejected: () => Promise.resolve(), + get: () => Promise.resolve(null), + }) + + expect(defineAIPersistence({ stores: { inbox } }).stores.inbox).toBe(inbox) + }) +}) diff --git a/packages/ai-persistence/tests/memory.test.ts b/packages/ai-persistence/tests/memory.test.ts index 14c14374ab..db479ec83b 100644 --- a/packages/ai-persistence/tests/memory.test.ts +++ b/packages/ai-persistence/tests/memory.test.ts @@ -16,6 +16,7 @@ describe('memoryPersistence', () => { 'artifacts', 'blobs', 'generationRuns', + 'inbox', 'interrupts', 'messages', 'metadata', diff --git a/packages/ai-persistence/tests/persistence-types.test-d.ts b/packages/ai-persistence/tests/persistence-types.test-d.ts index 7db2404d4b..467e17a3f0 100644 --- a/packages/ai-persistence/tests/persistence-types.test-d.ts +++ b/packages/ai-persistence/tests/persistence-types.test-d.ts @@ -21,6 +21,7 @@ import type { ChatTranscriptPersistence, ChatTranscriptStores, GenerationRunStore, + InboxStore, InterruptStore, MessagePage, MessageStore, @@ -139,6 +140,7 @@ expectTypeOf(memoryPersistence()).toEqualTypeOf< metadata: MetadataStore artifacts: ArtifactStore blobs: BlobStore + inbox: InboxStore }> >() const transcript: ChatTranscriptPersistence = messagesOnly diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index bb2055aa0c..8e2889ab62 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -109,7 +109,14 @@ export type { SubagentRouterPlan, SubagentStep, SubagentStepsPlan, + SubagentsBag, } from './activities/chat/agents/spawn' +// For hosts (a harness session) that run an agent outside a parent chat turn. +export { + spawnAgentStream as runAgentStream, + createSubagentId, +} from './activities/chat/agents/spawn' +export { compactForModel } from './activities/chat/tools/tool-calls' // Tool definition export { @@ -343,6 +350,7 @@ export type { // Capability primitives + middleware builder export { createCapability, + CapabilityRegistry, defineChatMiddleware, createChatMiddleware, MetadataCapability, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index f850d1c3eb..ca0c6edf73 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2520,6 +2520,21 @@ importers: specifier: ^8.2.1 version: 8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0) + packages/ai-harness: + devDependencies: + '@tanstack/ai': + specifier: workspace:* + version: link:../ai + '@tanstack/ai-persistence': + specifier: workspace:* + version: link:../ai-persistence + '@vitest/coverage-v8': + specifier: 4.1.10 + version: 4.1.10(vitest@4.1.11) + zod: + specifier: ^4.2.0 + version: 4.3.6 + packages/ai-isolate-cloudflare: dependencies: '@tanstack/ai-code-mode': @@ -3574,6 +3589,9 @@ importers: '@tanstack/ai-groq': specifier: workspace:* version: link:../../packages/ai-groq + '@tanstack/ai-harness': + specifier: workspace:* + version: link:../../packages/ai-harness '@tanstack/ai-llmgateway': specifier: workspace:* version: link:../../packages/ai-llmgateway diff --git a/testing/e2e/fixtures/harness/basic.json b/testing/e2e/fixtures/harness/basic.json new file mode 100644 index 0000000000..252c90de69 --- /dev/null +++ b/testing/e2e/fixtures/harness/basic.json @@ -0,0 +1,20 @@ +{ + "fixtures": [ + { + "match": { "userMessage": "[harness-turns] first" }, + "response": { "content": "First answer." } + }, + { + "match": { "userMessage": "[harness-turns] second" }, + "response": { "content": "Second answer." } + }, + { + "match": { "userMessage": "[harness-agent] price vendor a" }, + "response": { "content": "Vendor A costs 12 dollars." } + }, + { + "match": { "userMessage": "[harness-agent] what did it cost?" }, + "response": { "content": "It cost 12 dollars." } + } + ] +} diff --git a/testing/e2e/package.json b/testing/e2e/package.json index 148f58052e..41b655bed1 100644 --- a/testing/e2e/package.json +++ b/testing/e2e/package.json @@ -31,6 +31,7 @@ "@tanstack/ai-gemini": "workspace:*", "@tanstack/ai-grok": "workspace:*", "@tanstack/ai-groq": "workspace:*", + "@tanstack/ai-harness": "workspace:*", "@tanstack/ai-llmgateway": "workspace:*", "@tanstack/ai-lovable": "workspace:*", "@tanstack/ai-mcp": "workspace:*", diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 7308d98796..87b7a9edc4 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -97,6 +97,7 @@ import { Route as ApiLazyToolsWireRouteImport } from './routes/api.lazy-tools-wi import { Route as ApiJoinRunClientToolRouteImport } from './routes/api.join-run-client-tool' import { Route as ApiInterruptsTestRouteImport } from './routes/api.interrupts-test' import { Route as ApiImageRouteImport } from './routes/api.image' +import { Route as ApiHarnessTestRouteImport } from './routes/api.harness-test' import { Route as ApiGenerationPersistenceServerRouteImport } from './routes/api.generation-persistence-server' import { Route as ApiGenerationPersistenceResumeRouteImport } from './routes/api.generation-persistence-resume' import { Route as ApiGeminiNativeImageWireRouteImport } from './routes/api.gemini-native-image-wire' @@ -585,6 +586,11 @@ const ApiImageRoute = ApiImageRouteImport.update({ path: '/api/image', getParentRoute: () => rootRouteImport, } as any) +const ApiHarnessTestRoute = ApiHarnessTestRouteImport.update({ + id: '/api/harness-test', + path: '/api/harness-test', + getParentRoute: () => rootRouteImport, +} as any) const ApiGenerationPersistenceServerRoute = ApiGenerationPersistenceServerRouteImport.update({ id: '/api/generation-persistence-server', @@ -795,6 +801,7 @@ export interface FileRoutesByFullPath { '/api/gemini-native-image-wire': typeof ApiGeminiNativeImageWireRoute '/api/generation-persistence-resume': typeof ApiGenerationPersistenceResumeRoute '/api/generation-persistence-server': typeof ApiGenerationPersistenceServerRoute + '/api/harness-test': typeof ApiHarnessTestRoute '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/join-run-client-tool': typeof ApiJoinRunClientToolRoute @@ -914,6 +921,7 @@ export interface FileRoutesByTo { '/api/gemini-native-image-wire': typeof ApiGeminiNativeImageWireRoute '/api/generation-persistence-resume': typeof ApiGenerationPersistenceResumeRoute '/api/generation-persistence-server': typeof ApiGenerationPersistenceServerRoute + '/api/harness-test': typeof ApiHarnessTestRoute '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/join-run-client-tool': typeof ApiJoinRunClientToolRoute @@ -1034,6 +1042,7 @@ export interface FileRoutesById { '/api/gemini-native-image-wire': typeof ApiGeminiNativeImageWireRoute '/api/generation-persistence-resume': typeof ApiGenerationPersistenceResumeRoute '/api/generation-persistence-server': typeof ApiGenerationPersistenceServerRoute + '/api/harness-test': typeof ApiHarnessTestRoute '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/join-run-client-tool': typeof ApiJoinRunClientToolRoute @@ -1155,6 +1164,7 @@ export interface FileRouteTypes { | '/api/gemini-native-image-wire' | '/api/generation-persistence-resume' | '/api/generation-persistence-server' + | '/api/harness-test' | '/api/image' | '/api/interrupts-test' | '/api/join-run-client-tool' @@ -1274,6 +1284,7 @@ export interface FileRouteTypes { | '/api/gemini-native-image-wire' | '/api/generation-persistence-resume' | '/api/generation-persistence-server' + | '/api/harness-test' | '/api/image' | '/api/interrupts-test' | '/api/join-run-client-tool' @@ -1393,6 +1404,7 @@ export interface FileRouteTypes { | '/api/gemini-native-image-wire' | '/api/generation-persistence-resume' | '/api/generation-persistence-server' + | '/api/harness-test' | '/api/image' | '/api/interrupts-test' | '/api/join-run-client-tool' @@ -1513,6 +1525,7 @@ export interface RootRouteChildren { ApiGeminiNativeImageWireRoute: typeof ApiGeminiNativeImageWireRoute ApiGenerationPersistenceResumeRoute: typeof ApiGenerationPersistenceResumeRoute ApiGenerationPersistenceServerRoute: typeof ApiGenerationPersistenceServerRoute + ApiHarnessTestRoute: typeof ApiHarnessTestRoute ApiImageRoute: typeof ApiImageRouteWithChildren ApiInterruptsTestRoute: typeof ApiInterruptsTestRoute ApiJoinRunClientToolRoute: typeof ApiJoinRunClientToolRoute @@ -2190,6 +2203,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiImageRouteImport parentRoute: typeof rootRouteImport } + '/api/harness-test': { + id: '/api/harness-test' + path: '/api/harness-test' + fullPath: '/api/harness-test' + preLoaderRoute: typeof ApiHarnessTestRouteImport + parentRoute: typeof rootRouteImport + } '/api/generation-persistence-server': { id: '/api/generation-persistence-server' path: '/api/generation-persistence-server' @@ -2511,6 +2531,7 @@ const rootRouteChildren: RootRouteChildren = { ApiGeminiNativeImageWireRoute: ApiGeminiNativeImageWireRoute, ApiGenerationPersistenceResumeRoute: ApiGenerationPersistenceResumeRoute, ApiGenerationPersistenceServerRoute: ApiGenerationPersistenceServerRoute, + ApiHarnessTestRoute: ApiHarnessTestRoute, ApiImageRoute: ApiImageRouteWithChildren, ApiInterruptsTestRoute: ApiInterruptsTestRoute, ApiJoinRunClientToolRoute: ApiJoinRunClientToolRoute, diff --git a/testing/e2e/src/routes/api.harness-test.ts b/testing/e2e/src/routes/api.harness-test.ts new file mode 100644 index 0000000000..6b7be92b52 --- /dev/null +++ b/testing/e2e/src/routes/api.harness-test.ts @@ -0,0 +1,80 @@ +import { createFileRoute } from '@tanstack/react-router' +import { defineAgent } from '@tanstack/ai' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { z } from 'zod' +import { createTextAdapter } from '@/lib/providers' + +/** + * Harness session. The main model and the agent are real OpenAI adapters + * against aimock. + * + * - `turns`: two prompts sent back to back. The second waits for the first, + * then runs with the first turn in its history. + * - `agent`: `pricer` runs from code with typed input. Its result goes into + * the transcript, then a prompt runs. + */ +export const Route = createFileRoute('/api/harness-test')({ + server: { + handlers: { + POST: async ({ request }) => { + const body = (await request.json()) as { + scenario?: string + testId?: string + aimockPort?: number + } + const testId = body.testId + const aimockPort = body.aimockPort + const openai = () => + createTextAdapter('openai', undefined, aimockPort, testId).adapter + + const pricer = defineAgent({ + name: 'pricer', + description: 'Prices one vendor', + inputSchema: z.object({ task: z.string() }), + run: (ctx) => + ctx.chat({ + adapter: openai(), + messages: [{ role: 'user', content: ctx.input.task }], + stream: false, + }), + }) + const harness = defineHarness({ + name: 'e2e/harness', + adapter: openai(), + agents: [pricer], + }) + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + try { + const session = await host.open(harness, { threadId: 'e2e-thread' }) + if (body.scenario === 'agent') { + const result = await session.agents.pricer.run({ + task: '[harness-agent] price vendor a', + }) + const turn = await session.prompt( + '[harness-agent] what did it cost?', + ) + return Response.json({ result, text: turn.text }) + } + const first = session.prompt('[harness-turns] first') + const second = session.prompt('[harness-turns] second') + const texts = [(await first).text, (await second).text] + const saved = + await persistence.stores.messages.loadThread('e2e-thread') + return Response.json({ + texts, + roles: saved.map((message) => message.role), + }) + } catch (error) { + return Response.json( + { error: error instanceof Error ? error.message : String(error) }, + { status: 500 }, + ) + } finally { + await host.close() + } + }, + }, + }, +}) diff --git a/testing/e2e/tests/harness.spec.ts b/testing/e2e/tests/harness.spec.ts new file mode 100644 index 0000000000..03e3c458a5 --- /dev/null +++ b/testing/e2e/tests/harness.spec.ts @@ -0,0 +1,31 @@ +import { test, expect } from './fixtures' + +test.describe('harness session', () => { + test('a second prompt waits for the first and keeps its history', async ({ + request, + testId, + aimockPort, + }) => { + const response = await request.post('/api/harness-test', { + data: { scenario: 'turns', testId, aimockPort }, + }) + expect(response.ok()).toBe(true) + const body = await response.json() + expect(body.texts).toEqual(['First answer.', 'Second answer.']) + expect(body.roles).toEqual(['user', 'assistant', 'user', 'assistant']) + }) + + test('a typed agent runs from code and the next turn sees its result', async ({ + request, + testId, + aimockPort, + }) => { + const response = await request.post('/api/harness-test', { + data: { scenario: 'agent', testId, aimockPort }, + }) + expect(response.ok()).toBe(true) + const body = await response.json() + expect(body.result).toBe('Vendor A costs 12 dollars.') + expect(body.text).toBe('It cost 12 dollars.') + }) +}) From 06b138dfe9be1b6af0e761bf6bc9579b6e47ff9d Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 12:41:03 +0200 Subject: [PATCH 007/286] feat(ai-harness): resume crashed turns and serve sessions to clients Resume: a session holds a lease on each running turn, saves the transcript around tool phases, and records tool calls that have no result. When a host opens a thread whose turn lost its lease, it continues the turn in a new run and emits harness.operation.resumed. toolDefinition({ replay }) decides whether an unfinished tool runs again. Protocol: createHarnessHandler serves capabilities, standard AG-UI runs, the session event stream with cursors, control inputs with receipts, and snapshots, behind a required authorize hook. handleHarnessSocket serves the same session tier over a WebSocket. @tanstack/ai-harness/client adds a typed, reconnecting client. harnessText runs a harness as the model of another chat() call. ACP: @tanstack/ai-acp/agent serves a harness as an ACP v2 agent (SDK 1.5), with tool approvals as permission requests. CLI: new @tanstack/ai-harness-cli. runCli(harness) gives an Ink terminal UI, a print mode with exit codes, NDJSON output, --acp, and --serve. --- .changeset/harness-p2-protocols.md | 14 + docs/config.json | 15 + docs/harness/cli.md | 83 +++++ docs/harness/connect.md | 127 ++++++++ docs/harness/durable-sessions.md | 71 ++++ packages/ai-acp/package.json | 14 +- packages/ai-acp/src/agent/index.ts | 251 ++++++++++++++ packages/ai-acp/tests/agent.test.ts | 175 ++++++++++ packages/ai-acp/vite.config.ts | 2 +- packages/ai-harness-cli/README.md | 63 ++++ packages/ai-harness-cli/package.json | 92 ++++++ packages/ai-harness-cli/src/args.ts | 68 ++++ packages/ai-harness-cli/src/commands.ts | 100 ++++++ packages/ai-harness-cli/src/index.ts | 111 +++++++ packages/ai-harness-cli/src/interactive.tsx | 170 ++++++++++ packages/ai-harness-cli/src/lines.ts | 80 +++++ packages/ai-harness-cli/src/print.ts | 54 +++ packages/ai-harness-cli/src/serve.ts | 113 +++++++ packages/ai-harness-cli/src/session-view.ts | 96 ++++++ packages/ai-harness-cli/tests/cli.test.ts | 265 +++++++++++++++ .../ai-harness-cli/tests/interactive.test.ts | 106 ++++++ packages/ai-harness-cli/tsconfig.json | 9 + packages/ai-harness-cli/vite.config.ts | 35 ++ packages/ai-harness/package.json | 4 + packages/ai-harness/src/client.ts | 198 +++++++++++ packages/ai-harness/src/harness-text.ts | 133 ++++++++ packages/ai-harness/src/host.ts | 2 + packages/ai-harness/src/http.ts | 308 ++++++++++++++++++ packages/ai-harness/src/index.ts | 20 ++ packages/ai-harness/src/protocol.ts | 200 ++++++++++++ packages/ai-harness/src/resume.ts | 146 +++++++++ packages/ai-harness/src/session.ts | 67 ++++ packages/ai-harness/src/types.ts | 1 + packages/ai-harness/tests/client.test.ts | 87 +++++ .../ai-harness/tests/harness-text.test.ts | 49 +++ packages/ai-harness/tests/protocol.test.ts | 287 ++++++++++++++++ packages/ai-harness/tests/resume.test.ts | 185 +++++++++++ packages/ai-harness/vite.config.ts | 2 +- .../activities/chat/tools/tool-definition.ts | 7 + packages/ai/src/types.ts | 7 + pnpm-lock.yaml | 299 ++++++++++++++--- testing/e2e/fixtures/harness/basic.json | 40 ++- testing/e2e/src/routeTree.gen.ts | 21 ++ .../e2e/src/routes/api.harness-protocol.$.ts | 57 ++++ testing/e2e/tests/harness-protocol.spec.ts | 83 +++++ 45 files changed, 4262 insertions(+), 55 deletions(-) create mode 100644 .changeset/harness-p2-protocols.md create mode 100644 docs/harness/cli.md create mode 100644 docs/harness/connect.md create mode 100644 docs/harness/durable-sessions.md create mode 100644 packages/ai-acp/src/agent/index.ts create mode 100644 packages/ai-acp/tests/agent.test.ts create mode 100644 packages/ai-harness-cli/README.md create mode 100644 packages/ai-harness-cli/package.json create mode 100644 packages/ai-harness-cli/src/args.ts create mode 100644 packages/ai-harness-cli/src/commands.ts create mode 100644 packages/ai-harness-cli/src/index.ts create mode 100644 packages/ai-harness-cli/src/interactive.tsx create mode 100644 packages/ai-harness-cli/src/lines.ts create mode 100644 packages/ai-harness-cli/src/print.ts create mode 100644 packages/ai-harness-cli/src/serve.ts create mode 100644 packages/ai-harness-cli/src/session-view.ts create mode 100644 packages/ai-harness-cli/tests/cli.test.ts create mode 100644 packages/ai-harness-cli/tests/interactive.test.ts create mode 100644 packages/ai-harness-cli/tsconfig.json create mode 100644 packages/ai-harness-cli/vite.config.ts create mode 100644 packages/ai-harness/src/client.ts create mode 100644 packages/ai-harness/src/harness-text.ts create mode 100644 packages/ai-harness/src/http.ts create mode 100644 packages/ai-harness/src/protocol.ts create mode 100644 packages/ai-harness/src/resume.ts create mode 100644 packages/ai-harness/tests/client.test.ts create mode 100644 packages/ai-harness/tests/harness-text.test.ts create mode 100644 packages/ai-harness/tests/protocol.test.ts create mode 100644 packages/ai-harness/tests/resume.test.ts create mode 100644 testing/e2e/src/routes/api.harness-protocol.$.ts create mode 100644 testing/e2e/tests/harness-protocol.spec.ts diff --git a/.changeset/harness-p2-protocols.md b/.changeset/harness-p2-protocols.md new file mode 100644 index 0000000000..cbdc0bbe0b --- /dev/null +++ b/.changeset/harness-p2-protocols.md @@ -0,0 +1,14 @@ +--- +'@tanstack/ai-harness': minor +'@tanstack/ai-harness-cli': minor +'@tanstack/ai-acp': minor +'@tanstack/ai': minor +--- + +Harness sessions now survive a crash and talk to clients. + +- **Resume.** A session holds a lease on each running turn, saves the transcript around each tool phase, and records tool calls that have no result yet. When a host opens a session whose turn lost its lease, it continues the turn in a new run and emits `harness.operation.resumed`. `toolDefinition({ replay: 'safe' | 'never' })` in `@tanstack/ai` decides whether an unfinished tool runs again. +- **Protocol.** `createHarnessHandler` serves capabilities, standard AG-UI runs, the session event stream with cursors, control inputs with receipts, and snapshots. `authorize` is required. `handleHarnessSocket` serves the same session tier over a WebSocket. `createHarnessClient` in `@tanstack/ai-harness/client` is a typed, reconnecting client. +- **`harnessText`** runs a harness as the text adapter of another `chat()` call. +- **ACP v2.** `@tanstack/ai-acp/agent` adds `createAcpAgent` and `serveAcp`, which serve a harness as an ACP v2 agent (experimental, like the draft protocol). `@tanstack/ai-acp` now uses `@agentclientprotocol/sdk` 1.5. The existing ACP client adapters keep working. +- **New package `@tanstack/ai-harness-cli`.** `runCli(harness)` gives a harness an Ink terminal UI, a print mode with exit codes, NDJSON output, `--acp`, and `--serve`. diff --git a/docs/config.json b/docs/config.json index e08f6892b5..4dac967f22 100644 --- a/docs/config.json +++ b/docs/config.json @@ -808,6 +808,21 @@ "label": "Build your first harness", "to": "harness/overview", "addedAt": "2026-09-26" + }, + { + "label": "Connect clients", + "to": "harness/connect", + "addedAt": "2026-09-26" + }, + { + "label": "Run in the terminal", + "to": "harness/cli", + "addedAt": "2026-09-26" + }, + { + "label": "Durable sessions", + "to": "harness/durable-sessions", + "addedAt": "2026-09-26" } ] }, diff --git a/docs/harness/cli.md b/docs/harness/cli.md new file mode 100644 index 0000000000..a572a2bc02 --- /dev/null +++ b/docs/harness/cli.md @@ -0,0 +1,83 @@ +--- +title: Run a harness in the terminal +id: harness-cli +order: 3 +description: "Give your harness a terminal UI, a print mode for scripts and CI, NDJSON output, an ACP mode for editors, and an HTTP server." +keywords: + - tanstack ai + - harness + - cli + - terminal + - ink +--- + +You built a harness and want to use it like Claude Code: type in a terminal, watch it work, approve tools. You also want to run it in CI. `runCli` gives one harness all of these modes. + +## Install + + + +react: @tanstack/ai-harness-cli +vue: @tanstack/ai-harness-cli +solid: @tanstack/ai-harness-cli +svelte: @tanstack/ai-harness-cli +preact: @tanstack/ai-harness-cli +angular: @tanstack/ai-harness-cli +vanilla: @tanstack/ai-harness-cli +octane: @tanstack/ai-harness-cli + + + +The interactive UI uses Ink, which needs Node 22 or later. + +## 1. Write the entry file + +```ts group=harness-cli +import { defineHarness } from '@tanstack/ai-harness' +import { runCli } from '@tanstack/ai-harness-cli' +import { openaiText } from '@tanstack/ai-openai' + +const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-5.6'), +}) + +process.exitCode = await runCli(assistant) +``` + +## 2. Pick a mode + +- No flags: the interactive UI. Type a message and press Enter. While the agent works, Enter steers it and Esc cancels. +- `-p "prompt"`: run one prompt, print the answer, and exit. +- `-p "prompt" --output ndjson`: print every AG-UI event as one JSON line. +- `--acp`: serve the harness as an ACP v2 agent over stdio, for editors. Needs `@tanstack/ai-acp`. +- `--serve`: serve the session protocol over HTTP on `127.0.0.1:8787`. Every request needs the bearer token. Pass `--token`, set `HARNESS_TOKEN`, or copy the token the CLI prints. + +When stdin is a pipe, the CLI reads one message or command per line and waits for each turn. + +## 3. Use it in CI + +`-p` exits with a code your script can check: + +| Code | Meaning | +|---|---| +| 0 | The turn finished. | +| 1 | The turn failed. | +| 2 | The turn waits for an approval. | +| 130 | The turn was cancelled. | + +## Commands in the interactive UI + +- `/agents`: list the agents. +- `/agent {"json":"input"}`: run an agent in the background. When it is done, a new turn starts with its result. +- `/cancel`: cancel the running turn. +- `/status`: show what runs and what waits. +- `/exit`: quit. + +When a turn stops for an approval, type `y` to approve or `n` to reject. + +## What you have now + +- One entry file that runs your harness as a terminal app, a script step, an editor agent, or a server. + +Next: keep long turns alive through crashes with [durable sessions](./durable-sessions). diff --git a/docs/harness/connect.md b/docs/harness/connect.md new file mode 100644 index 0000000000..f38633fb99 --- /dev/null +++ b/docs/harness/connect.md @@ -0,0 +1,127 @@ +--- +title: Connect clients to a harness +id: harness-connect +order: 2 +description: "Serve a harness session over HTTP, SSE, or WebSocket. Talk to it from a web app, an editor over ACP, or another chat() call." +keywords: + - tanstack ai + - harness + - AG-UI + - ACP + - websocket + - harnessText +--- + +Your harness runs on a server, but the people who use it are in a browser, an editor, or another agent. `createHarnessHandler` serves a session over HTTP. A web app talks to it with `createHarnessClient`, an editor with ACP, and another `chat()` call with `harnessText`. + +## Serve the session over HTTP + +Mount one fetch handler on a route. `authorize` is required, so no endpoint is open by accident. + +```ts group=harness-connect +import { defineHarness, createHarnessHandler, createHarnessHost } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { openaiText } from '@tanstack/ai-openai' + +export const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-5.6'), +}) + +const host = createHarnessHost({ persistence: memoryPersistence() }) + +export const handler = createHarnessHandler({ + host, + harness: assistant, + authorize: (request) => + request.headers.get('authorization') === `Bearer ${process.env.HARNESS_TOKEN}` + ? { id: 'user-1' } + : null, + canAccess: (principal, threadId) => threadId.startsWith(principal.id), +}) +``` + +The handler answers these paths under your route: + +- `GET capabilities`: the AG-UI capabilities, with the agents in `expose.agents`. +- `POST run`: standard AG-UI. One request runs one prompt and streams it as SSE. Any AG-UI client works. +- `GET events?threadId=`: every event of the session as SSE. Each event id is a cursor, so a reconnect with `Last-Event-ID` continues where it stopped. +- `POST control`: send `{ threadId, input }`, for example `{ op: 'prompt', message }`. You get a receipt back. +- `GET snapshot?threadId=`: the status, running operations, and waiting approvals. + +## Talk to it from a web app + +`createHarnessClient` wraps those endpoints. Import the harness with `import type`, so no server code reaches the browser. + +```ts group=harness-connect-client +import { createHarnessClient } from '@tanstack/ai-harness/client' +import type { assistant } from './harness' + +const client = createHarnessClient({ + url: '/api/harness', + threadId: 'user-1-thread', + headers: { authorization: 'Bearer my-token' }, +}) + +await client.prompt('Summarize my open tickets.') + +for await (const entry of client.events()) { + if (entry.event.type === 'TEXT_MESSAGE_CONTENT') { + console.log(entry.event.delta) + } +} +``` + +`events()` reconnects after a network error and continues from the last cursor. A second tab, a phone, or a reload all see the same session. + +## Use a WebSocket + +For one connection that carries events and inputs, authorize the upgrade, then hand the socket to `handleHarnessSocket`: + +```ts group=harness-connect +import { handleHarnessSocket } from '@tanstack/ai-harness' +import type { WebSocketLike } from '@tanstack/ai' + +export function onUpgrade(socket: WebSocketLike) { + handleHarnessSocket({ host, harness: assistant, socket, principal: { id: 'user-1' } }) +} +``` + +The client sends `{ type: 'harness.subscribe', threadId }` first. After that, each `harness.input` frame gets a `harness.receipt`, and every event arrives as a `harness.event` frame with its cursor. + +## Use it from an editor (ACP) + +Editors such as Zed start agents as a process and talk ACP over stdio. `serveAcp` makes the harness an ACP v2 agent: + +```ts group=harness-connect +import { serveAcp } from '@tanstack/ai-acp/agent' + +serveAcp({ host, harness: assistant }) +``` + +Tool approvals become permission requests in the editor. ACP v2 is still a draft, so this API is experimental. + +## Use it as the model of another chat + +`harnessText` turns a harness into a text adapter. The outer chat sends a message, and the harness runs a full turn with its own tools, plugins, and agents: + +```ts group=harness-connect +import { chat } from '@tanstack/ai' +import { harnessText } from '@tanstack/ai-harness' + +const stream = chat({ + adapter: harnessText(assistant, { host }), + messages: [{ role: 'user', content: 'Fix the failing test.' }], + threadId: 'outer-thread', +}) +``` + +Each outer thread gets its own inner session, so the harness keeps its own history. + +## What you have now + +- One handler that serves standard AG-UI and the session stream. +- A typed web client that reconnects and resumes from a cursor. +- A WebSocket, an ACP agent for editors, and a harness you can call from `chat()`. + +Next: run the same harness in a terminal with the [CLI](./cli). diff --git a/docs/harness/durable-sessions.md b/docs/harness/durable-sessions.md new file mode 100644 index 0000000000..a97a22270b --- /dev/null +++ b/docs/harness/durable-sessions.md @@ -0,0 +1,71 @@ +--- +title: Durable sessions +id: harness-durable-sessions +order: 4 +description: "Keep a harness turn alive through a crash or a deploy. The next host continues it from the last checkpoint." +keywords: + - tanstack ai + - harness + - durability + - resume + - checkpoint +--- + +A long agent turn can outlive the process that started it: a deploy restarts the server, or the machine crashes in the middle of a tool call. With durable stores, the next host that opens the session continues the turn from its last checkpoint. + +## Give the host durable stores + +Use your own stores built with `defineAIPersistence`. Resume needs these: + +- `messages`: the transcript. +- `runs`: the run records, with leases and checkpoints. Implement `listByThread`. +- `inbox`: messages that were accepted but did not run yet. + +```ts group=harness-durable +import { createHarnessHost } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' + +// Swap in your own stores. Memory stores do not survive a restart. +const host = createHarnessHost({ persistence: memoryPersistence() }) +``` + +## What the session saves while a turn runs + +- A lease on the run record. The host renews it every 10 seconds. It expires after 30 seconds. +- The transcript before and after each tool phase. +- Each tool call that started and has no result yet. + +## What happens after a crash + +When a host opens the session and finds a turn whose lease expired, it continues that turn in a new run: + +1. The old run is marked failed. +2. Tool calls without a result get handled by their `replay` setting. +3. A new chat run starts from the saved transcript. +4. Clients see a `harness.operation.resumed` event, then the new events. + +## Mark tools that are safe to run again + +A crash can stop a run after a tool started but before its result was saved. Tell the harness what to do with `replay`: + +```ts group=harness-durable +import { toolDefinition } from '@tanstack/ai' +import { z } from 'zod' + +const lookup = toolDefinition({ + name: 'lookup_order', + description: 'Read an order', + inputSchema: z.object({ id: z.string() }), + replay: 'safe', +}).server(async ({ id }) => ({ id, status: 'shipped' })) +``` + +- `replay: 'safe'`: the tool runs again. Use it for reads and other idempotent calls. +- `replay: 'never'` (default): the model gets a note that the tool may or may not have run, and checks before it tries again. + +A model response that was cut in the middle runs again after resume. The provider can bill that call twice. + +## What you have now + +- Turns that continue after a crash or a deploy. +- Tools that either run again safely or tell the model to check first. diff --git a/packages/ai-acp/package.json b/packages/ai-acp/package.json index 0e63ad861f..37012401e0 100644 --- a/packages/ai-acp/package.json +++ b/packages/ai-acp/package.json @@ -26,6 +26,10 @@ ".": { "types": "./dist/esm/index.d.ts", "import": "./dist/esm/index.js" + }, + "./agent": { + "types": "./dist/esm/agent/index.d.ts", + "import": "./dist/esm/agent/index.js" } }, "files": [ @@ -44,14 +48,22 @@ "test:types": "tsc" }, "dependencies": { - "@agentclientprotocol/sdk": "^0.25.0" + "@agentclientprotocol/sdk": "^1.5.0" }, "peerDependencies": { "@tanstack/ai": "workspace:^", + "@tanstack/ai-harness": "workspace:^", "@tanstack/ai-sandbox": "workspace:^" }, + "peerDependenciesMeta": { + "@tanstack/ai-harness": { + "optional": true + } + }, "devDependencies": { "@tanstack/ai": "workspace:*", + "@tanstack/ai-harness": "workspace:*", + "@tanstack/ai-persistence": "workspace:*", "@tanstack/ai-sandbox": "workspace:*", "@tanstack/ai-sandbox-local-process": "workspace:*", "@vitest/coverage-v8": "4.1.10" diff --git a/packages/ai-acp/src/agent/index.ts b/packages/ai-acp/src/agent/index.ts new file mode 100644 index 0000000000..a9d3fc1af0 --- /dev/null +++ b/packages/ai-acp/src/agent/index.ts @@ -0,0 +1,251 @@ +import { EventType } from '@tanstack/ai' +import { + PROTOCOL_VERSION, + agent, + ndJsonStream, +} from '@agentclientprotocol/sdk/experimental/v2' +import type { + AgentApp, + AgentContext, + AgentConnection, + SessionUpdate, + Stream, +} from '@agentclientprotocol/sdk/experimental/v2' +import type { Interrupt, RunAgentResumeItem, StreamChunk } from '@tanstack/ai' +import type { + AnyHarness, + HarnessHost, + HarnessSession, + Operation, +} from '@tanstack/ai-harness' + +export interface AcpAgentOptions { + host: HarnessHost + harness: AnyHarness + /** Reported in `initialize`. Default `'0.0.0'`. */ + version?: string +} + +/** The text of an ACP prompt. Non-text blocks are skipped. */ +function promptText( + prompt: ReadonlyArray<{ type: string; text?: unknown }>, +): string { + return prompt + .map((block) => + block.type === 'text' && typeof block.text === 'string' ? block.text : '', + ) + .filter((text) => text !== '') + .join('\n') +} + +/** One AG-UI chunk as an ACP session update, or `undefined` to skip it. */ +export function toSessionUpdate(chunk: StreamChunk): SessionUpdate | undefined { + // Child agent work stays inside the harness. ACP shows the main turn. + if ('subagentRunId' in chunk && chunk.subagentRunId) return undefined + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) { + return { + sessionUpdate: 'agent_message_chunk', + messageId: chunk.messageId, + content: { type: 'text', text: chunk.delta }, + } + } + if (chunk.type === EventType.REASONING_MESSAGE_CONTENT) { + return { + sessionUpdate: 'agent_thought_chunk', + messageId: chunk.messageId, + content: { type: 'text', text: chunk.delta }, + } + } + if (chunk.type === EventType.TOOL_CALL_START) { + return { + sessionUpdate: 'tool_call_update', + toolCallId: chunk.toolCallId, + name: chunk.toolCallName, + title: chunk.toolCallName, + status: 'in_progress', + } + } + if (chunk.type === EventType.TOOL_CALL_RESULT) { + return { + sessionUpdate: 'tool_call_update', + toolCallId: chunk.toolCallId, + status: 'completed', + } + } + return undefined +} + +const APPROVAL_OPTIONS = [ + { optionId: 'allow', name: 'Allow', kind: 'allow_once' }, + { optionId: 'reject', name: 'Reject', kind: 'reject_once' }, +] as const + +/** + * An ACP v2 agent backed by a harness. Each ACP session is one harness + * session. Prompts follow the harness `busy` setting, tool approvals become + * `session/request_permission` requests, and cancel cancels the running turn. + * + * ACP v2 is a draft, so this API is experimental too. + */ +export function createAcpAgent(options: AcpAgentOptions): AgentApp { + const { host, harness } = options + const sessions = new Map() + + const sessionFor = async (sessionId: string) => { + let session = sessions.get(sessionId) + if (!session) { + session = await host.open(harness, { threadId: sessionId }) + sessions.set(sessionId, session) + } + return session + } + + const notify = ( + client: AgentContext, + sessionId: string, + update: SessionUpdate, + ) => client.notify('session/update', { sessionId, update }).catch(() => {}) + + /** + * Ask the client about each approval, then continue the turn. The next + * operation is wrapped, because an operation is thenable and an async + * function would unwrap it. + */ + const approve = async ( + client: AgentContext, + sessionId: string, + session: HarnessSession, + interrupts: ReadonlyArray, + ): Promise<{ next: Operation | undefined }> => { + await notify(client, sessionId, { + sessionUpdate: 'state_update', + state: 'requires_action', + }) + const resume: Array = [] + for (const interrupt of interrupts) { + const toolCallId = interrupt.toolCallId ?? interrupt.id + const response = await client.request('session/request_permission', { + sessionId, + title: interrupt.message ?? 'Allow this tool call?', + subject: { type: 'tool_call', toolCall: { toolCallId } }, + options: [...APPROVAL_OPTIONS], + }) + const allowed = + response.outcome.outcome === 'selected' && + 'optionId' in response.outcome && + response.outcome.optionId === 'allow' + resume.push({ + interruptId: interrupt.id, + status: 'resolved', + payload: allowed, + }) + } + const receipt = await session.resolve(resume) + return { + next: receipt.operationId + ? session.operation(receipt.operationId) + : undefined, + } + } + + /** Stream one turn to the client, through any approvals, then report idle. */ + const pump = async ( + client: AgentContext, + sessionId: string, + session: HarnessSession, + first: Operation, + ) => { + await notify(client, sessionId, { + sessionUpdate: 'state_update', + state: 'running', + }) + let operation: Operation | undefined = first + let stopReason = 'end_turn' + while (operation) { + for await (const chunk of operation.stream()) { + const update = toSessionUpdate(chunk) + if (update) await notify(client, sessionId, update) + } + const result: unknown = await operation.then( + (value) => value, + () => undefined, + ) + if (operation.status() === 'cancelled') stopReason = 'cancelled' + const interrupts = + typeof result === 'object' && + result !== null && + 'interrupts' in result && + Array.isArray(result.interrupts) + ? (result.interrupts as Array) + : [] + operation = + interrupts.length > 0 + ? (await approve(client, sessionId, session, interrupts)).next + : undefined + } + await notify(client, sessionId, { + sessionUpdate: 'state_update', + state: 'idle', + stopReason, + }) + } + + return agent() + .onRequest('initialize', () => ({ + protocolVersion: PROTOCOL_VERSION, + info: { name: harness.name, version: options.version ?? '0.0.0' }, + capabilities: {}, + })) + .onRequest('session/new', async () => { + const sessionId = `acp-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 9)}` + await sessionFor(sessionId) + return { sessionId } + }) + .onRequest('session/resume', async ({ params }) => { + await sessionFor(params.sessionId) + return {} + }) + .onRequest('session/prompt', async ({ params, client }) => { + const session = await sessionFor(params.sessionId) + const operation = session.prompt(promptText(params.prompt)) + void pump(client, params.sessionId, session, operation) + return { messageId: operation.id } + }) + .onRequest('session/close', async ({ params }) => { + const session = sessions.get(params.sessionId) + sessions.delete(params.sessionId) + await session?.close() + }) + .onNotification('session/cancel', async ({ params }) => { + await sessions.get(params.sessionId)?.cancel() + }) +} + +/** + * Serve a harness as an ACP v2 agent over a stream. With no stream, it uses + * stdin and stdout as newline-delimited JSON (for editors that start the + * agent as a process). + */ +export function serveAcp( + options: AcpAgentOptions & { stream?: Stream }, +): AgentConnection { + return createAcpAgent(options).connect(options.stream ?? stdioStream()) +} + +function stdioStream(): Stream { + const output = new WritableStream({ + write: (chunk) => + new Promise((resolve, reject) => + process.stdout.write(chunk, (error) => + error ? reject(error) : resolve(), + ), + ), + }) + const input = new ReadableStream({ + start(controller) { + process.stdin.on('data', (data: Uint8Array) => controller.enqueue(data)) + process.stdin.on('end', () => controller.close()) + }, + }) + return ndJsonStream(output, input) +} diff --git a/packages/ai-acp/tests/agent.test.ts b/packages/ai-acp/tests/agent.test.ts new file mode 100644 index 0000000000..0fcfe8bf57 --- /dev/null +++ b/packages/ai-acp/tests/agent.test.ts @@ -0,0 +1,175 @@ +import { describe, expect, it, vi } from 'vitest' +import { EventType, toolDefinition } from '@tanstack/ai' +import { client } from '@agentclientprotocol/sdk/experimental/v2' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { createAcpAgent } from '../src/agent' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import type { SessionUpdate } from '@agentclientprotocol/sdk/experimental/v2' + +const now = () => Date.now() + +function scripted(turns: Array>) { + let call = 0 + const adapter: AnyTextAdapter = { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + chatStream: () => { + const chunks = turns[call] ?? [] + call += 1 + return (async function* () { + yield* chunks + })() + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } + return adapter +} + +const textTurn = (text: string): Array => [ + { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, + { + type: EventType.TEXT_MESSAGE_START, + messageId: 'm', + role: 'assistant', + timestamp: now(), + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm', + delta: text, + timestamp: now(), + }, + { type: EventType.TEXT_MESSAGE_END, messageId: 'm', timestamp: now() }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'stop' } }, + }, +] + +const toolTurn = (name: string, args: string): Array => [ + { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, + { + type: EventType.TOOL_CALL_START, + toolCallId: 'call_1', + toolCallName: name, + timestamp: now(), + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call_1', + delta: args, + timestamp: now(), + }, + { type: EventType.TOOL_CALL_END, toolCallId: 'call_1', timestamp: now() }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'tool_calls' } }, + }, +] + +describe('ACP v2 agent', () => { + it('streams a prompt, asks for approval, and reports idle', async () => { + const remove = vi.fn(async () => ({ removed: true })) + const harness = defineHarness({ + name: 'test/acp', + adapter: scripted([ + toolTurn('remove', '{"path":"a.txt"}'), + textTurn('Removed a.txt.'), + ]), + tools: [ + toolDefinition({ + name: 'remove', + description: 'Remove a file', + needsApproval: true, + inputSchema: { + type: 'object', + properties: { path: { type: 'string' } }, + required: ['path'], + }, + }).server(remove), + ], + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const agentApp = createAcpAgent({ host, harness, version: '1.0.0' }) + + const updates: Array = [] + const permissions: Array = [] + const editor = client() + .onNotification('session/update', ({ params }) => { + updates.push(params.update) + }) + .onRequest('session/request_permission', ({ params }) => { + permissions.push(params.title) + return { outcome: { outcome: 'selected', optionId: 'allow' } } + }) + + await editor.connectWith(agentApp, async (ctx) => { + const init = await ctx.request('initialize', { + protocolVersion: 2, + info: { name: 'test-editor', version: '1.0.0' }, + }) + expect(init.info.name).toBe('test/acp') + + const { sessionId } = await ctx.request('session/new', { + cwd: '/tmp', + mcpServers: [], + }) + const prompted = await ctx.request('session/prompt', { + sessionId, + prompt: [{ type: 'text', text: 'remove a.txt' }], + }) + expect(prompted.messageId).toMatch(/^op-chat-/) + + await vi.waitFor(() => + expect( + updates.some( + (update) => + update.sessionUpdate === 'state_update' && + update.state === 'idle', + ), + ).toBe(true), + ) + }) + + expect(permissions).toHaveLength(1) + expect(remove).toHaveBeenCalledTimes(1) + const states = updates + .filter((update) => update.sessionUpdate === 'state_update') + .map((update) => ('state' in update ? update.state : '')) + expect(states).toEqual(['running', 'requires_action', 'idle']) + const text = updates + .map((update) => + update.sessionUpdate === 'agent_message_chunk' + ? JSON.stringify(update) + : '', + ) + .join('') + expect(text).toContain('Removed a.txt.') + expect( + updates.some((update) => update.sessionUpdate === 'tool_call_update'), + ).toBe(true) + await host.close() + }) +}) diff --git a/packages/ai-acp/vite.config.ts b/packages/ai-acp/vite.config.ts index 588ff924d2..ab0a9e68fc 100644 --- a/packages/ai-acp/vite.config.ts +++ b/packages/ai-acp/vite.config.ts @@ -33,7 +33,7 @@ const config = defineConfig({ // Put package test settings last so they win over the shared build config. export default mergeConfig( tanstackViteConfig({ - entry: ['./src/index.ts'], + entry: ['./src/index.ts', './src/agent/index.ts'], srcDir: './src', cjs: false, }), diff --git a/packages/ai-harness-cli/README.md b/packages/ai-harness-cli/README.md new file mode 100644 index 0000000000..1ff4d8ea4d --- /dev/null +++ b/packages/ai-harness-cli/README.md @@ -0,0 +1,63 @@ +
+ + + + TanStack AI + +
+ +
+ +# @tanstack/ai-harness-cli + +Run a TanStack AI harness from the terminal: an interactive UI, a print mode for scripts and CI, NDJSON output, an ACP mode for editors, and an HTTP server. + +## Installation + +```bash +npm install @tanstack/ai-harness @tanstack/ai-harness-cli +``` + +The interactive UI uses Ink, which needs Node 22 or later. + +## Usage + +```ts +#!/usr/bin/env node +import { defineHarness } from '@tanstack/ai-harness' +import { runCli } from '@tanstack/ai-harness-cli' +import { openaiText } from '@tanstack/ai-openai' + +const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-5.6'), +}) + +process.exitCode = await runCli(assistant) +``` + +| Flag | Mode | +| ----------------------------- | -------------------------------------------------- | +| (none) | Interactive UI | +| `-p "prompt"` | Run one prompt and print the answer | +| `-p "prompt" --output ndjson` | Print every AG-UI event as JSON lines | +| `--acp` | ACP v2 agent over stdio (needs `@tanstack/ai-acp`) | +| `--serve` | Session protocol over HTTP, with a bearer token | + +## Documentation + +Read [Run a harness in the terminal](https://tanstack.com/ai/latest/docs/harness/cli). + +## License + +MIT diff --git a/packages/ai-harness-cli/package.json b/packages/ai-harness-cli/package.json new file mode 100644 index 0000000000..e9224933f2 --- /dev/null +++ b/packages/ai-harness-cli/package.json @@ -0,0 +1,92 @@ +{ + "name": "@tanstack/ai-harness-cli", + "version": "0.0.0", + "description": "Run a TanStack AI harness from the terminal: interactive UI, print mode, NDJSON, ACP, and an HTTP server", + "author": "", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/TanStack/ai.git", + "directory": "packages/ai-harness-cli" + }, + "type": "module", + "module": "./dist/esm/index.js", + "types": "./dist/esm/index.d.ts", + "exports": { + ".": { + "types": "./dist/esm/index.d.ts", + "import": "./dist/esm/index.js" + } + }, + "sideEffects": false, + "files": [ + "dist", + "src" + ], + "scripts": { + "build": "vite build", + "clean": "premove ./build ./dist", + "lint:fix": "oxlint src --type-aware --fix", + "test:build": "publint --strict", + "test:oxlint": "oxlint src --type-aware", + "test:coverage": "vitest run --passWithNoTests --coverage --coverage.include='src/**' --coverage.reporter=text-summary --coverage.reporter=json-summary", + "test:lib": "vitest --passWithNoTests", + "test:lib:dev": "pnpm test:lib --watch", + "test:types": "tsc" + }, + "nx": { + "targets": { + "test:lib": { + "dependsOn": [ + "build", + "^build" + ] + }, + "test:coverage": { + "dependsOn": [ + "build", + "^build" + ] + }, + "test:types": { + "dependsOn": [ + "build", + "^build" + ] + } + } + }, + "keywords": [ + "ai", + "tanstack", + "harness", + "agent", + "cli", + "terminal", + "acp" + ], + "dependencies": { + "ink": "^7.1.1", + "react": "^19.2.3" + }, + "peerDependencies": { + "@tanstack/ai": "workspace:^", + "@tanstack/ai-acp": "workspace:^", + "@tanstack/ai-harness": "workspace:^", + "@tanstack/ai-persistence": "workspace:^" + }, + "peerDependenciesMeta": { + "@tanstack/ai-acp": { + "optional": true + } + }, + "devDependencies": { + "@tanstack/ai": "workspace:*", + "@tanstack/ai-acp": "workspace:*", + "@tanstack/ai-harness": "workspace:*", + "@tanstack/ai-persistence": "workspace:*", + "@tanstack/ai-sandbox": "workspace:*", + "@types/react": "^19.2.7", + "@vitest/coverage-v8": "4.1.10" + } +} diff --git a/packages/ai-harness-cli/src/args.ts b/packages/ai-harness-cli/src/args.ts new file mode 100644 index 0000000000..c83474e715 --- /dev/null +++ b/packages/ai-harness-cli/src/args.ts @@ -0,0 +1,68 @@ +import { parseArgs } from 'node:util' + +export interface CliArgs { + help: boolean + print?: string + output: 'text' | 'ndjson' + acp: boolean + serve: boolean + port: number + hostname: string + token?: string + thread: string +} + +export const USAGE = `Usage: [options] + +With no options, starts the interactive UI. + +Options: + -p, --print Run one prompt, print the answer, and exit + --output Print format for --print: text (default) or ndjson + --acp Serve the harness as an ACP v2 agent over stdio + --serve Serve the session protocol over HTTP + --port Port for --serve (default 8787) + --host Host name for --serve (default 127.0.0.1) + --token Bearer token for --serve (default: HARNESS_TOKEN, or a new random token) + --thread Conversation id (default: main) + -h, --help Show this help + +Exit codes for --print: 0 done, 1 failed, 2 waiting for approval, 130 cancelled.` + +/** Parse CLI flags. Throws with a short message on an unknown flag. */ +export function parseCliArgs(argv: ReadonlyArray): CliArgs { + const { values } = parseArgs({ + args: [...argv], + options: { + help: { type: 'boolean', short: 'h' }, + print: { type: 'string', short: 'p' }, + output: { type: 'string' }, + acp: { type: 'boolean' }, + serve: { type: 'boolean' }, + port: { type: 'string' }, + host: { type: 'string' }, + token: { type: 'string' }, + thread: { type: 'string' }, + }, + strict: true, + }) + const output = values.output ?? 'text' + if (output !== 'text' && output !== 'ndjson') { + throw new Error(`--output must be text or ndjson, got "${output}".`) + } + const port = Number(values.port ?? '8787') + if (!Number.isInteger(port) || port < 0 || port > 65535) { + throw new Error(`--port must be a port number, got "${values.port}".`) + } + return { + help: values.help ?? false, + ...(values.print !== undefined ? { print: values.print } : {}), + output, + acp: values.acp ?? false, + serve: values.serve ?? false, + port, + hostname: values.host ?? '127.0.0.1', + ...(values.token !== undefined ? { token: values.token } : {}), + thread: values.thread ?? 'main', + } +} diff --git a/packages/ai-harness-cli/src/commands.ts b/packages/ai-harness-cli/src/commands.ts new file mode 100644 index 0000000000..4c0ce10d28 --- /dev/null +++ b/packages/ai-harness-cli/src/commands.ts @@ -0,0 +1,100 @@ +import type { HarnessSession } from '@tanstack/ai-harness' + +export const HELP_TEXT = [ + 'Type a message and press Enter. While the agent works, a new message steers it.', + 'Commands:', + ' /agents list the agents', + ' /agent [json] run an agent in the background', + ' /cancel cancel the running turn (or press Esc)', + ' /status show the session status', + ' /exit quit', +].join('\n') + +/** What the UI does after a line of input. */ +export type LineResult = + | { type: 'notice'; text: string } + | { type: 'exit' } + | { type: 'sent' } + +/** + * Handle one line of user input: a slash command, a steer while a turn runs, + * or a new prompt. Shared by the interactive UI and the line mode. + */ +export async function handleLine( + session: HarnessSession, + line: string, +): Promise { + const text = line.trim() + if (text === '') return { type: 'notice', text: '' } + if (!text.startsWith('/')) { + if (session.snapshot().status === 'running') await session.steer(text) + else void session.prompt(text) + return { type: 'sent' } + } + const [name = '', ...rest] = text.slice(1).split(' ') + const argText = rest.join(' ').trim() + switch (name) { + case 'help': + return { type: 'notice', text: HELP_TEXT } + case 'exit': + case 'quit': + return { type: 'exit' } + case 'cancel': { + const receipt = await session.cancel() + return { + type: 'notice', + text: + receipt.status === 'rejected' ? 'Nothing is running.' : 'Cancelled.', + } + } + case 'status': { + const snapshot = session.snapshot() + const active = snapshot.activeOperations + .map((operation) => operation.agent ?? operation.kind) + .join(', ') + return { + type: 'notice', + text: `Status: ${snapshot.status}. Running: ${active || 'nothing'}. Queued turns: ${snapshot.queuedTurns}.`, + } + } + case 'agents': { + const agents = session.registry.list() + return { + type: 'notice', + text: + agents.length === 0 + ? 'This harness has no agents.' + : agents + .map((agent) => ` ${agent.name}: ${agent.description}`) + .join('\n'), + } + } + case 'agent': { + const [agentName = '', ...inputParts] = argText.split(' ') + const handle = session.agent(agentName) + if (!handle) + return { + type: 'notice', + text: `Unknown agent: ${agentName || '(none)'}`, + } + const inputText = inputParts.join(' ').trim() + let input: unknown + try { + input = inputText === '' ? undefined : JSON.parse(inputText) + } catch { + return { + type: 'notice', + text: 'The agent input must be JSON, for example {"vendor":"acme"}.', + } + } + const operation = handle.start(input, { wake: true }) + operation.then( + () => {}, + () => {}, + ) + return { type: 'notice', text: `Started ${agentName} in the background.` } + } + default: + return { type: 'notice', text: `Unknown command: /${name}. Type /help.` } + } +} diff --git a/packages/ai-harness-cli/src/index.ts b/packages/ai-harness-cli/src/index.ts new file mode 100644 index 0000000000..d91dd0d9d7 --- /dev/null +++ b/packages/ai-harness-cli/src/index.ts @@ -0,0 +1,111 @@ +import { createHarnessHost } from '@tanstack/ai-harness' +import { USAGE, parseCliArgs } from './args' +import { runLines } from './lines' +import { EXIT, runPrint } from './print' +import { createToken, serve } from './serve' +import type { AnyHarness, HarnessPersistence } from '@tanstack/ai-harness' + +export interface RunCliOptions { + /** Where sessions keep state. Default: in memory. */ + persistence?: HarnessPersistence + /** Default: `process.argv.slice(2)`. */ + argv?: ReadonlyArray + stdin?: NodeJS.ReadStream + stdout?: { write: (text: string) => unknown } + stderr?: { write: (text: string) => unknown } + env?: Record +} + +/** + * Run a harness from the terminal. Resolves to the process exit code. + * + * @example + * ```ts + * #!/usr/bin/env node + * process.exitCode = await runCli(studio) + * ``` + */ +export async function runCli( + harness: AnyHarness, + options: RunCliOptions = {}, +): Promise { + const stdout = options.stdout ?? process.stdout + const stderr = options.stderr ?? process.stderr + const stdin = options.stdin ?? process.stdin + const env = options.env ?? process.env + + let args + try { + args = parseCliArgs(options.argv ?? process.argv.slice(2)) + } catch (error) { + stderr.write( + `${error instanceof Error ? error.message : String(error)}\n\n${USAGE}\n`, + ) + return EXIT.failed + } + if (args.help) { + stdout.write(`${USAGE}\n`) + return EXIT.ok + } + + const host = createHarnessHost( + options.persistence ? { persistence: options.persistence } : {}, + ) + try { + if (args.acp) { + const { serveAcp } = await import('@tanstack/ai-acp/agent').catch(() => { + throw new Error( + '--acp needs @tanstack/ai-acp. Install it next to @tanstack/ai-harness-cli.', + ) + }) + const connection = serveAcp({ host, harness }) + await connection.closed + return EXIT.ok + } + + if (args.serve) { + const token = args.token ?? env.HARNESS_TOKEN ?? createToken() + const server = await serve({ + host, + harness, + port: args.port, + hostname: args.hostname, + token, + }) + stderr.write(`Serving ${harness.name} at ${server.url}\n`) + if (!args.token && !env.HARNESS_TOKEN) stderr.write(`Token: ${token}\n`) + await new Promise((resolve) => { + process.once('SIGINT', resolve) + process.once('SIGTERM', resolve) + }) + await server.close() + return EXIT.ok + } + + const session = await host.open(harness, { threadId: args.thread }) + if (args.print !== undefined) { + return await runPrint(session, args.print, { + output: args.output, + stdout, + stderr, + }) + } + if (stdin.isTTY) { + // Loaded only here, so the other modes never load React or Ink. + const { runInteractive } = await import('./interactive') + await runInteractive(session, harness) + } else { + await runLines(session, stdin, stdout) + } + return EXIT.ok + } catch (error) { + stderr.write(`${error instanceof Error ? error.message : String(error)}\n`) + return EXIT.failed + } finally { + await host.close() + } +} + +export { parseCliArgs, USAGE } from './args' +export type { CliArgs } from './args' +export { EXIT } from './print' diff --git a/packages/ai-harness-cli/src/interactive.tsx b/packages/ai-harness-cli/src/interactive.tsx new file mode 100644 index 0000000000..d05540ed51 --- /dev/null +++ b/packages/ai-harness-cli/src/interactive.tsx @@ -0,0 +1,170 @@ +import { useEffect, useState } from 'react' +import { Box, Text, render, useApp, useInput } from 'ink' +import { HARNESS_EVENTS } from '@tanstack/ai-harness' +import { EventType } from '@tanstack/ai' +import { handleLine } from './commands' +import { applyEvent, approvalQuestion, resolveAll } from './session-view' +import type { + AnyHarness, + HarnessSession, + SessionSnapshot, +} from '@tanstack/ai-harness' +import type { ViewEntry } from './session-view' + +const MAX_ENTRIES = 200 + +function Entry({ entry }: { entry: ViewEntry }) { + switch (entry.kind) { + case 'user': + return {`> ${entry.text}`} + case 'assistant': + return {entry.text} + case 'tool': + return {` - ${entry.text}`} + case 'notice': + return {entry.text} + } +} + +function App({ + session, + harness, +}: { + session: HarnessSession + harness: AnyHarness +}) { + const { exit } = useApp() + const [entries, setEntries] = useState>([ + { kind: 'notice', text: `${harness.name}. Type /help for commands.` }, + ]) + const [input, setInput] = useState('') + const [snapshot, setSnapshot] = useState(() => + session.snapshot(), + ) + + useEffect(() => { + const reader = new AbortController() + void (async () => { + for await (const entry of session.events({ + from: session.snapshot().cursor, + signal: reader.signal, + })) { + setEntries((current) => applyEvent(current, entry).slice(-MAX_ENTRIES)) + const event = entry.event + if ( + event.type === EventType.CUSTOM && + (event.name === HARNESS_EVENTS.operationStarted || + event.name === HARNESS_EVENTS.operationFinished) + ) { + // Let the operation settle before reading the status. + setTimeout(() => setSnapshot(session.snapshot()), 0) + } + } + })() + return () => reader.abort() + }, [session]) + + const notice = (text: string) => + setEntries((current) => + [...current, { kind: 'notice' as const, text }].slice(-MAX_ENTRIES), + ) + + const submit = async (line: string) => { + const current = session.snapshot() + if (current.status === 'requires_action') { + const approved = /^y(es)?$/i.test(line.trim()) + await resolveAll(session, current.pendingInterrupts, approved) + notice(approved ? 'Approved.' : 'Rejected.') + return + } + if (line.trim() !== '' && !line.trim().startsWith('/')) { + setEntries((list) => + [...list, { kind: 'user' as const, text: line.trim() }].slice( + -MAX_ENTRIES, + ), + ) + } + const result = await handleLine(session, line) + if (result.type === 'exit') exit() + else if (result.type === 'notice' && result.text) notice(result.text) + setSnapshot(session.snapshot()) + } + + useInput((character, key) => { + if (key.escape) { + void session.cancel().then(() => setSnapshot(session.snapshot())) + return + } + if (key.return) { + const line = input + setInput('') + void submit(line) + return + } + if (key.backspace || key.delete) { + setInput((current) => current.slice(0, -1)) + return + } + if (key.ctrl || key.meta || !character) return + // A paste (or fast typing) can arrive as one chunk with line breaks in it. + const [first = '', ...rest] = character.split(/\r\n|\r|\n/) + if (rest.length === 0) { + setInput((current) => current + first) + return + } + const lines = [input + first, ...rest.slice(0, -1)] + setInput(rest.at(-1) ?? '') + void (async () => { + for (const line of lines) await submit(line) + })() + }) + + const status = + snapshot.status === 'running' + ? 'working (Esc to cancel, Enter steers)' + : snapshot.status === 'requires_action' + ? approvalQuestion(snapshot.pendingInterrupts) + : 'ready' + const background = snapshot.activeOperations.filter( + (operation) => operation.kind === 'agent', + ) + + return ( + + {entries.map((entry, index) => ( + + ))} + + + {status} + {background.length > 0 + ? ` | agents: ${background.map((operation) => operation.agent).join(', ')}` + : ''} + {snapshot.queuedTurns > 0 ? ` | queued: ${snapshot.queuedTurns}` : ''} + + + + {'> '} + {input} + + + + ) +} + +/** Start the interactive UI. Resolves when the user quits. */ +export async function runInteractive( + session: HarnessSession, + harness: AnyHarness, + streams: { + stdin?: NodeJS.ReadStream + stdout?: NodeJS.WriteStream + /** Route `console` output above the UI. Default true. */ + patchConsole?: boolean + } = {}, +): Promise { + const instance = render(, streams) + await instance.waitUntilExit() +} diff --git a/packages/ai-harness-cli/src/lines.ts b/packages/ai-harness-cli/src/lines.ts new file mode 100644 index 0000000000..c16ff9e0ba --- /dev/null +++ b/packages/ai-harness-cli/src/lines.ts @@ -0,0 +1,80 @@ +import { createInterface } from 'node:readline' +import { handleLine } from './commands' +import { + applyEvent, + approvalQuestion, + resolveAll, + waitIdle, +} from './session-view' +import type { HarnessSession } from '@tanstack/ai-harness' +import type { ViewEntry } from './session-view' + +interface Output { + write: (text: string) => unknown +} + +/** + * The line mode, for piped input: one message or command per line. Each + * line waits for the turn it started. When a turn stops for approval, the + * next line answers it (`y` approves, anything else rejects). + */ +export async function runLines( + session: HarnessSession, + input: NodeJS.ReadableStream, + stdout: Output, +): Promise { + const reader = new AbortController() + let entries: Array = [] + let printed = 0 + const printing = (async () => { + for await (const entry of session.events({ + from: session.snapshot().cursor, + signal: reader.signal, + })) { + const next = applyEvent(entries, entry) + const last = next.at(-1) + if (next === entries || !last) continue + // Print the growth of the last assistant entry, or each new entry. + const previous = entries.at(-1) + if ( + last.kind === 'assistant' && + previous?.kind === 'assistant' && + previous.operationId === last.operationId + ) { + stdout.write(last.text.slice(previous.text.length)) + } else { + if (printed > 0) stdout.write('\n') + stdout.write(last.kind === 'assistant' ? last.text : `[${last.text}]`) + } + printed += 1 + entries = next + } + })() + + const lines = createInterface({ input, crlfDelay: Infinity }) + for await (const line of lines) { + const snapshot = session.snapshot() + if (snapshot.status === 'requires_action') { + await resolveAll( + session, + snapshot.pendingInterrupts, + /^y(es)?$/i.test(line.trim()), + ) + } else { + const result = await handleLine(session, line) + if (result.type === 'exit') break + if (result.type === 'notice' && result.text) + stdout.write(`${result.text}\n`) + } + await waitIdle(session) + const after = session.snapshot() + if (after.status === 'requires_action') { + stdout.write(`\n${approvalQuestion(after.pendingInterrupts)}\n`) + } + } + lines.close() + await waitIdle(session) + reader.abort() + await printing + if (printed > 0) stdout.write('\n') +} diff --git a/packages/ai-harness-cli/src/print.ts b/packages/ai-harness-cli/src/print.ts new file mode 100644 index 0000000000..757c278d94 --- /dev/null +++ b/packages/ai-harness-cli/src/print.ts @@ -0,0 +1,54 @@ +import { EventType } from '@tanstack/ai' +import type { HarnessSession } from '@tanstack/ai-harness' + +export const EXIT = { + ok: 0, + failed: 1, + needsAction: 2, + cancelled: 130, +} as const + +interface Output { + write: (text: string) => unknown +} + +/** + * Run one prompt and print it. `text` prints the answer as it streams. + * `ndjson` prints every AG-UI event as one JSON line. + */ +export async function runPrint( + session: HarnessSession, + prompt: string, + options: { output: 'text' | 'ndjson'; stdout: Output; stderr: Output }, +): Promise { + const { stdout, stderr } = options + const operation = session.prompt(prompt) + for await (const chunk of operation.stream()) { + if (options.output === 'ndjson') { + stdout.write(`${JSON.stringify(chunk)}\n`) + } else if ( + chunk.type === EventType.TEXT_MESSAGE_CONTENT && + !('subagentRunId' in chunk && chunk.subagentRunId) + ) { + stdout.write(chunk.delta) + } + } + const failure: unknown = await operation.then( + () => undefined, + (error: unknown) => error, + ) + if (options.output === 'text') stdout.write('\n') + const status = operation.status() + if (status === 'completed') return EXIT.ok + if (status === 'interrupted') { + stderr.write( + 'The turn stopped for an approval. Run the interactive mode to answer it.\n', + ) + return EXIT.needsAction + } + if (status === 'cancelled') return EXIT.cancelled + stderr.write( + `The turn failed: ${failure instanceof Error ? failure.message : String(failure)}\n`, + ) + return EXIT.failed +} diff --git a/packages/ai-harness-cli/src/serve.ts b/packages/ai-harness-cli/src/serve.ts new file mode 100644 index 0000000000..71bd801804 --- /dev/null +++ b/packages/ai-harness-cli/src/serve.ts @@ -0,0 +1,113 @@ +import { createServer } from 'node:http' +import { Readable } from 'node:stream' +import { randomBytes, timingSafeEqual } from 'node:crypto' +import { createHarnessHandler } from '@tanstack/ai-harness' +import type { IncomingMessage, ServerResponse } from 'node:http' +import type { ReadableStream as NodeReadableStream } from 'node:stream/web' +import type { AnyHarness, HarnessHost } from '@tanstack/ai-harness' + +export interface ServeOptions { + host: HarnessHost + harness: AnyHarness + port: number + hostname: string + /** Every request must send `Authorization: Bearer `. */ + token: string +} + +/** A random token for `--serve` when none is given. */ +export function createToken(): string { + return randomBytes(24).toString('base64url') +} + +function sameToken(given: string | null, expected: string): boolean { + if (!given?.startsWith('Bearer ')) return false + const a = Buffer.from(given.slice(7)) + const b = Buffer.from(expected) + return a.length === b.length && timingSafeEqual(a, b) +} + +async function toRequest( + req: IncomingMessage, + base: string, + signal: AbortSignal, +): Promise { + const headers = new Headers() + for (const [key, value] of Object.entries(req.headers)) { + if (typeof value === 'string') headers.set(key, value) + else if (Array.isArray(value)) headers.set(key, value.join(', ')) + } + const hasBody = req.method !== 'GET' && req.method !== 'HEAD' + const chunks: Array = [] + if (hasBody) for await (const chunk of req) chunks.push(Buffer.from(chunk)) + return new Request(new URL(req.url ?? '/', base), { + method: req.method ?? 'GET', + headers, + signal, + ...(hasBody ? { body: Buffer.concat(chunks) } : {}), + }) +} + +async function send(res: ServerResponse, response: Response): Promise { + res.writeHead(response.status, Object.fromEntries(response.headers)) + if (!response.body) { + res.end() + return + } + // `Readable.fromWeb` takes the stream type from node:stream/web. + Readable.fromWeb(response.body as NodeReadableStream).pipe(res) +} + +/** + * Serve the harness session protocol over HTTP (capabilities, run, events, + * control, snapshot). Every request needs the bearer token. + */ +export function serve( + options: ServeOptions, +): Promise<{ url: string; close: () => Promise }> { + const handler = createHarnessHandler({ + host: options.host, + harness: options.harness, + authorize: (request) => + sameToken(request.headers.get('authorization'), options.token) + ? { id: 'cli' } + : null, + }) + const server = createServer((req, res) => { + const aborter = new AbortController() + res.on('close', () => aborter.abort()) + void (async () => { + try { + const base = `http://${req.headers.host ?? `${options.hostname}:${options.port}`}` + await send( + res, + await handler(await toRequest(req, base, aborter.signal)), + ) + } catch (error) { + if (!res.headersSent) + res.writeHead(500, { 'Content-Type': 'application/json' }) + res.end( + JSON.stringify({ + error: error instanceof Error ? error.message : String(error), + }), + ) + } + })() + }) + return new Promise((resolve, reject) => { + server.once('error', reject) + server.listen(options.port, options.hostname, () => { + const address = server.address() + const port = + typeof address === 'object' && address ? address.port : options.port + resolve({ + url: `http://${options.hostname}:${port}`, + close: () => + new Promise((done) => { + server.closeAllConnections() + server.close(() => done()) + }), + }) + }) + }) +} diff --git a/packages/ai-harness-cli/src/session-view.ts b/packages/ai-harness-cli/src/session-view.ts new file mode 100644 index 0000000000..24325f072d --- /dev/null +++ b/packages/ai-harness-cli/src/session-view.ts @@ -0,0 +1,96 @@ +import { EventType } from '@tanstack/ai' +import { HARNESS_EVENTS } from '@tanstack/ai-harness' +import type { Interrupt, RunAgentResumeItem } from '@tanstack/ai' +import type { HarnessSession, SessionEvent } from '@tanstack/ai-harness' + +/** One line in the transcript view. */ +export type ViewEntry = + | { kind: 'user'; text: string } + | { kind: 'assistant'; operationId: string; text: string } + | { kind: 'tool'; text: string } + | { kind: 'notice'; text: string } + +/** + * Fold one session event into the view entries. Returns the new list, or the + * same list when the event changes nothing on screen. + */ +export function applyEvent( + entries: Array, + entry: SessionEvent, +): Array { + const { event, operationId } = entry + if ('subagentRunId' in event && event.subagentRunId) { + if (event.type === EventType.SUBAGENT_STARTED) { + return [ + ...entries, + { kind: 'tool', text: `agent ${event.name ?? ''} started` }, + ] + } + return entries + } + if (event.type === EventType.TEXT_MESSAGE_CONTENT) { + const last = entries.at(-1) + if (last?.kind === 'assistant' && last.operationId === operationId) { + return [ + ...entries.slice(0, -1), + { ...last, text: last.text + event.delta }, + ] + } + return [...entries, { kind: 'assistant', operationId, text: event.delta }] + } + if (event.type === EventType.TOOL_CALL_START) { + return [...entries, { kind: 'tool', text: `tool ${event.toolCallName}` }] + } + if (event.type === EventType.RUN_ERROR) { + return [...entries, { kind: 'notice', text: `Error: ${event.message}` }] + } + if ( + event.type === EventType.CUSTOM && + event.name === HARNESS_EVENTS.operationResumed + ) { + return [ + ...entries, + { kind: 'notice', text: 'Resumed a turn that a crash stopped.' }, + ] + } + return entries +} + +/** Answer every open interrupt of the last turn with one decision. */ +export function resolveAll( + session: HarnessSession, + interrupts: ReadonlyArray, + approved: boolean, +) { + const resume: Array = interrupts.map((interrupt) => ({ + interruptId: interrupt.id, + status: 'resolved', + payload: approved, + })) + return session.resolve(resume) +} + +/** A short question for the open interrupts. */ +export function approvalQuestion(interrupts: ReadonlyArray): string { + const names = interrupts.map( + (interrupt) => interrupt.message ?? interrupt.toolCallId ?? interrupt.id, + ) + return `Approve ${names.join(', ')}? [y/n]` +} + +/** Resolves when no chat turn runs or waits in the queue. */ +export async function waitIdle(session: HarnessSession): Promise { + while (true) { + const snapshot = session.snapshot() + const chatActive = snapshot.activeOperations.some( + (operation) => operation.kind === 'chat', + ) + if ( + snapshot.status !== 'running' && + !chatActive && + snapshot.queuedTurns === 0 + ) + return + await new Promise((resolve) => setTimeout(resolve, 20)) + } +} diff --git a/packages/ai-harness-cli/tests/cli.test.ts b/packages/ai-harness-cli/tests/cli.test.ts new file mode 100644 index 0000000000..645e5b9537 --- /dev/null +++ b/packages/ai-harness-cli/tests/cli.test.ts @@ -0,0 +1,265 @@ +import { Readable } from 'node:stream' +import { describe, expect, it, vi } from 'vitest' +import { EventType, defineAgent, toolDefinition } from '@tanstack/ai' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { EXIT, parseCliArgs, runCli } from '../src' +import { serve } from '../src/serve' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' + +const now = () => Date.now() +const textTurn = (text: string): Array => [ + { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, + { + type: EventType.TEXT_MESSAGE_START, + messageId: 'm', + role: 'assistant', + timestamp: now(), + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm', + delta: text, + timestamp: now(), + }, + { type: EventType.TEXT_MESSAGE_END, messageId: 'm', timestamp: now() }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'stop' } }, + }, +] +const toolTurn = (name: string, args: string): Array => [ + { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, + { + type: EventType.TOOL_CALL_START, + toolCallId: 'call_1', + toolCallName: name, + timestamp: now(), + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call_1', + delta: args, + timestamp: now(), + }, + { type: EventType.TOOL_CALL_END, toolCallId: 'call_1', timestamp: now() }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'tool_calls' } }, + }, +] + +function scripted(turns: Array>) { + let call = 0 + const seen: Array> = [] + const adapter: AnyTextAdapter = { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + chatStream: (options) => { + seen.push(options.messages) + const chunks = turns[call] ?? textTurn('') + call += 1 + return (async function* () { + yield* chunks + })() + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } + return { adapter, seen } +} + +function capture() { + let text = '' + return { + write: (chunk: string) => (text += chunk), + get text() { + return text + }, + } +} + +describe('parseCliArgs', () => { + it('reads flags and rejects bad values', () => { + expect(parseCliArgs(['-p', 'hi', '--output', 'ndjson'])).toMatchObject({ + print: 'hi', + output: 'ndjson', + thread: 'main', + }) + expect(() => parseCliArgs(['--output', 'xml'])).toThrow('--output') + expect(() => parseCliArgs(['--nope'])).toThrow() + }) +}) + +describe('print mode', () => { + it('prints the answer and exits 0', async () => { + const { adapter } = scripted([textTurn('Hello from the harness.')]) + const stdout = capture() + const code = await runCli(defineHarness({ name: 'test/print', adapter }), { + argv: ['-p', 'hi'], + stdout, + stderr: capture(), + persistence: memoryPersistence(), + }) + expect(code).toBe(EXIT.ok) + expect(stdout.text).toBe('Hello from the harness.\n') + }) + + it('prints NDJSON AG-UI events', async () => { + const { adapter } = scripted([textTurn('x')]) + const stdout = capture() + await runCli(defineHarness({ name: 'test/ndjson', adapter }), { + argv: ['-p', 'hi', '--output', 'ndjson'], + stdout, + stderr: capture(), + persistence: memoryPersistence(), + }) + const events = stdout.text + .trim() + .split('\n') + .map((line) => JSON.parse(line)) + expect(events.map((event) => event.type)).toContain('TEXT_MESSAGE_CONTENT') + expect(events.at(-1).name).toBe('harness.operation.finished') + }) + + it('exits 2 when the turn waits for an approval', async () => { + const { adapter } = scripted([toolTurn('remove', '{}')]) + const stderr = capture() + const code = await runCli( + defineHarness({ + name: 'test/approval', + adapter, + tools: [ + toolDefinition({ + name: 'remove', + description: 'Remove', + needsApproval: true, + }).server(async () => 'ok'), + ], + }), + { + argv: ['-p', 'remove it'], + stdout: capture(), + stderr, + persistence: memoryPersistence(), + }, + ) + expect(code).toBe(EXIT.needsAction) + expect(stderr.text).toContain('approval') + }) +}) + +describe('line mode', () => { + it('runs piped lines as turns, answers approvals, and runs commands', async () => { + const remove = vi.fn(async () => 'removed') + const pricer = defineAgent({ + name: 'pricer', + description: 'Prices a vendor', + run: async () => 'twelve dollars', + }) + const { adapter } = scripted([ + textTurn('First answer.'), + toolTurn('remove', '{}'), + textTurn('Done removing.'), + ]) + const stdout = capture() + const input = Readable.from([ + 'hello\n', + 'remove the file\n', + 'y\n', + '/agents\n', + '/exit\n', + ]) + const code = await runCli( + defineHarness({ + name: 'test/lines', + adapter, + agents: [pricer], + tools: [ + toolDefinition({ + name: 'remove', + description: 'Remove', + needsApproval: true, + }).server(remove), + ], + }), + { + argv: [], + stdin: Object.assign(input, { + isTTY: false, + }) as unknown as NodeJS.ReadStream, + stdout, + stderr: capture(), + persistence: memoryPersistence(), + }, + ) + expect(code).toBe(EXIT.ok) + expect(remove).toHaveBeenCalledTimes(1) + expect(stdout.text).toContain('First answer.') + expect(stdout.text).toContain('[y/n]') + expect(stdout.text).toContain('Done removing.') + expect(stdout.text).toContain('pricer: Prices a vendor') + }) +}) + +describe('serve mode', () => { + it('serves the session protocol and requires the token', async () => { + const { adapter } = scripted([textTurn('served')]) + const harness = defineHarness({ name: 'test/serve', adapter }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const server = await serve({ + host, + harness, + port: 0, + hostname: '127.0.0.1', + token: 'secret', + }) + try { + const denied = await fetch(`${server.url}/capabilities`) + expect(denied.status).toBe(401) + const wrong = await fetch(`${server.url}/capabilities`, { + headers: { authorization: 'Bearer nope' }, + }) + expect(wrong.status).toBe(401) + const capabilities = await fetch(`${server.url}/capabilities`, { + headers: { authorization: 'Bearer secret' }, + }) + expect((await capabilities.json()).identity.name).toBe('test/serve') + const receipt = await fetch(`${server.url}/control`, { + method: 'POST', + headers: { + authorization: 'Bearer secret', + 'content-type': 'application/json', + }, + body: JSON.stringify({ + threadId: 't', + input: { op: 'prompt', message: 'hi' }, + }), + }) + expect((await receipt.json()).status).toBe('accepted') + } finally { + await server.close() + await host.close() + } + }) +}) diff --git a/packages/ai-harness-cli/tests/interactive.test.ts b/packages/ai-harness-cli/tests/interactive.test.ts new file mode 100644 index 0000000000..9eaeb2845b --- /dev/null +++ b/packages/ai-harness-cli/tests/interactive.test.ts @@ -0,0 +1,106 @@ +import { PassThrough } from 'node:stream' +import { describe, expect, it } from 'vitest' +import { EventType } from '@tanstack/ai' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { runInteractive } from '../src/interactive' +import type { AnyTextAdapter } from '@tanstack/ai' + +const now = () => Date.now() +const adapter: AnyTextAdapter = { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + chatStream: () => + (async function* () { + yield { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: now(), + } + yield { + type: EventType.TEXT_MESSAGE_START, + messageId: 'm', + role: 'assistant', + timestamp: now(), + } + yield { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm', + delta: 'Hi from the Ink UI', + timestamp: now(), + } + yield { + type: EventType.TEXT_MESSAGE_END, + messageId: 'm', + timestamp: now(), + } + yield { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'stop' } }, + } + })(), + structuredOutput: async () => ({ data: {}, rawText: '{}' }), +} + +const wait = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)) + +describe('interactive UI', () => { + it('sends a typed line, shows the answer, and quits on /exit', async () => { + const harness = defineHarness({ name: 'test/ink', adapter }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(harness, { threadId: 'ink' }) + + const stdin = Object.assign(new PassThrough(), { + isTTY: true, + setRawMode: () => stdin, + ref: () => stdin, + unref: () => stdin, + }) + let output = '' + const stdout = Object.assign(new PassThrough(), { + isTTY: true, + columns: 100, + rows: 30, + }) + stdout.on('data', (data: Buffer) => { + output += data.toString() + }) + + const done = runInteractive(session, harness, { + stdin: stdin as unknown as NodeJS.ReadStream, + stdout: stdout as unknown as NodeJS.WriteStream, + // Vitest's console cannot be patched. + patchConsole: false, + }) + await wait(100) + // One chunk with a line break, the way a paste arrives. + stdin.write('hello\r') + await wait(300) + stdin.write('/exit\r') + await done + + const plain = output.replace(/\x1b\[[0-9;?]*[A-Za-z]/g, '') + expect(plain).toContain('> hello') + expect(plain).toContain('Hi from the Ink UI') + await host.close() + }) +}) diff --git a/packages/ai-harness-cli/tsconfig.json b/packages/ai-harness-cli/tsconfig.json new file mode 100644 index 0000000000..b9045b913b --- /dev/null +++ b/packages/ai-harness-cli/tsconfig.json @@ -0,0 +1,9 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "dist", + "jsx": "react-jsx" + }, + "include": ["vite.config.ts", "./src", "./tests"], + "exclude": ["node_modules", "dist"] +} diff --git a/packages/ai-harness-cli/vite.config.ts b/packages/ai-harness-cli/vite.config.ts new file mode 100644 index 0000000000..be372645e8 --- /dev/null +++ b/packages/ai-harness-cli/vite.config.ts @@ -0,0 +1,35 @@ +import { defineConfig, mergeConfig } from 'vitest/config' +import { tanstackViteConfig } from '@tanstack/vite-config' +import packageJson from './package.json' + +const config = defineConfig({ + test: { + name: packageJson.name, + dir: './', + watch: false, + globals: true, + environment: 'node', + include: ['tests/**/*.test.ts'], + coverage: { + provider: 'v8', + reporter: ['text', 'json', 'html', 'lcov'], + exclude: [ + 'node_modules/', + 'dist/', + 'tests/', + '**/*.test.ts', + '**/*.config.ts', + ], + include: ['src/**/*.{ts,tsx}'], + }, + }, +}) + +export default mergeConfig( + config, + tanstackViteConfig({ + entry: ['./src/index.ts'], + srcDir: './src', + cjs: false, + }), +) diff --git a/packages/ai-harness/package.json b/packages/ai-harness/package.json index 5d45fa9b9d..fd59be8e31 100644 --- a/packages/ai-harness/package.json +++ b/packages/ai-harness/package.json @@ -29,6 +29,10 @@ ".": { "types": "./dist/esm/index.d.ts", "import": "./dist/esm/index.js" + }, + "./client": { + "types": "./dist/esm/client.d.ts", + "import": "./dist/esm/client.js" } }, "scripts": { diff --git a/packages/ai-harness/src/client.ts b/packages/ai-harness/src/client.ts new file mode 100644 index 0000000000..cd7f6d055d --- /dev/null +++ b/packages/ai-harness/src/client.ts @@ -0,0 +1,198 @@ +// Browser-safe: this module has type-only imports from the rest of the +// package, so no server code (adapters, stores, secrets) reaches a client. +import type { RunAgentResumeItem } from '@tanstack/ai' +import type { AgentInputOf } from './agents' +import type { AnyHarness, HarnessAgentsOf } from './define' +import type { SessionSnapshot } from './session' +import type { + BusyPolicy, + Cursor, + HarnessInput, + Receipt, + SessionEvent, + UserInput, +} from './types' + +export interface HarnessClientOptions { + /** The base URL of `createHarnessHandler`, for example `/api/harness`. */ + url: string + threadId: string + /** Extra headers, for example `Authorization`. */ + headers?: Record | (() => Record) + fetch?: typeof fetch + /** Wait before a reconnect of `events()`. Default 1000 ms. */ + reconnectDelayMs?: number +} + +type StartArgs = + AgentInputOf extends undefined + ? [input?: undefined, options?: { detached?: boolean }] + : [input: AgentInputOf, options?: { detached?: boolean }] + +/** `client.agents`: start an exposed agent, typed from the harness. */ +export type ClientAgentHandles = { + [TAgent in HarnessAgentsOf as TAgent['name']]: { + start: (...args: StartArgs) => Promise + } +} + +export interface HarnessClient { + prompt: ( + message: UserInput, + options?: { busy?: BusyPolicy }, + ) => Promise + steer: (message: UserInput) => Promise + followUp: (message: UserInput) => Promise + resolve: (resume: Array) => Promise + cancel: (operationId?: string) => Promise + agents: ClientAgentHandles + /** + * The session events from `from` (exclusive). Reconnects after a network + * error and resumes from the last cursor. Ends when `signal` aborts. + */ + events: (options?: { + from?: Cursor + signal?: AbortSignal + }) => AsyncIterable + snapshot: () => Promise +} + +/** + * A client for a harness session served by `createHarnessHandler`. Pass the + * harness type for typed agents: `createHarnessClient(...)`. + * Import the harness with `import type`, so it stays out of the bundle. + */ +export function createHarnessClient( + options: HarnessClientOptions, +): HarnessClient { + const base = options.url.replace(/\/$/, '') + const doFetch = options.fetch ?? fetch + const headers = () => + typeof options.headers === 'function' + ? options.headers() + : (options.headers ?? {}) + + const send = async (input: HarnessInput): Promise => { + const response = await doFetch(`${base}/control`, { + method: 'POST', + headers: { 'Content-Type': 'application/json', ...headers() }, + body: JSON.stringify({ threadId: options.threadId, input }), + }) + const body: unknown = await response.json() + if (!response.ok) { + const message = + typeof body === 'object' && body !== null && 'error' in body + ? String(body.error) + : response.statusText + throw new Error(`Harness request failed (${response.status}): ${message}`) + } + // The handler answers /control with a Receipt. + return body as Receipt + } + + async function* events( + eventOptions: { from?: Cursor; signal?: AbortSignal } = {}, + ) { + let cursor = eventOptions.from + const signal = eventOptions.signal + while (!signal?.aborted) { + try { + const query = new URLSearchParams({ threadId: options.threadId }) + if (cursor) query.set('from', cursor) + const response = await doFetch(`${base}/events?${query}`, { + headers: headers(), + ...(signal ? { signal } : {}), + }) + if (!response.ok || !response.body) { + throw new Error(`Harness events failed (${response.status})`) + } + const reader = response.body.getReader() + // Some fetch shims ignore the signal once the body streams. + signal?.addEventListener('abort', () => void reader.cancel(), { + once: true, + }) + const decoder = new TextDecoder() + let buffer = '' + while (true) { + const { value, done } = await reader.read() + if (done) break + buffer += decoder.decode(value, { stream: true }) + const blocks = buffer.split('\n\n') + buffer = blocks.pop() ?? '' + for (const block of blocks) { + const data = block + .split('\n') + .find((line) => line.startsWith('data: ')) + ?.slice(6) + if (!data) continue + const frame: unknown = JSON.parse(data) + if ( + typeof frame === 'object' && + frame !== null && + 'type' in frame && + frame.type === 'harness.event' && + 'cursor' in frame && + typeof frame.cursor === 'string' + ) { + cursor = frame.cursor + // The handler sends `{ type, cursor, operationId, event }`. + const { type: _type, ...entry } = frame as SessionEvent & { + type: string + } + yield entry + } + } + } + } catch (error) { + if (signal?.aborted) return + if (error instanceof Error && /\((401|403|404)\)/.test(error.message)) + throw error + } + if (signal?.aborted) return + await new Promise((resolve) => + setTimeout(resolve, options.reconnectDelayMs ?? 1000), + ) + } + } + + const agents = new Proxy({} as ClientAgentHandles, { + get: (_target, name) => + typeof name === 'string' + ? { + start: (input?: unknown, startOptions?: { detached?: boolean }) => + send({ + op: 'agent', + agent: name, + input, + ...(startOptions?.detached ? { detached: true } : {}), + }), + } + : undefined, + }) + + return { + prompt: (message, promptOptions) => + send({ + op: 'prompt', + message, + ...(promptOptions?.busy ? { busy: promptOptions.busy } : {}), + }), + steer: (message) => send({ op: 'steer', message }), + followUp: (message) => send({ op: 'followUp', message }), + resolve: (resume) => send({ op: 'resolve', resume }), + cancel: (operationId) => + send({ op: 'cancel', ...(operationId ? { operationId } : {}) }), + agents, + events, + snapshot: async () => { + const query = new URLSearchParams({ threadId: options.threadId }) + const response = await doFetch(`${base}/snapshot?${query}`, { + headers: headers(), + }) + if (!response.ok) + throw new Error(`Harness snapshot failed (${response.status})`) + // The handler answers /snapshot with a SessionSnapshot. + return (await response.json()) as SessionSnapshot + }, + } +} diff --git a/packages/ai-harness/src/harness-text.ts b/packages/ai-harness/src/harness-text.ts new file mode 100644 index 0000000000..51b0aa7a21 --- /dev/null +++ b/packages/ai-harness/src/harness-text.ts @@ -0,0 +1,133 @@ +import { EventType } from '@tanstack/ai' +import { createHarnessHost } from './host' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import type { AnyHarness } from './define' +import type { HarnessHost } from './host' + +export interface HarnessTextOptions { + /** The host that runs the inner sessions. Default: a memory host. */ + host?: HarnessHost +} + +const isRecord = (value: unknown): value is Record => + typeof value === 'object' && value !== null + +function lastUserText(messages: unknown): string { + const list: ReadonlyArray = Array.isArray(messages) ? messages : [] + const message = list.findLast( + (entry) => isRecord(entry) && entry.role === 'user', + ) + if (!isRecord(message)) return '' + if (typeof message.content === 'string') return message.content + const parts: ReadonlyArray = Array.isArray(message.content) + ? message.content + : [] + return parts + .map((part) => + isRecord(part) && part.type === 'text' && typeof part.content === 'string' + ? part.content + : '', + ) + .join('') +} + +/** Events of the inner turn that the outer chat shows as its own model output. */ +const FORWARDED = new Set([ + EventType.TEXT_MESSAGE_START, + EventType.TEXT_MESSAGE_CONTENT, + EventType.TEXT_MESSAGE_END, + EventType.REASONING_START, + EventType.REASONING_MESSAGE_START, + EventType.REASONING_MESSAGE_CONTENT, + EventType.REASONING_MESSAGE_END, + EventType.REASONING_END, +]) + +/** + * Use a harness as the model of a `chat()` call, the way you would call a + * coding agent. Each outer thread gets its own inner session, which keeps + * its own transcript, tools, plugins, and agents. The outer chat sees the + * inner turn's text and reasoning. + * + * @example + * ```ts + * const stream = chat({ adapter: harnessText(studio), messages, threadId }) + * ``` + */ +export function harnessText( + harness: AnyHarness, + options: HarnessTextOptions = {}, +): AnyTextAdapter { + let host = options.host + return { + kind: 'text', + name: 'harness', + model: harness.name, + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + chatStream: (chatOptions) => + (async function* (): AsyncGenerator { + host ??= createHarnessHost() + const threadId = `${chatOptions.threadId ?? 'default'}:${harness.name}` + const runId = chatOptions.runId ?? `harness-${Date.now().toString(36)}` + const signal = + isRecord(chatOptions.request) && + chatOptions.request.signal instanceof AbortSignal + ? chatOptions.request.signal + : undefined + const session = await host.open(harness, { threadId }) + const operation = session.prompt(lastUserText(chatOptions.messages)) + signal?.addEventListener('abort', () => void operation.cancel(), { + once: true, + }) + + yield { + type: EventType.RUN_STARTED, + runId, + threadId, + timestamp: Date.now(), + } + let failed: string | undefined + for await (const chunk of operation.stream()) { + if ( + FORWARDED.has(chunk.type) && + !('subagentRunId' in chunk && chunk.subagentRunId) + ) { + yield chunk + } + if (chunk.type === EventType.RUN_ERROR) failed = chunk.message + } + if (failed !== undefined) { + yield { + type: EventType.RUN_ERROR, + message: failed, + timestamp: Date.now(), + } + return + } + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + metadata: { tanstack: { finishReason: 'stop' } }, + } + })(), + structuredOutput: () => + Promise.reject( + new Error('harnessText does not support structured output.'), + ), + } +} diff --git a/packages/ai-harness/src/host.ts b/packages/ai-harness/src/host.ts index dda6dd628e..dd8dafb6c7 100644 --- a/packages/ai-harness/src/host.ts +++ b/packages/ai-harness/src/host.ts @@ -69,6 +69,7 @@ export function createHarnessHost( // to keep accepted inputs across restarts. const inbox = persistence.stores.inbox ?? memoryPersistence().stores.inbox const sessions = new Map>() + const hostId = `host-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 9)}` return { open(harness, { threadId, principal }) { @@ -80,6 +81,7 @@ export function createHarnessHost( threadId, persistence, inbox, + hostId, ...(principal ? { principal } : {}), onClose: () => sessions.delete(key), }) diff --git a/packages/ai-harness/src/http.ts b/packages/ai-harness/src/http.ts new file mode 100644 index 0000000000..e9be9a5e24 --- /dev/null +++ b/packages/ai-harness/src/http.ts @@ -0,0 +1,308 @@ +import { + EventType, + chatParamsFromRequestBody, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { + HARNESS_PROTOCOL_VERSION, + applyInput, + capabilitiesOf, + parseControlFrame, + parseHarnessInput, +} from './protocol' +import type { StreamChunk, WebSocketLike } from '@tanstack/ai' +import type { AnyHarness } from './define' +import type { HarnessHost } from './host' +import type { HostFrame } from './protocol' +import type { Principal } from './types' + +/** + * Decide who sends a request. Return `null` to refuse it with 401. Every + * endpoint calls it, so there is no unauthenticated route. + */ +export type Authorize = ( + request: Request, +) => Principal | null | Promise + +export interface HarnessHandlerOptions { + host: HarnessHost + harness: AnyHarness + authorize: Authorize + /** + * May this principal use this thread? Default: yes. Use it to keep users + * out of each other's threads. + */ + canAccess?: ( + principal: Principal, + threadId: string, + ) => boolean | Promise +} + +const json = (body: unknown, status = 200) => + new Response(JSON.stringify(body), { + status, + headers: { 'Content-Type': 'application/json' }, + }) + +const isRecord = (value: unknown): value is Record => + typeof value === 'object' && value !== null + +/** The text of the last user message of an AG-UI request. */ +function lastUserText(messages: ReadonlyArray): string | undefined { + const message = messages.findLast( + (entry) => isRecord(entry) && entry.role === 'user', + ) + if (!isRecord(message)) return undefined + if (typeof message.content === 'string') return message.content + const parts = Array.isArray(message.parts) ? message.parts : [] + return parts + .map((part: unknown) => + isRecord(part) && part.type === 'text' && typeof part.content === 'string' + ? part.content + : '', + ) + .join('') +} + +/** + * A fetch handler for a harness. Mount it on a route that ends with any of + * these paths: + * + * - `GET .../capabilities`: the AG-UI capabilities document. + * - `POST .../run`: standard AG-UI. One request is one prompt, streamed as SSE. + * - `GET .../events?threadId=&from=`: the session stream as SSE. Each event id + * is a cursor, so `Last-Event-ID` resumes it. + * - `POST .../control`: `{ threadId, input }`. Returns the receipt. + * - `GET .../snapshot?threadId=`: the session snapshot. + */ +export function createHarnessHandler( + options: HarnessHandlerOptions, +): (request: Request) => Promise { + const { host, harness, authorize } = options + const canAccess = options.canAccess ?? (() => true) + + const openFor = async (principal: Principal, threadId: string) => { + if (!(await canAccess(principal, threadId))) return null + return host.open(harness, { threadId, principal }) + } + + return async (request) => { + const principal = await authorize(request) + if (!principal) return json({ error: 'unauthorized' }, 401) + const url = new URL(request.url) + const route = url.pathname.split('/').at(-1) + + try { + if (request.method === 'GET' && route === 'capabilities') { + return json(capabilitiesOf(harness)) + } + + if (request.method === 'POST' && route === 'run') { + const params = await chatParamsFromRequestBody(await request.json()) + const message = lastUserText(params.messages) + const session = await openFor(principal, params.threadId) + if (!session) return json({ error: 'forbidden' }, 403) + if (params.resume && params.resume.length > 0) { + const receipt = await session.resolve(params.resume) + if (receipt.status === 'rejected' || !receipt.operationId) { + return json({ error: receipt.reason ?? 'rejected' }, 409) + } + // Not `request.signal`: some servers abort it once the body is + // read. The response aborts this controller when the client leaves. + const reader = new AbortController() + const events = session.events({ signal: reader.signal }) + return toServerSentEventsResponse( + followOperation(events, receipt.operationId), + { + abortController: reader, + }, + ) + } + if (message === undefined) + return json({ error: 'no user message' }, 400) + const operation = session.prompt(message) + const reader = new AbortController() + return toServerSentEventsResponse( + operation.stream({ signal: reader.signal }), + { + abortController: reader, + }, + ) + } + + if (request.method === 'GET' && route === 'events') { + const threadId = url.searchParams.get('threadId') + if (!threadId) return json({ error: 'threadId is required' }, 400) + const session = await openFor(principal, threadId) + if (!session) return json({ error: 'forbidden' }, 403) + const from = + request.headers.get('Last-Event-ID') ?? + url.searchParams.get('from') ?? + undefined + const encoder = new TextEncoder() + // Stops when the client cancels the response stream. + const reader = new AbortController() + const body = new ReadableStream({ + cancel: () => reader.abort(), + async start(controller) { + try { + for await (const entry of session.events({ + ...(from ? { from } : {}), + signal: reader.signal, + })) { + const frame: HostFrame = { type: 'harness.event', ...entry } + controller.enqueue( + encoder.encode( + `id: ${entry.cursor}\ndata: ${JSON.stringify(frame)}\n\n`, + ), + ) + } + } finally { + if (!reader.signal.aborted) controller.close() + } + }, + }) + return new Response(body, { + headers: { + 'Content-Type': 'text/event-stream', + 'Cache-Control': 'no-cache', + Connection: 'keep-alive', + }, + }) + } + + if (request.method === 'POST' && route === 'control') { + const body: unknown = await request.json() + if ( + typeof body !== 'object' || + body === null || + !('threadId' in body) + ) { + return json({ error: 'threadId is required' }, 400) + } + const threadId = body.threadId + if (typeof threadId !== 'string') + return json({ error: 'threadId is required' }, 400) + const input = parseHarnessInput( + 'input' in body ? body.input : undefined, + ) + const session = await openFor(principal, threadId) + if (!session) return json({ error: 'forbidden' }, 403) + return json(await applyInput(harness, session, input)) + } + + if (request.method === 'GET' && route === 'snapshot') { + const threadId = url.searchParams.get('threadId') + if (!threadId) return json({ error: 'threadId is required' }, 400) + const session = await openFor(principal, threadId) + if (!session) return json({ error: 'forbidden' }, 403) + return json(session.snapshot()) + } + } catch (error) { + return json( + { error: error instanceof Error ? error.message : String(error) }, + 400, + ) + } + return json({ error: 'not found' }, 404) + } +} + +/** Session events of one operation, until it finishes. */ +async function* followOperation( + events: AsyncIterable<{ operationId: string; event: StreamChunk }>, + operationId: string, +) { + for await (const entry of events) { + if (entry.operationId !== operationId) continue + yield entry.event + if ( + entry.event.type === EventType.CUSTOM && + entry.event.name === 'harness.operation.finished' + ) { + return + } + } +} + +export interface HarnessSocketOptions { + host: HarnessHost + harness: AnyHarness + socket: WebSocketLike + /** The principal your upgrade handler authorized. */ + principal: Principal + canAccess?: ( + principal: Principal, + threadId: string, + ) => boolean | Promise +} + +/** + * Serve the session tier over one WebSocket. The first frame must be + * `harness.subscribe`. Authorize the upgrade request before you call this. + */ +export function handleHarnessSocket(options: HarnessSocketOptions): void { + const { host, harness, socket, principal } = options + const canAccess = options.canAccess ?? (() => true) + const reader = new AbortController() + let session: Awaited> | undefined + const send = (frame: HostFrame) => { + try { + socket.send(JSON.stringify(frame)) + } catch { + reader.abort() + } + } + + socket.addEventListener('close', () => reader.abort()) + socket.addEventListener('error', () => reader.abort()) + socket.addEventListener('message', (message) => { + void (async () => { + try { + const frame = parseControlFrame(String(message.data)) + if (frame.type === 'harness.subscribe') { + if (session) throw new Error('Already subscribed.') + if (!(await canAccess(principal, frame.threadId))) { + send({ type: 'harness.error', message: 'forbidden' }) + socket.close(4403, 'forbidden') + return + } + session = await host.open(harness, { + threadId: frame.threadId, + principal, + }) + send({ + type: 'harness.hello', + v: HARNESS_PROTOCOL_VERSION, + threadId: frame.threadId, + }) + const events = session.events({ + ...(frame.from ? { from: frame.from } : {}), + signal: reader.signal, + }) + void (async () => { + for await (const entry of events) + send({ type: 'harness.event', ...entry }) + })() + return + } + if (!session) throw new Error('Send harness.subscribe first.') + if (frame.type === 'harness.snapshot') { + send({ type: 'harness.snapshot', snapshot: session.snapshot() }) + return + } + const receipt = await applyInput(harness, session, frame.input) + send({ + type: 'harness.receipt', + requestId: frame.requestId, + ...receipt, + }) + } catch (error) { + send({ + type: 'harness.error', + message: error instanceof Error ? error.message : String(error), + }) + } + })() + }) +} diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts index ac889765b3..e1de83cca8 100644 --- a/packages/ai-harness/src/index.ts +++ b/packages/ai-harness/src/index.ts @@ -38,6 +38,7 @@ export type { AgentHandles, AgentRunOptions, AgentStartOptions, + DynamicAgentHandle, SessionSnapshot, } from './session' @@ -55,3 +56,22 @@ export type { SessionEvent, UserInput, } from './types' + +export { + HARNESS_PROTOCOL_VERSION, + applyInput, + capabilitiesOf, + parseControlFrame, + parseHarnessInput, +} from './protocol' +export type { ControlFrame, HostFrame } from './protocol' + +export { createHarnessHandler, handleHarnessSocket } from './http' +export type { + Authorize, + HarnessHandlerOptions, + HarnessSocketOptions, +} from './http' + +export { harnessText } from './harness-text' +export type { HarnessTextOptions } from './harness-text' diff --git a/packages/ai-harness/src/protocol.ts b/packages/ai-harness/src/protocol.ts new file mode 100644 index 0000000000..3532454792 --- /dev/null +++ b/packages/ai-harness/src/protocol.ts @@ -0,0 +1,200 @@ +import { convertSchemaToJsonSchema } from '@tanstack/ai' +import type { StreamChunk } from '@tanstack/ai' +import type { AnyAgent } from './agents' +import type { AnyHarness } from './define' +import type { HarnessSession, SessionSnapshot } from './session' +import type { Cursor, HarnessInput, Receipt } from './types' + +/** The session-tier protocol version. Sent in `subscribe` and `hello`. */ +export const HARNESS_PROTOCOL_VERSION = 1 + +/** Client to host. */ +export type ControlFrame = + | { type: 'harness.subscribe'; threadId: string; from?: Cursor; v?: number } + | { type: 'harness.input'; requestId: string; input: HarnessInput } + | { type: 'harness.snapshot' } + +/** Host to client. */ +export type HostFrame = + | { type: 'harness.hello'; v: number; threadId: string } + | { + type: 'harness.receipt' + requestId: string + status: Receipt['status'] + inputId?: string + operationId?: string + reason?: string + } + | { + type: 'harness.event' + cursor: Cursor + operationId: string + event: StreamChunk + } + | { type: 'harness.snapshot'; snapshot: SessionSnapshot } + | { type: 'harness.error'; message: string } + +const INPUT_OPS = new Set([ + 'prompt', + 'steer', + 'followUp', + 'resolve', + 'agent', + 'cancel', +]) + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null +} + +/** Check the shape of a client input. Throws with a short reason. */ +export function parseHarnessInput(value: unknown): HarnessInput { + if ( + !isRecord(value) || + typeof value.op !== 'string' || + !INPUT_OPS.has(value.op) + ) { + throw new Error('Invalid input: expected { op } with a known op.') + } + const needsMessage = + value.op === 'prompt' || value.op === 'steer' || value.op === 'followUp' + if ( + needsMessage && + typeof value.message !== 'string' && + !Array.isArray(value.message) + ) { + throw new Error(`Invalid input: ${value.op} needs a message.`) + } + if (value.op === 'resolve' && !Array.isArray(value.resume)) { + throw new Error('Invalid input: resolve needs a resume array.') + } + if (value.op === 'agent' && typeof value.agent !== 'string') { + throw new Error('Invalid input: agent needs an agent name.') + } + // The checks above cover every field the session reads. + return value as HarnessInput +} + +/** Parse one text frame from a client. Throws with a short reason. */ +export function parseControlFrame(data: string): ControlFrame { + const value: unknown = JSON.parse(data) + if (!isRecord(value)) throw new Error('Invalid frame.') + if ( + value.type === 'harness.subscribe' && + typeof value.threadId === 'string' + ) { + return { + type: 'harness.subscribe', + threadId: value.threadId, + ...(typeof value.from === 'string' ? { from: value.from } : {}), + } + } + if (value.type === 'harness.input' && typeof value.requestId === 'string') { + return { + type: 'harness.input', + requestId: value.requestId, + input: parseHarnessInput(value.input), + } + } + if (value.type === 'harness.snapshot') return { type: 'harness.snapshot' } + throw new Error('Invalid frame: unknown type.') +} + +/** + * Apply a client input to a session. Only agents in `expose.agents` can run + * from a client. Resolves to the receipt. + */ +export async function applyInput( + harness: AnyHarness, + session: HarnessSession, + input: HarnessInput, +): Promise { + switch (input.op) { + case 'prompt': { + const operation = session.prompt( + input.message, + input.busy ? { busy: input.busy } : {}, + ) + const queued = session + .snapshot() + .activeOperations.some( + (active) => active.kind === 'chat' && active.id !== operation.id, + ) + return { + inputId: operation.id, + status: queued ? 'queued' : 'accepted', + operationId: operation.id, + } + } + case 'steer': + return session.steer(input.message) + case 'followUp': + return session.followUp(input.message) + case 'resolve': + return session.resolve(input.resume) + case 'cancel': + return session.cancel(input.operationId) + case 'agent': { + const exposed = (harness.expose?.agents ?? []).includes(input.agent) + if (!exposed) { + return { inputId: '', status: 'rejected', reason: 'not_exposed' } + } + const handle = session.agent(input.agent) + if (!handle) { + return { inputId: '', status: 'rejected', reason: 'unknown_agent' } + } + const operation = handle.start(input.input, { + wake: input.detached === true, + }) + return { + inputId: operation.id, + status: 'accepted', + operationId: operation.id, + } + } + } +} + +/** The AG-UI capabilities document of a harness. */ +export function capabilitiesOf(harness: AnyHarness) { + const exposed = new Set(harness.expose?.agents ?? []) + const subagents: ReadonlyArray = harness.subagents?.agents ?? [] + const agents: ReadonlyArray = [ + ...(harness.agents ?? []), + ...subagents, + ] + return { + identity: { name: harness.name, type: 'tanstack-ai-harness' }, + transport: { streaming: true, websocket: true }, + tools: { + supported: true, + items: (harness.tools ?? []).map((tool) => ({ + name: tool.name, + description: tool.description, + })), + }, + multiAgent: { + supported: agents.length > 0, + subagents: subagents.map((agent) => ({ + name: agent.name, + description: agent.description, + })), + }, + humanInTheLoop: { supported: true, interrupts: true }, + custom: { + tanstack: { + protocol: HARNESS_PROTOCOL_VERSION, + agents: agents + .filter((agent) => exposed.has(agent.name)) + .map((agent) => ({ + name: agent.name, + description: agent.description, + ...(agent.produces ? { produces: agent.produces } : {}), + ...(agent.inputSchema + ? { inputSchema: convertSchemaToJsonSchema(agent.inputSchema) } + : {}), + })), + }, + }, + } +} diff --git a/packages/ai-harness/src/resume.ts b/packages/ai-harness/src/resume.ts new file mode 100644 index 0000000000..ea695840e0 --- /dev/null +++ b/packages/ai-harness/src/resume.ts @@ -0,0 +1,146 @@ +import type { AnyChatMiddleware, ModelMessage, RunRecord } from '@tanstack/ai' +import type { HarnessPersistence } from './host' + +/** How often a running host renews its lease, and when a lease expires. */ +export const LEASE = { renewMs: 10_000, ttlMs: 30_000 } + +/** The tool result a crash leaves for a tool that must not run twice. */ +export const INTERRUPTED_TOOL_RESULT = { + interrupted: true, + note: 'The tool may or may not have run. Check before you retry.', +} + +type PendingTool = { + toolCallId: string + name: string + replay: 'safe' | 'never' +} + +/** + * Chat middleware that makes a turn resumable after a crash: + * + * - holds a lease on the run record and renews it while the turn runs; + * - saves the transcript before and after each tool phase; + * - records each tool call that started but has no result yet. + */ +export function checkpointMiddleware( + persistence: HarnessPersistence, + hostId: string, +): AnyChatMiddleware { + const runs = persistence.stores.runs + const messages = persistence.stores.messages + const state = new WeakMap< + object, + { timer: ReturnType; pending: Array } + >() + + const saveCheckpoint = async (runId: string, pending: Array) => { + await runs?.update(runId, { + checkpoint: { at: Date.now(), pendingTools: [...pending] }, + }) + } + const stop = (ctx: object) => { + const entry = state.get(ctx) + if (entry) clearInterval(entry.timer) + state.delete(ctx) + } + + return { + name: 'harness:checkpoint', + async onStart(ctx) { + const renew = () => + runs?.update(ctx.runId, { + leaseOwner: hostId, + leaseExpiresAt: Date.now() + LEASE.ttlMs, + }) + await renew() + const timer = setInterval(() => void renew(), LEASE.renewMs) + // A lease timer must not keep a CLI or a test process alive. + if (typeof timer === 'object' && 'unref' in timer) timer.unref() + state.set(ctx, { timer, pending: [] }) + }, + async onBeforeToolCall(ctx, hook) { + const entry = state.get(ctx) + if (!entry) return + // The first call of a phase: the transcript now holds the assistant + // message with the tool calls. Save it, so resume can find them. + if (entry.pending.length === 0) { + await messages.saveThread(ctx.threadId, [...ctx.messages]) + } + entry.pending.push({ + toolCallId: hook.toolCallId, + name: hook.toolName, + replay: hook.tool?.replay ?? 'never', + }) + await saveCheckpoint(ctx.runId, entry.pending) + }, + async onAfterToolCall(ctx, info) { + const entry = state.get(ctx) + if (!entry) return + entry.pending = entry.pending.filter( + (tool) => tool.toolCallId !== info.toolCallId, + ) + }, + async onToolPhaseComplete(ctx) { + const entry = state.get(ctx) + if (!entry) return + await messages.saveThread(ctx.threadId, [...ctx.messages]) + entry.pending = [] + await saveCheckpoint(ctx.runId, []) + }, + onFinish: (ctx) => stop(ctx), + onAbort: (ctx) => stop(ctx), + onError: (ctx) => stop(ctx), + } +} + +/** Chat runs of this thread that a crashed host left `running`. */ +export async function findCrashedRuns( + persistence: HarnessPersistence, + threadId: string, + now = Date.now(), +): Promise> { + const runs = persistence.stores.runs + if (!runs?.listByThread) return [] + const records = await runs.listByThread(threadId) + return records.filter( + (record) => + record.status === 'running' && + (record.kind === undefined || record.kind === 'chat') && + record.leaseExpiresAt !== undefined && + record.leaseExpiresAt < now, + ) +} + +/** + * Prepare the transcript of a crashed run for a new run. A tool with + * `replay: 'never'` gets {@link INTERRUPTED_TOOL_RESULT}. A tool with + * `replay: 'safe'` stays without a result, so the engine runs it again. + */ +export async function repairTranscript( + persistence: HarnessPersistence, + crashed: RunRecord, +): Promise { + const pending = crashed.checkpoint?.pendingTools ?? [] + const never = pending.filter((tool) => tool.replay === 'never') + if (never.length === 0) return + const store = persistence.stores.messages + const history = await store.loadThread(crashed.threadId) + const answered = new Set( + history.flatMap((message) => + message.role === 'tool' && message.toolCallId ? [message.toolCallId] : [], + ), + ) + const notes = never + .filter((tool) => !answered.has(tool.toolCallId)) + .map( + (tool): ModelMessage => ({ + role: 'tool', + toolCallId: tool.toolCallId, + content: JSON.stringify(INTERRUPTED_TOOL_RESULT), + }), + ) + if (notes.length > 0) { + await store.saveThread(crashed.threadId, [...history, ...notes]) + } +} diff --git a/packages/ai-harness/src/session.ts b/packages/ai-harness/src/session.ts index 9e5a74b5c3..3e638d0f31 100644 --- a/packages/ai-harness/src/session.ts +++ b/packages/ai-harness/src/session.ts @@ -12,6 +12,11 @@ import { AgentRegistry } from './agents' import { SessionFeed } from './feed' import { OperationImpl } from './operation' import { mountPlugins } from './plugins' +import { + checkpointMiddleware, + findCrashedRuns, + repairTranscript, +} from './resume' import { HARNESS_EVENTS } from './types' import type { AnyChatMiddleware, @@ -74,6 +79,12 @@ export interface AgentHandle { ) => Operation> } +/** A handle for an agent picked by name at runtime. */ +export interface DynamicAgentHandle { + run: (input?: unknown, options?: AgentRunOptions) => Operation + start: (input?: unknown, options?: AgentStartOptions) => Operation +} + /** `session.agents`: one typed handle per registered agent name. */ export type AgentHandles = { [TAgent in HarnessAgentsOf as TAgent['name']]: AgentHandle @@ -101,6 +112,8 @@ export interface SessionDependencies { persistence: HarnessPersistence inbox: InboxStore principal?: Principal + /** Identifies this host on run leases. */ + hostId: string onClose: () => void } @@ -173,6 +186,7 @@ export class HarnessSession { private sessionPlugins: MountedPlugins | undefined private closing: Promise | undefined private readonly onClose: () => void + private readonly checkpoint: AnyChatMiddleware constructor(deps: SessionDependencies) { this.harness = deps.harness as THarness @@ -181,6 +195,7 @@ export class HarnessSession { this.inbox = deps.inbox this.principal = deps.principal this.onClose = deps.onClose + this.checkpoint = checkpointMiddleware(deps.persistence, deps.hostId) this.registry = this.agentRegistry for (const agent of this.harness.agents ?? []) { this.agentRegistry.add(agent, 'the harness') @@ -201,6 +216,24 @@ export class HarnessSession { }) } + /** + * The agent named `name`, for names known only at runtime (a slash + * command, a protocol input). `undefined` when no such agent exists. + */ + agent(name: string): DynamicAgentHandle | undefined { + if (!this.agentRegistry.get(name)) return undefined + return { + run: (input, options) => + this.runAgent(name, input, { ...options, wake: false }), + start: (input, options) => this.runAgent(name, input, options ?? {}), + } + } + + /** An operation of this session by id, running or settled. */ + operation(id: string): Operation | undefined { + return this.operations.get(id) + } + /** @internal Mount session plugins and replay inputs left in the inbox. */ async open(): Promise { this.plugins = this.harness.plugins?.() ?? [] @@ -217,6 +250,7 @@ export class HarnessSession { ), }, ) + await this.recoverCrashedTurn() await this.recoverInbox() } @@ -513,6 +547,7 @@ export class HarnessSession { middleware: [ ...bridges, withPersistence(this.persistence), + this.checkpoint, ...(this.harness.middleware ?? []), ...(session?.middleware ?? []), ...(runPlugins?.middleware ?? []), @@ -844,6 +879,38 @@ export class HarnessSession { ) } + /** + * Continue the newest chat turn that a crashed host left running. Older + * crashed turns are marked failed. + */ + private async recoverCrashedTurn(): Promise { + const crashed = await findCrashedRuns(this.persistence, this.threadId) + const newest = crashed.sort((a, b) => b.startedAt - a.startedAt)[0] + for (const record of crashed) { + await this.persistence.stores.runs?.update(record.runId, { + status: 'failed', + finishedAt: Date.now(), + error: { + message: + record === newest + ? 'The host stopped. The session continued this turn in a new run.' + : 'The host stopped during this turn.', + }, + }) + } + if (!newest) return + await repairTranscript(this.persistence, newest) + const operation = this.createTurnOperation() + this.feed.publish( + operation.id, + customEvent(HARNESS_EVENTS.operationResumed, { + operationId: operation.id, + resumedFrom: newest.runId, + }), + ) + this.enqueueTurn({ operation }) + } + /** Re-run turns that were accepted but never applied before a restart. */ private async recoverInbox(): Promise { const pending: Array = await this.inbox.listPending( diff --git a/packages/ai-harness/src/types.ts b/packages/ai-harness/src/types.ts index 5d6b766ae9..43c65515b7 100644 --- a/packages/ai-harness/src/types.ts +++ b/packages/ai-harness/src/types.ts @@ -93,6 +93,7 @@ export interface Operation extends PromiseLike { export const HARNESS_EVENTS = { operationStarted: 'harness.operation.started', operationFinished: 'harness.operation.finished', + operationResumed: 'harness.operation.resumed', inputAccepted: 'harness.input.accepted', inputApplied: 'harness.input.applied', inputRejected: 'harness.input.rejected', diff --git a/packages/ai-harness/tests/client.test.ts b/packages/ai-harness/tests/client.test.ts new file mode 100644 index 0000000000..8f45fd9961 --- /dev/null +++ b/packages/ai-harness/tests/client.test.ts @@ -0,0 +1,87 @@ +import { describe, expect, it } from 'vitest' +import { z } from 'zod' +import { defineAgent } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHandler, createHarnessHost, defineHarness } from '../src' +import { createHarnessClient } from '../src/client' +import { mockAdapter, text } from './helpers' +import type { SessionEvent } from '../src' + +describe('createHarnessClient', () => { + it('sends inputs, reads events, and gets a snapshot through the handler', async () => { + const pricer = defineAgent({ + name: 'pricer', + description: 'Prices a vendor', + inputSchema: z.object({ vendor: z.string() }), + run: async (ctx) => ({ cents: ctx.input.vendor.length }), + }) + const { adapter } = mockAdapter([() => text('hello from the host')]) + const studio = defineHarness({ + name: 'test/client', + adapter, + agents: [pricer], + expose: { agents: ['pricer'] }, + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const handler = createHarnessHandler({ + host, + harness: studio, + authorize: (request) => + request.headers.get('authorization') === 'Bearer t' + ? { id: 'u' } + : null, + }) + const client = createHarnessClient({ + url: 'http://local/api/harness/', + threadId: 'thread-c', + headers: { authorization: 'Bearer t' }, + fetch: (input, init) => handler(new Request(input, init)), + }) + + const receipt = await client.prompt('hi') + expect(receipt.status).toBe('accepted') + + const controller = new AbortController() + const seen: Array = [] + for await (const entry of client.events({ signal: controller.signal })) { + seen.push(entry) + if ( + entry.event.type === 'CUSTOM' && + entry.event.name === 'harness.operation.finished' + ) { + controller.abort() + } + } + expect(JSON.stringify(seen)).toContain('hello from the host') + + const agent = await client.agents.pricer.start({ vendor: 'acme' }) + expect(agent.status).toBe('accepted') + // @ts-expect-error the input is typed from the harness + void client.agents.pricer.start({ vendor: 1 }) + + const snapshot = await client.snapshot() + expect(snapshot.threadId).toBe('thread-c') + await host.close() + }) + + it('throws on a refused request instead of retrying', async () => { + const { adapter } = mockAdapter([]) + const studio = defineHarness({ name: 'test/client-denied', adapter }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const handler = createHarnessHandler({ + host, + harness: studio, + authorize: () => null, + }) + const client = createHarnessClient({ + url: 'http://local/api/harness', + threadId: 't', + fetch: (input, init) => handler(new Request(input, init)), + }) + await expect(client.prompt('hi')).rejects.toThrow('401') + await expect(async () => { + for await (const _ of client.events()) break + }).rejects.toThrow('401') + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/harness-text.test.ts b/packages/ai-harness/tests/harness-text.test.ts new file mode 100644 index 0000000000..86bc7619ef --- /dev/null +++ b/packages/ai-harness/tests/harness-text.test.ts @@ -0,0 +1,49 @@ +import { describe, expect, it } from 'vitest' +import { EventType, chat } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, harnessText } from '../src' +import { messageTexts, mockAdapter, text } from './helpers' +import type { StreamChunk } from '@tanstack/ai' + +describe('harnessText', () => { + it('runs a harness as the model of chat() and keeps its own transcript', async () => { + const inner = mockAdapter([ + () => text('inner answer'), + () => text('second answer'), + ]) + const studio = defineHarness({ name: 'test/inner', adapter: inner.adapter }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + + const collect = async (content: string) => { + const chunks: Array = [] + for await (const chunk of chat({ + adapter: harnessText(studio, { host }), + messages: [{ role: 'user', content }], + threadId: 'outer-1', + })) { + chunks.push(chunk) + } + return chunks + } + + const first = await collect('hello') + const deltas = first + .filter((chunk) => chunk.type === EventType.TEXT_MESSAGE_CONTENT) + .map((chunk) => + chunk.type === EventType.TEXT_MESSAGE_CONTENT ? chunk.delta : '', + ) + expect(deltas.join('')).toBe('inner answer') + expect(first.some((chunk) => chunk.type === EventType.RUN_ERROR)).toBe( + false, + ) + + await collect('and then?') + // The inner session saw the first turn as history. + expect(messageTexts(inner.calls[1])).toEqual([ + 'hello', + 'inner answer', + 'and then?', + ]) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/protocol.test.ts b/packages/ai-harness/tests/protocol.test.ts new file mode 100644 index 0000000000..b37209e84c --- /dev/null +++ b/packages/ai-harness/tests/protocol.test.ts @@ -0,0 +1,287 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { defineAgent } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + createHarnessHandler, + createHarnessHost, + defineHarness, + handleHarnessSocket, + parseControlFrame, +} from '../src' +import { mockAdapter, text } from './helpers' +import type { WebSocketLike } from '@tanstack/ai' +import type { HostFrame } from '../src' + +const pricer = defineAgent({ + name: 'pricer', + description: 'Prices a vendor', + inputSchema: z.object({ vendor: z.string() }), + run: async (ctx) => ({ vendor: ctx.input.vendor, cents: 1200 }), +}) +const secret = defineAgent({ + name: 'secret', + description: 'Not for clients', + run: async () => 'hidden', +}) + +function setup(replies = [() => text('hello'), () => text('again')]) { + const { adapter, calls } = mockAdapter(replies) + const harness = defineHarness({ + name: 'test/protocol', + adapter, + agents: [pricer, secret], + expose: { agents: ['pricer'] }, + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const handler = createHarnessHandler({ + host, + harness, + authorize: (request) => + request.headers.get('authorization') === 'Bearer good' + ? { id: 'user-1' } + : null, + canAccess: (principal, threadId) => threadId.startsWith(principal.id), + }) + return { handler, host, harness, calls } +} + +const auth = { authorization: 'Bearer good' } + +async function readSse(response: Response, until: (data: any) => boolean) { + const reader = response.body!.getReader() + const decoder = new TextDecoder() + const frames: Array = [] + let buffer = '' + while (true) { + const { value, done } = await reader.read() + if (done) break + buffer += decoder.decode(value, { stream: true }) + const blocks = buffer.split('\n\n') + buffer = blocks.pop() ?? '' + for (const block of blocks) { + const data = block + .split('\n') + .find((line) => line.startsWith('data: ')) + ?.slice(6) + if (!data) continue + const parsed = JSON.parse(data) + frames.push(parsed) + if (until(parsed)) { + await reader.cancel() + return frames + } + } + } + return frames +} + +describe('HTTP handler', () => { + it('refuses requests that authorize rejects', async () => { + const { handler, host } = setup() + const response = await handler( + new Request('http://x/api/harness/capabilities'), + ) + expect(response.status).toBe(401) + await host.close() + }) + + it('refuses threads the principal cannot access', async () => { + const { handler, host } = setup() + const response = await handler( + new Request('http://x/api/harness/snapshot?threadId=other-thread', { + headers: auth, + }), + ) + expect(response.status).toBe(403) + await host.close() + }) + + it('lists only exposed agents in the capabilities', async () => { + const { handler, host } = setup() + const response = await handler( + new Request('http://x/api/harness/capabilities', { headers: auth }), + ) + const body = await response.json() + expect(body.identity.name).toBe('test/protocol') + expect( + body.custom.tanstack.agents.map((agent: { name: string }) => agent.name), + ).toEqual(['pricer']) + expect( + body.custom.tanstack.agents[0].inputSchema.properties.vendor.type, + ).toBe('string') + await host.close() + }) + + it('serves a standard AG-UI run as SSE', async () => { + const { handler, host } = setup() + const response = await handler( + new Request('http://x/api/harness/run', { + method: 'POST', + headers: { ...auth, 'content-type': 'application/json' }, + body: JSON.stringify({ + threadId: 'user-1-thread', + runId: 'client-run', + messages: [{ id: 'm1', role: 'user', content: 'hi' }], + tools: [], + context: [], + state: {}, + forwardedProps: {}, + }), + }), + ) + expect(response.headers.get('content-type')).toContain('text/event-stream') + const text = await response.text() + expect(text).toContain('"delta":"hello"') + expect(text).toContain('RUN_FINISHED') + await host.close() + }) + + it('takes control inputs, streams session events, and resumes from a cursor', async () => { + const { handler, host } = setup() + const control = (input: unknown) => + handler( + new Request('http://x/api/harness/control', { + method: 'POST', + headers: { ...auth, 'content-type': 'application/json' }, + body: JSON.stringify({ threadId: 'user-1-t', input }), + }), + ) + + const receipt = await ( + await control({ op: 'prompt', message: 'hi' }) + ).json() + expect(receipt.status).toBe('accepted') + + const first = await readSse( + await handler( + new Request('http://x/api/harness/events?threadId=user-1-t', { + headers: auth, + }), + ), + (frame) => frame.event?.name === 'harness.operation.finished', + ) + expect(first.every((frame) => frame.type === 'harness.event')).toBe(true) + const cursor = first.at(-1).cursor + + await (await control({ op: 'prompt', message: 'more' })).json() + const second = await readSse( + await handler( + new Request('http://x/api/harness/events?threadId=user-1-t', { + headers: { ...auth, 'Last-Event-ID': cursor }, + }), + ), + (frame) => frame.event?.name === 'harness.operation.finished', + ) + expect(second.every((frame) => Number(frame.cursor) > Number(cursor))).toBe( + true, + ) + expect(JSON.stringify(second)).toContain('again') + await host.close() + }) + + it('runs exposed agents from a client and refuses the others', async () => { + const { handler, host } = setup() + const control = async (input: unknown) => + ( + await handler( + new Request('http://x/api/harness/control', { + method: 'POST', + headers: { ...auth, 'content-type': 'application/json' }, + body: JSON.stringify({ threadId: 'user-1-a', input }), + }), + ) + ).json() + + expect(await control({ op: 'agent', agent: 'secret' })).toMatchObject({ + status: 'rejected', + reason: 'not_exposed', + }) + expect( + await control({ op: 'agent', agent: 'pricer', input: { vendor: 'x' } }), + ).toMatchObject({ + status: 'accepted', + }) + expect(await control({ op: 'launch' })).toMatchObject({ + error: expect.stringContaining('Invalid input'), + }) + await host.close() + }) +}) + +describe('WebSocket', () => { + function fakeSocket() { + const handlers: Record void>> = {} + const sent: Array = [] + const socket: WebSocketLike = { + send: (data) => sent.push(JSON.parse(data)), + close: () => handlers.close?.forEach((handler) => handler()), + addEventListener: (type: string, handler: (event?: any) => void) => { + ;(handlers[type] ??= []).push(handler) + }, + } + const receive = (frame: unknown) => + handlers.message?.forEach((handler) => + handler({ data: JSON.stringify(frame) }), + ) + return { socket, sent, receive } + } + + it('subscribes, applies inputs with receipts, and streams events', async () => { + const { host, harness } = setup() + const { socket, sent, receive } = fakeSocket() + handleHarnessSocket({ host, harness, socket, principal: { id: 'user-1' } }) + + receive({ + type: 'harness.input', + requestId: 'r0', + input: { op: 'prompt', message: 'x' }, + }) + await vi.waitFor(() => expect(sent.at(-1)?.type).toBe('harness.error')) + + receive({ type: 'harness.subscribe', threadId: 'user-1-ws' }) + await vi.waitFor(() => + expect(sent.some((frame) => frame.type === 'harness.hello')).toBe(true), + ) + receive({ + type: 'harness.input', + requestId: 'r1', + input: { op: 'prompt', message: 'hi' }, + }) + + await vi.waitFor(() => + expect( + sent.some( + (frame) => + frame.type === 'harness.event' && + frame.event.type === 'CUSTOM' && + frame.event.name === 'harness.operation.finished', + ), + ).toBe(true), + ) + expect( + sent.find((frame) => frame.type === 'harness.receipt'), + ).toMatchObject({ + requestId: 'r1', + status: 'accepted', + }) + receive({ type: 'harness.snapshot' }) + await vi.waitFor(() => + expect(sent.at(-1)).toMatchObject({ + type: 'harness.snapshot', + snapshot: { status: 'idle' }, + }), + ) + socket.close() + await host.close() + }) +}) + +describe('parseControlFrame', () => { + it('rejects unknown frames', () => { + expect(() => parseControlFrame('{"type":"nope"}')).toThrow('unknown type') + expect(() => + parseControlFrame('{"type":"harness.input","requestId":"1","input":{}}'), + ).toThrow('Invalid input') + }) +}) diff --git a/packages/ai-harness/tests/resume.test.ts b/packages/ai-harness/tests/resume.test.ts new file mode 100644 index 0000000000..c57c045baf --- /dev/null +++ b/packages/ai-harness/tests/resume.test.ts @@ -0,0 +1,185 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { EventType, toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { HARNESS_EVENTS, createHarnessHost, defineHarness } from '../src' +import { INTERRUPTED_TOOL_RESULT, LEASE } from '../src/resume' +import { gate, mockAdapter, text, toolCall } from './helpers' +import type { SessionEvent } from '../src' + +describe('checkpoints and leases', () => { + it('holds a lease and records the running tool with its replay mode', async () => { + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + const started = gate() + const release = gate() + const lookup = toolDefinition({ + name: 'lookup', + description: 'Look up a fact', + inputSchema: z.object({ q: z.string() }), + replay: 'safe', + }).server(async () => { + started.open() + await release.opened + return 'found' + }) + const { adapter } = mockAdapter([ + () => toolCall('lookup', { q: 'x' }), + () => text('done'), + ]) + const session = await host.open( + defineHarness({ name: 'test/lease', adapter, tools: [lookup] }), + { threadId: 't1' }, + ) + + const turn = session.prompt('go') + await started.opened + const running = await persistence.stores.runs.get(turn.id) + expect(running?.leaseOwner).toMatch(/^host-/) + expect(running?.leaseExpiresAt).toBeGreaterThan(Date.now()) + expect(running?.leaseExpiresAt).toBeLessThanOrEqual( + Date.now() + LEASE.ttlMs, + ) + expect(running?.checkpoint?.pendingTools).toEqual([ + { toolCallId: 'call-1', name: 'lookup', replay: 'safe' }, + ]) + // The assistant message with the tool call is saved before the tool runs. + const saved = await persistence.stores.messages.loadThread('t1') + expect(saved.at(-1)?.toolCalls?.[0]?.function.name).toBe('lookup') + + release.open() + await turn + const finished = await persistence.stores.runs.get(turn.id) + expect(finished?.checkpoint?.pendingTools).toEqual([]) + await host.close() + }) +}) + +describe('crash resume', () => { + it('continues a crashed turn, runs safe tools again, and notes the others', async () => { + const persistence = memoryPersistence() + await persistence.stores.messages.saveThread('t1', [ + { id: 'u1', role: 'user', content: 'refund and look up' }, + { + id: 'a1', + role: 'assistant', + content: '', + toolCalls: [ + { + id: 'call-charge', + type: 'function', + function: { name: 'charge', arguments: '{"cents":100}' }, + }, + { + id: 'call-lookup', + type: 'function', + function: { name: 'lookup', arguments: '{"q":"x"}' }, + }, + ], + }, + ]) + await persistence.stores.runs.createOrResume({ + runId: 'crashed-run', + threadId: 't1', + startedAt: Date.now() - 60_000, + }) + await persistence.stores.runs.update('crashed-run', { + leaseOwner: 'host-gone', + leaseExpiresAt: Date.now() - 1_000, + checkpoint: { + at: Date.now() - 2_000, + pendingTools: [ + { toolCallId: 'call-charge', name: 'charge', replay: 'never' }, + { toolCallId: 'call-lookup', name: 'lookup', replay: 'safe' }, + ], + }, + }) + + const charge = vi.fn(async () => 'charged') + const lookup = vi.fn(async () => 'found') + const tools = [ + toolDefinition({ + name: 'charge', + description: 'Charge a card', + inputSchema: z.object({ cents: z.number() }), + }).server(charge), + toolDefinition({ + name: 'lookup', + description: 'Look up a fact', + inputSchema: z.object({ q: z.string() }), + replay: 'safe', + }).server(lookup), + ] + const { adapter, calls } = mockAdapter([() => text('recovered')]) + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ name: 'test/resume', adapter, tools }), + { threadId: 't1' }, + ) + + const seen: Array = [] + const controller = new AbortController() + const reading = (async () => { + for await (const entry of session.events({ + from: '0', + signal: controller.signal, + })) { + seen.push(entry) + } + })() + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + await vi.waitFor(() => expect(calls).toHaveLength(1)) + controller.abort() + await reading + + expect(lookup).toHaveBeenCalledTimes(1) + expect(charge).not.toHaveBeenCalled() + const transcript = await persistence.stores.messages.loadThread('t1') + const chargeResult = transcript.find( + (message) => + message.role === 'tool' && message.toolCallId === 'call-charge', + ) + expect(chargeResult?.content).toBe(JSON.stringify(INTERRUPTED_TOOL_RESULT)) + expect(transcript.at(-1)?.content).toBe('recovered') + + expect((await persistence.stores.runs.get('crashed-run'))?.status).toBe( + 'failed', + ) + const resumed = seen.find( + (entry) => + entry.event.type === EventType.CUSTOM && + entry.event.name === HARNESS_EVENTS.operationResumed, + ) + expect(resumed?.event).toMatchObject({ + value: { resumedFrom: 'crashed-run' }, + }) + await host.close() + }) + + it('leaves a run alone while its lease is still valid', async () => { + const persistence = memoryPersistence() + await persistence.stores.runs.createOrResume({ + runId: 'live-run', + threadId: 't1', + startedAt: Date.now(), + }) + await persistence.stores.runs.update('live-run', { + leaseOwner: 'host-other', + leaseExpiresAt: Date.now() + 20_000, + }) + const { adapter, calls } = mockAdapter([]) + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ name: 'test/live', adapter }), + { + threadId: 't1', + }, + ) + expect(session.snapshot().status).toBe('idle') + expect(calls).toHaveLength(0) + expect((await persistence.stores.runs.get('live-run'))?.status).toBe( + 'running', + ) + await host.close() + }) +}) diff --git a/packages/ai-harness/vite.config.ts b/packages/ai-harness/vite.config.ts index 1f3542380f..2fabdfea04 100644 --- a/packages/ai-harness/vite.config.ts +++ b/packages/ai-harness/vite.config.ts @@ -28,7 +28,7 @@ const config = defineConfig({ export default mergeConfig( config, tanstackViteConfig({ - entry: ['./src/index.ts'], + entry: ['./src/index.ts', './src/client.ts'], srcDir: './src', cjs: false, }), diff --git a/packages/ai/src/activities/chat/tools/tool-definition.ts b/packages/ai/src/activities/chat/tools/tool-definition.ts index cbe4db76ca..a141e5d478 100644 --- a/packages/ai/src/activities/chat/tools/tool-definition.ts +++ b/packages/ai/src/activities/chat/tools/tool-definition.ts @@ -128,6 +128,7 @@ export interface ClientTool< needsApproval?: TNeedsApproval approvalSchema?: TApprovalSchema lazy?: boolean + replay?: 'safe' | 'never' metadata?: Record execute?: ToolExecuteFunction } @@ -220,6 +221,12 @@ export type ToolDefinitionConfig< inputSchema?: TInput outputSchema?: TOutput lazy?: boolean + /** + * Whether a harness may run this tool again after a crash cut the run + * between the call and its result. `'safe'` runs it again. `'never'` + * (default) gives the model a note that the tool may or may not have run. + */ + replay?: 'safe' | 'never' metadata?: Record } & ApprovalConfig diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 3d29433e71..33b6687fed 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -874,6 +874,13 @@ export interface Tool< /** If true, this tool is lazy and will only be sent to the LLM after being discovered via the lazy tool discovery mechanism. Works with both chat() (the synthetic discovery tool) and Code Mode (kept out of the system prompt and revealed via discover_tools). */ lazy?: boolean + /** + * Whether a harness may run this tool again after a crash cut the run + * between the call and its result. `'safe'` runs it again. `'never'` + * (default) gives the model a note that the tool may or may not have run. + */ + replay?: 'safe' | 'never' + /** Additional metadata for adapters or custom extensions */ metadata?: Record | undefined } diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index ca0c6edf73..4a914a15c6 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1979,12 +1979,18 @@ importers: packages/ai-acp: dependencies: '@agentclientprotocol/sdk': - specifier: ^0.25.0 - version: 0.25.1(zod@4.3.6) + specifier: ^1.5.0 + version: 1.5.0(zod@4.3.6) devDependencies: '@tanstack/ai': specifier: workspace:* version: link:../ai + '@tanstack/ai-harness': + specifier: workspace:* + version: link:../ai-harness + '@tanstack/ai-persistence': + specifier: workspace:* + version: link:../ai-persistence '@tanstack/ai-sandbox': specifier: workspace:* version: link:../ai-sandbox @@ -2535,6 +2541,37 @@ importers: specifier: ^4.2.0 version: 4.3.6 + packages/ai-harness-cli: + dependencies: + ink: + specifier: ^7.1.1 + version: 7.1.1(@types/react@19.3.0)(react-devtools-core@6.1.5)(react@19.2.3) + react: + specifier: ^19.2.3 + version: 19.2.3 + devDependencies: + '@tanstack/ai': + specifier: workspace:* + version: link:../ai + '@tanstack/ai-acp': + specifier: workspace:* + version: link:../ai-acp + '@tanstack/ai-harness': + specifier: workspace:* + version: link:../ai-harness + '@tanstack/ai-persistence': + specifier: workspace:* + version: link:../ai-persistence + '@tanstack/ai-sandbox': + specifier: workspace:* + version: link:../ai-sandbox + '@types/react': + specifier: ^19.2.7 + version: 19.3.0 + '@vitest/coverage-v8': + specifier: 4.1.10 + version: 4.1.10(vitest@4.1.11) + packages/ai-isolate-cloudflare: dependencies: '@tanstack/ai-code-mode': @@ -3900,11 +3937,15 @@ packages: zod: optional: true - '@agentclientprotocol/sdk@0.25.1': - resolution: {integrity: sha512-jx2rF3bdpGwZ75Q/meyEDLLbYmbtxk82Uh9hDCdxDvcEedBnNSF5hZAnL/kJR5VNz56JqwOmqnAqasC84MwwkQ==} + '@agentclientprotocol/sdk@1.5.0': + resolution: {integrity: sha512-524jwbB2iYWA+kWWyv9fhKbhU89dH/lu9u5EXwVNmfYzopV8BujCDxByBDZhxRUkB7RWIJzISCnwivdDk+bdVg==} peerDependencies: zod: ^3.25.0 || ^4.0.0 + '@alcalzone/ansi-tokenize@0.3.1': + resolution: {integrity: sha512-cUD0PT0YX+rImRhfn7BfMRduP/tX/Zqt/N0dHlAyn2RlcJssCb5KsCJkhln61eh9vuJDkH8wKn3L8zJzzWXPFg==} + engines: {node: '>=18'} + '@ampproject/remapping@2.3.0': resolution: {integrity: sha512-30iZtAPgz+LTIYoeivqYo853f02jBYSd5uGnGpkFV0M3xOt9aN73erkgYAmZU43x4VfqcnLxW9Kpg3R5LC4YYw==} engines: {node: '>=6.0.0'} @@ -11853,10 +11894,6 @@ packages: resolution: {integrity: sha512-quJQXlTSUGL2LH9SUXo8VwsY4soanhgo6LNSm84E1LBcE8s3O0wpdiRzyR9z/ZZJMlMWv37qOOb9pdJlMUEKFQ==} engines: {node: '>=8'} - ansi-regex@6.2.2: - resolution: {integrity: sha512-Bq3SmSpyFHaWjPk8If9yc6svM8c56dB5BAtW4Qbw5jHTwwXXcTLoRMkpDJp6VL0XzlWaCHTXrkFURMYmD0sLqg==} - engines: {node: '>=12'} - ansi-regex@6.3.0: resolution: {integrity: sha512-WpDfL7NO6j7tH88IDBNVdUJxDh9nmCteAVW9dsep846XdwF4naCBK+/tGLX3KJgcpgMRXCFlTM2hKGoK9FsdrQ==} engines: {node: '>=12'} @@ -11970,6 +12007,10 @@ packages: asynckit@0.4.0: resolution: {integrity: sha512-Oei9OH4tRh0YqU3GxhX79dM/mwVgvbZJaSNaRk+bshkj0S5cfHcgYakreBjrHwatXKbz+IoIdYLxrKim2MjW0Q==} + auto-bind@5.0.1: + resolution: {integrity: sha512-ooviqdwwgfIfNmDwo94wlshcdzfO64XV0Cg6oDsDYBJfITDz1EngD2z7DkbvCWn+XIMsIqW27sEVF6qcpJrRcg==} + engines: {node: ^12.20.0 || ^14.13.1 || >=16.0.0} + autoprefixer@10.4.22: resolution: {integrity: sha512-ARe0v/t9gO28Bznv6GgqARmVqcWOV3mfgUPn9becPHMiD3o9BwlRgaeccZnwTpZ7Zwqrm+c1sUSsMxIzQzc8Xg==} engines: {node: ^10 || ^12 || >=14} @@ -12442,6 +12483,10 @@ packages: resolution: {integrity: sha512-/lzGpEWL/8PfI0BmBOPRwp0c/wFNX1RdUML3jK/RcSBA9T8mZDdQpqYBKtCFTOfQbwPqWEOpjqW+Fnayc0969g==} engines: {node: '>=10'} + cli-boxes@4.0.1: + resolution: {integrity: sha512-5IOn+jcCEHEraYolBPs/sT4BxYCe2nHg374OPiItB1O96KZFseS2gthU4twyYzeDcFew4DaUM/xwc5BQf08JJw==} + engines: {node: '>=18.20 <19 || >=20.10'} + cli-cursor@2.1.0: resolution: {integrity: sha512-8lgKz8LmCRYZZQDpRyT2m5rKJ08TnU4tR9FFFW2rxpxR1FzWi4PQ/NfyODchAatHaUgnSPVcx/R5w6NuTBzFiw==} engines: {node: '>=4'} @@ -12450,6 +12495,10 @@ packages: resolution: {integrity: sha512-I/zHAwsKf9FqGoXM4WWRACob9+SNukZTd94DWF57E4toouRulbCxcUh6RKUEOQlYTHJnzkPMySvPNaaSLNfLZw==} engines: {node: '>=8'} + cli-cursor@4.0.0: + resolution: {integrity: sha512-VGtlMu3x/4DOtIUwEkRezxUZ2lBacNJCHash0N0WeZDBS+7Ux1dm3XWAgWYxLJFMMdOeXMHXorshEFhbMSGelg==} + engines: {node: ^12.20.0 || ^14.13.1 || >=16.0.0} + cli-cursor@5.0.0: resolution: {integrity: sha512-aCj4O5wKyszjMmDT4tZj93kxyydN/K5zPWSCe6/0AV/AA1pqe5ZBIw0a2ZfPQV7lL5/yb5HsUreJ6UFAF1tEQw==} engines: {node: '>=18'} @@ -12466,6 +12515,10 @@ packages: resolution: {integrity: sha512-xRwvIOMGrfOAnM1JYtqQImuaNtDEv9v6oIYAs4LIHwTiKee8uwvIi363igssOC0O5U04i4AlENs79LQLu9tEMw==} engines: {node: '>=20'} + cli-truncate@6.1.1: + resolution: {integrity: sha512-06p9vyLahLa4zkGcgsGxU6iEkSOiuI4fhCH6Emhe2lPAcoUv73n72DnODsnHA+5wwXGnV0n9M9/qOQJSjYhFhw==} + engines: {node: '>=22'} + cli-width@4.1.0: resolution: {integrity: sha512-ouuZd4/dm2Sw5Gmqy6bGyNNNe1qt9RpmxveLSO7KcgsTnU7RXfsw+/bukWGo1abgBiMAic068rclZsO4IWmmxQ==} engines: {node: '>= 12'} @@ -12501,6 +12554,10 @@ packages: code-block-writer@12.0.0: resolution: {integrity: sha512-q4dMFMlXtKR3XNBHyMHt/3pwYNA69EDk00lloMOaaUMKPUXBw6lpXtbu3MMVG6/uOihGnRDOlkyqsONEUj60+w==} + code-excerpt@4.0.0: + resolution: {integrity: sha512-xxodCmBen3iy2i0WtAK8FlFNrRzjUqjRsMfho58xT/wvZU1YTM3fCnRjcy1gJPMepaRlgm/0e6w8SpWHpn3/cA==} + engines: {node: ^12.20.0 || ^14.13.1 || >=16.0.0} + color-convert@1.9.3: resolution: {integrity: sha512-QfAUtd+vFdAtFQcC8CCyYt1fYWxSqAiK2cSD6zDB8N3cpsEBAvRxp9zOGg6G/SHHJYAT88/az/IuDGALsNVbGg==} @@ -12631,6 +12688,10 @@ packages: convert-source-map@2.0.0: resolution: {integrity: sha512-Kvp459HrV2FEJ1CAsi1Ku+MY3kasH19TFykTz2xWmMeq6bk2NU3XXvfJ+Q61m0xktWwt+1HSYf3JZsTms3aRJg==} + convert-to-spaces@2.0.1: + resolution: {integrity: sha512-rcQ1bsQO9799wq24uE5AM2tAILy4gXGIK/njFWcVQkGNZ96edlpY+A7bjwvzjYvLDyzmG1MmMLZhpcsb+klNMQ==} + engines: {node: ^12.20.0 || ^14.13.1 || >=16.0.0} + cookie-es@1.2.2: resolution: {integrity: sha512-+W7VmiVINB+ywl1HGXJXmrqkOhpKrIiVZV6tQuV54ZyQC7MMuBt81Vc336GMLoHBq5hV/F9eXgt5Mnx0Rha5Fg==} @@ -13355,6 +13416,9 @@ packages: resolution: {integrity: sha512-j6vWzfrGVfyXxge+O0x5sh6cvxAog0a/4Rdd2K36zCMV5eJ+/+tOAngRO8cODMNWbVRdVlmGZQL2YS3yR8bIUA==} engines: {node: '>= 0.4'} + es-toolkit@1.52.0: + resolution: {integrity: sha512-XTNEJQh1tY1ZJVcf6ayP/2n4ZPyaHlW2FWs7xvw5ddPuhUVjLD3olQVQS7kf58JbAB48iL0uL/jerTrjtV3lDA==} + esbuild-plugin-solid@0.5.0: resolution: {integrity: sha512-ITK6n+0ayGFeDVUZWNMxX+vLsasEN1ILrg4pISsNOQ+mq4ljlJJiuXotInd+HE0MzwTcA9wExT1yzDE2hsqPsg==} peerDependencies: @@ -13392,6 +13456,10 @@ packages: resolution: {integrity: sha512-vbRorB5FUQWvla16U8R/qgaFIya2qGzwDrNmCZuYKrbdSUMG6I1ZCGQRefkRVhuOkIGVne7BQ35DSfo1qvJqFg==} engines: {node: '>=0.8.0'} + escape-string-regexp@2.0.0: + resolution: {integrity: sha512-UpzcLCXolUWcNu5HtVMHYdXJjArjsF9C0aNnquZYY4uW/Vu0miy5YoWvbV345HauVvcAUnpRuhMMcqTcGOY2+w==} + engines: {node: '>=8'} + escape-string-regexp@4.0.0: resolution: {integrity: sha512-TtpcNJ3XAzx3Gq8sWRzJaVajRs0uVxA2YAkdb1jm2YkPz4G6egUFAyA3n5vtEIZefPk5Wa4UXbKuS5fKkJWdgA==} engines: {node: '>=10'} @@ -13939,10 +14007,6 @@ packages: resolution: {integrity: sha512-DyFP3BM/3YHTQOCUL/w0OZHR0lpKeGrxotcHWcqNEdnltqFwXVfhEBQ94eIo34AfQpo0rGki4cyIiftY06h2Fg==} engines: {node: 6.* || 8.* || >= 10.*} - get-east-asian-width@1.6.0: - resolution: {integrity: sha512-QRbvDIbx6YklUe6RxeTeleMR0yv3cYH6PsPZHcnVn7xv7zO1BHN8r0XETu8n6Ye3Q+ahtSarc3WgtNWmehIBfA==} - engines: {node: '>=18'} - get-east-asian-width@1.7.0: resolution: {integrity: sha512-XjH1AECxf0giL2V1aU8vKyRR2ppRUb5c0EvT7zuJTokQ74bNo52zOtghqdWIqrhUD79fo3x0WfKZdOqxF6LG1Q==} engines: {node: '>=18'} @@ -14346,6 +14410,10 @@ packages: resolution: {integrity: sha512-JmXMZ6wuvDmLiHEml9ykzqO6lwFbof0GG4IkcGaENdCRDDmMVnny7s5HsIgHCbaq0w2MyPhDqkhTUgS2LU2PHA==} engines: {node: '>=0.8.19'} + indent-string@5.0.0: + resolution: {integrity: sha512-m6FAo/spmsW2Ab2fU35JTYwtOKa2yAwXSwgjSv1TJzh4Mh7mC3lzAOVLBprb72XsTrgkEIsl7YrFNAiDiRhIGg==} + engines: {node: '>=12'} + inherits@2.0.4: resolution: {integrity: sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==} @@ -14355,6 +14423,19 @@ packages: injection-js@2.6.1: resolution: {integrity: sha512-dbR5bdhi7TWDoCye9cByZqeg/gAfamm8Vu3G1KZOTYkOif8WkuM8CD0oeDPtZYMzT5YH76JAFB7bkmyY9OJi2A==} + ink@7.1.1: + resolution: {integrity: sha512-Y43xxa1ZSPvpmfLHcN5o+OdP8Rf8ykkNJEuKYOUNZKT8wXVNLFTtEm1nSDMQkfBH+YANF4Xuu0hhZ4ejqAtN2w==} + engines: {node: '>=22'} + peerDependencies: + '@types/react': '>=19.2.0' + react: '>=19.2.0' + react-devtools-core: '>=6.1.2' + peerDependenciesMeta: + '@types/react': + optional: true + react-devtools-core: + optional: true + inline-style-parser@0.1.1: resolution: {integrity: sha512-7NXolsK4CAS5+xvdj5OMMbI962hU/wvwoxk+LWR9Ek9bVtyuuYScDN6eS0rUm6TxApFpw7CX1o4uJzcd4AyD3Q==} @@ -14479,6 +14560,11 @@ packages: resolution: {integrity: sha512-S+OpgB5i7wzIue/YSE5hg0e5ZYfG3hhpNh9KGl6ayJ38p7ED6wxQLd1TV91xHpcTvw90KMJ9EwN3F/iNflHBVg==} engines: {node: '>=8'} + is-in-ci@2.0.0: + resolution: {integrity: sha512-cFeerHriAnhrQSbpAxL37W1wcJKUUX07HyLWZCW1URJT/ra3GyUTzBgUnh24TMVfNTV2Hij2HLxkPHFZfOZy5w==} + engines: {node: '>=20'} + hasBin: true + is-in-ssh@1.0.0: resolution: {integrity: sha512-jYa6Q9rH90kR1vKB6NM7qqd1mge3Fx4Dhw5TVlK1MUBqhEOuCagrEHMevNuCcbECmXZ0ThXkRm+Ymr51HwEPAw==} engines: {node: '>=20'} @@ -16182,6 +16268,10 @@ packages: partial-json@0.1.7: resolution: {integrity: sha512-Njv/59hHaokb/hRUjce3Hdv12wd60MtM9Z5Olmn+nehe0QDAsRtRbJPvJ0Z91TusF0SuZRIvnM+S4l6EIP8leA==} + patch-console@2.0.0: + resolution: {integrity: sha512-0YNdUceMdaQwoKce1gatDScmMo5pu/tfABfnzEqeG0gtTmd7mh/WcwgUjtAeOU7N8nFFlbQBnFK2gXW5fGvmMA==} + engines: {node: ^12.20.0 || ^14.13.1 || >=16.0.0} + path-browserify@1.0.1: resolution: {integrity: sha512-b7uo2UCUOYZcnF/3ID0lulOJi/bafxa1xPe7ZPsammBSpjSWQkjNxlt635YGS2MiR9GjvuXCtz2emr3jbsz98g==} @@ -16618,6 +16708,12 @@ packages: '@types/react': optional: true + react-reconciler@0.33.0: + resolution: {integrity: sha512-KetWRytFv1epdpJc3J4G75I4WrplZE5jOL7Yq0p34+OVOKF4Se7WrdIdVC45XsSSmUTlht2FM/fM1FZb1mfQeA==} + engines: {node: '>=0.10.0'} + peerDependencies: + react: ^19.2.0 + react-refresh@0.14.2: resolution: {integrity: sha512-jCvmsr+1IUSMUyzOkRcvnVbX3ZYC6g9TDrDbFuFmRDq7PD4yaGbLKNQL6k2jnArV8hjYxh7hVhAZB6s9HDGpZA==} engines: {node: '>=0.10.0'} @@ -16858,6 +16954,10 @@ packages: resolution: {integrity: sha512-l+sSefzHpj5qimhFSE5a8nufZYAM3sBSVMAPtYkmC+4EH2anSGaEMXSD0izRQbu9nfyQ9y5JrVmp7E8oZrUjvA==} engines: {node: '>=8'} + restore-cursor@4.0.0: + resolution: {integrity: sha512-I9fPXU9geO9bHOt9pHHOhOkYerIMsmVaWB0rA2AI9ERh/+x/i7MV5HKBNrg+ljO5eoPVgCcnFuRjJ9uH6I/3eg==} + engines: {node: ^12.20.0 || ^14.13.1 || >=16.0.0} + restore-cursor@5.1.0: resolution: {integrity: sha512-oMA2dcrw6u0YfxJQXm342bFKX/E4sG9rbTzO9ptUcR/e8A33cHuvStiYOwH7fszkZlZ1z/ta9AAoPk2F4qIOHA==} engines: {node: '>=18'} @@ -17278,6 +17378,10 @@ packages: resolution: {integrity: sha512-stxByr12oeeOyY2BlviTNQlYV5xOj47GirPr4yA1hE9JCtxfQN0+tVbkxwCtYDQWhEKWFHsEK48ORg5jrouCAg==} engines: {node: '>=20'} + slice-ansi@9.0.1: + resolution: {integrity: sha512-aBY19bn/XA+hKOKX0qL7Z0UfdBsbQvc9hn993U8ALUjzxCvDcuZqLoRXjGJrUARWAlwMnsRVD9sw2AFNOYvalA==} + engines: {node: '>=22'} + slugify@1.6.9: resolution: {integrity: sha512-vZ7rfeehZui7wQs438JXBckYLkIIdfHOXsaVEUMyS5fHo1483l1bMdo0EDSWYclY0yZKFOipDy4KHuKs6ssvdg==} engines: {node: '>=8.0.0'} @@ -17375,6 +17479,10 @@ packages: resolution: {integrity: sha512-wPldCk3asibAjQ/kziWQQt1Wh3PgDFpC0XpwclzKcdT1vql6KeYxf5LIt4nlFkUeR8WuphYMKqUA56X4rjbfgQ==} engines: {node: '>=10.16.0'} + stack-utils@2.0.6: + resolution: {integrity: sha512-XlkWvfIm6RmsWtNJx+uqtKLS8eqFbxUg0ZzLXqY0caEy9l7hruX8IpiDnjsLavoBgqCCR71TqWO8MaXYheJ3RQ==} + engines: {node: '>=10'} + stackback@0.0.2: resolution: {integrity: sha512-1XMJE5fQo1jGH6Y/7ebnwPOBEkIEnT4QF32d5R1+VXdXveM0IBMJt8zfaxX1P3QhVwrYe+576+jkANtSS2mBbw==} @@ -17450,10 +17558,6 @@ packages: resolution: {integrity: sha512-tsaTIkKW9b4N+AEj+SVA+WhJzV7/zMhcSu78mLKWSk7cXMOSHsBKFWUs0fWwq8QyK3MgJBQRX6Gbi4kYbdvGkQ==} engines: {node: '>=18'} - string-width@8.2.1: - resolution: {integrity: sha512-IIaP0g3iy9Cyy18w3M9YcaDudujEAVHKt3a3QJg1+sr/oX96TbaGUubG0hJyCjCBThFH+tFpcIyoUHUn1ogaLA==} - engines: {node: '>=20'} - string-width@8.2.2: resolution: {integrity: sha512-GaPUh5gfdrYzqeVNZvUfT23vYYxXzKYidUcnMtJg/3rxRV63EFZy3k6xfKlmfeJD0176lnUV/Usr3XcwSvFzpg==} engines: {node: '>=20'} @@ -17475,10 +17579,6 @@ packages: resolution: {integrity: sha512-Y38VPSHcqkFrCpFnQ9vuSXmquuv5oXOKpGeT6aGrr3o3Gc9AlVa6JBfUSOCnbxGGZF+/0ooI7KrPuUSztUdU5A==} engines: {node: '>=8'} - strip-ansi@7.1.2: - resolution: {integrity: sha512-gmBGslpoQJtgnMAvOVqGZpEz9dyoKTCzy2nfz/n8aIFhN/jCE/rCmcxabB6jOOHV+0WNnylOxaxBQPSvcWklhA==} - engines: {node: '>=12'} - strip-ansi@7.2.0: resolution: {integrity: sha512-yDPMNjp4WyfYBkHnjIRLfca1i6KMyGCtsVgoKe/z1+6vukgaENdgGBZt+ZmKPc4gavvEZ5OgHfHdrazhgNyG7w==} engines: {node: '>=12'} @@ -17625,6 +17725,10 @@ packages: resolution: {integrity: sha512-un0FmiRUQNr5PJqy9kP7c40F5BOfpGlYTrxonDChEZB7pzZxRNp/bt+ymiy9/npwXya9KH99nJ/GXFIiUkYGFQ==} engines: {node: '>=8'} + terminal-size@4.0.1: + resolution: {integrity: sha512-avMLDQpUI9I5XFrklECw1ZEUPJhqzcwSWsyyI8blhRLT+8N1jLJWLWWYQpB2q2xthq8xDvjZPISVh53T/+CLYQ==} + engines: {node: '>=18'} + terser@5.44.1: resolution: {integrity: sha512-t/R3R/n0MSwnnazuPpPNVO60LX0SKL45pyl9YlvxIdkH0Of7D5qM2EVe+yASRIlY5pZ73nclYJfNANGWPwFDZw==} engines: {node: '>=10'} @@ -17866,6 +17970,10 @@ packages: resolution: {integrity: sha512-RAH822pAdBgcNMAfWnCBU3CFZcfZ/i1eZjwFU/dsLKumyuuP3niueg2UAukXYF0E2AAoc82ZSSf9J0WQBinzHA==} engines: {node: '>=12.20'} + type-fest@5.10.0: + resolution: {integrity: sha512-NoSdpq/WEiAg5sjmBkmV/hfxv6HJH4NqPNrqjtSO5CwRmpsDfaf4begxW34KdJykH/l1yHtwBWQkCRdoXO8mPA==} + engines: {node: '>=20'} + type-fest@5.3.1: resolution: {integrity: sha512-VCn+LMHbd4t6sF3wfU/+HKT63C9OoyrSIf4b+vtWHpt2U7/4InZG467YDNMFMR70DdHjAdpPWmw2lzRdg0Xqqg==} engines: {node: '>=20'} @@ -18786,6 +18894,10 @@ packages: resolution: {integrity: sha512-o0cyEG0e8GPzT4iGHphIOh0cJOV8fivsXxddQasHPHfoZf1ZexrfeA21w2NaEN1RHE+fXlfISmOE8R9N3u3Qig==} engines: {node: '>=12'} + widest-line@6.0.0: + resolution: {integrity: sha512-U89AsyEeAsyoF0zVJBkG9zBgekjgjK7yk9sje3F4IQpXBJ10TF6ByLlIfjMhcmHMJgHZI4KHt4rdNfktzxIAMA==} + engines: {node: '>=20'} + word-wrap@1.2.5: resolution: {integrity: sha512-BN22B5eaMMI9UMtjrGd5g5eCYPpCPDUy0FJXbYsaT5zYxjFOckS53SQDE3pWkVoWpHXVb3BrYcEN4Twa55B5cA==} engines: {node: '>=0.10.0'} @@ -18810,6 +18922,10 @@ packages: '@cloudflare/workers-types': optional: true + wrap-ansi@10.0.2: + resolution: {integrity: sha512-6OLRNZqRntVGBm26ghZ5eZhCuWWyOkRCSn8XsOIuq5stpCE/tjLbdozd93naILZig0lmD8ANfRwnhn0DLeXGlg==} + engines: {node: '>=20'} + wrap-ansi@6.2.0: resolution: {integrity: sha512-r6lPcBGxZXlIcymEu7InxDMhdW0KDxpLgoFLcguasxCaJ/SOIZwINatK9KY/tf+ZrlywOKU0UDj3ATXUBfxJXA==} engines: {node: '>=8'} @@ -18986,6 +19102,9 @@ packages: resolution: {integrity: sha512-CzhO+pFNo8ajLM2d2IW/R93ipy99LWjtwblvC1RsoSUMZgyLbYFr221TnSNT7GjGdYui6P459mw9JH/g/zW2ug==} engines: {node: '>=18'} + yoga-layout@3.2.1: + resolution: {integrity: sha512-0LPOt3AxKqMdFBZA3HBAt/t/8vIKq7VaQYbuA8WxCgung+p9TVyKRYdpvCb80HcdTN2NkbIKbhNwKUfm3tQywQ==} + youch-core@0.3.3: resolution: {integrity: sha512-ho7XuGjLaJ2hWHoK8yFnsUGy2Y5uDpqSTq1FkHLK4/oqKtyUU1AFbOOxY4IpC9f0fTLjwYbslUz0Po5BpD1wrA==} @@ -19039,10 +19158,15 @@ snapshots: optionalDependencies: zod: 4.2.1 - '@agentclientprotocol/sdk@0.25.1(zod@4.3.6)': + '@agentclientprotocol/sdk@1.5.0(zod@4.3.6)': dependencies: zod: 4.3.6 + '@alcalzone/ansi-tokenize@0.3.1': + dependencies: + ansi-styles: 6.2.3 + is-fullwidth-code-point: 5.1.0 + '@ampproject/remapping@2.3.0': dependencies: '@jridgewell/gen-mapping': 0.3.13 @@ -22599,7 +22723,7 @@ snapshots: dependencies: string-width: 5.1.2 string-width-cjs: string-width@4.2.3 - strip-ansi: 7.1.2 + strip-ansi: 7.2.0 strip-ansi-cjs: strip-ansi@6.0.1 wrap-ansi: 8.1.0 wrap-ansi-cjs: wrap-ansi@7.0.0 @@ -27077,7 +27201,7 @@ snapshots: '@tanstack/devtools-event-bus@0.3.3': dependencies: - ws: 8.21.0 + ws: 8.21.3 transitivePeerDependencies: - bufferutil - utf-8-validate @@ -29666,8 +29790,6 @@ snapshots: ansi-regex@5.0.1: {} - ansi-regex@6.2.2: {} - ansi-regex@6.3.0: {} ansi-styles@3.2.1: @@ -29790,6 +29912,8 @@ snapshots: asynckit@0.4.0: {} + auto-bind@5.0.1: {} + autoprefixer@10.4.22(postcss@8.5.26): dependencies: browserslist: 4.28.1 @@ -30455,6 +30579,8 @@ snapshots: cli-boxes@3.0.0: {} + cli-boxes@4.0.1: {} + cli-cursor@2.1.0: dependencies: restore-cursor: 2.0.0 @@ -30463,6 +30589,10 @@ snapshots: dependencies: restore-cursor: 3.1.0 + cli-cursor@4.0.0: + dependencies: + restore-cursor: 4.0.0 + cli-cursor@5.0.0: dependencies: restore-cursor: 5.1.0 @@ -30474,7 +30604,12 @@ snapshots: cli-truncate@5.2.0: dependencies: slice-ansi: 8.0.0 - string-width: 8.2.1 + string-width: 8.2.2 + + cli-truncate@6.1.1: + dependencies: + slice-ansi: 9.0.1 + string-width: 8.2.2 cli-width@4.1.0: {} @@ -30510,6 +30645,10 @@ snapshots: code-block-writer@12.0.0: {} + code-excerpt@4.0.0: + dependencies: + convert-to-spaces: 2.0.1 + color-convert@1.9.3: dependencies: color-name: 1.1.3 @@ -30625,6 +30764,8 @@ snapshots: convert-source-map@2.0.0: {} + convert-to-spaces@2.0.1: {} + cookie-es@1.2.2: {} cookie-es@1.2.3: {} @@ -31259,6 +31400,8 @@ snapshots: has-tostringtag: 1.0.2 hasown: 2.0.4 + es-toolkit@1.52.0: {} + esbuild-plugin-solid@0.5.0(esbuild@0.28.1)(solid-js@1.9.10): dependencies: '@babel/core': 7.29.7 @@ -31388,6 +31531,8 @@ snapshots: escape-string-regexp@1.0.5: {} + escape-string-regexp@2.0.0: {} + escape-string-regexp@4.0.0: {} escape-string-regexp@5.0.0: {} @@ -32087,8 +32232,6 @@ snapshots: get-caller-file@2.0.5: {} - get-east-asian-width@1.6.0: {} - get-east-asian-width@1.7.0: {} get-intrinsic@1.3.0: @@ -32599,6 +32742,8 @@ snapshots: imurmurhash@0.1.4: {} + indent-string@5.0.0: {} + inherits@2.0.4: {} ini@1.3.8: {} @@ -32607,6 +32752,41 @@ snapshots: dependencies: tslib: 2.8.1 + ink@7.1.1(@types/react@19.3.0)(react-devtools-core@6.1.5)(react@19.2.3): + dependencies: + '@alcalzone/ansi-tokenize': 0.3.1 + ansi-escapes: 7.3.0 + ansi-styles: 6.2.3 + auto-bind: 5.0.1 + chalk: 5.6.2 + cli-boxes: 4.0.1 + cli-cursor: 4.0.0 + cli-truncate: 6.1.1 + code-excerpt: 4.0.0 + es-toolkit: 1.52.0 + indent-string: 5.0.0 + is-in-ci: 2.0.0 + patch-console: 2.0.0 + react: 19.2.3 + react-reconciler: 0.33.0(react@19.2.3) + scheduler: 0.27.0 + signal-exit: 3.0.7 + slice-ansi: 9.0.1 + stack-utils: 2.0.6 + string-width: 8.2.2 + terminal-size: 4.0.1 + type-fest: 5.10.0 + widest-line: 6.0.0 + wrap-ansi: 10.0.2 + ws: 8.21.3 + yoga-layout: 3.2.1 + optionalDependencies: + '@types/react': 19.3.0 + react-devtools-core: 6.1.5 + transitivePeerDependencies: + - bufferutil + - utf-8-validate + inline-style-parser@0.1.1: {} inline-style-parser@0.2.7: {} @@ -32715,7 +32895,7 @@ snapshots: is-fullwidth-code-point@5.1.0: dependencies: - get-east-asian-width: 1.6.0 + get-east-asian-width: 1.7.0 is-glob@4.0.3: dependencies: @@ -32727,6 +32907,8 @@ snapshots: dependencies: html-tags: 3.3.1 + is-in-ci@2.0.0: {} + is-in-ssh@1.0.0: {} is-inside-container@1.0.0: @@ -33401,7 +33583,7 @@ snapshots: ansi-escapes: 7.3.0 cli-cursor: 5.0.0 slice-ansi: 7.1.2 - strip-ansi: 7.1.2 + strip-ansi: 7.2.0 wrap-ansi: 9.0.2 loglevel@1.9.2: {} @@ -35389,6 +35571,8 @@ snapshots: partial-json@0.1.7: {} + patch-console@2.0.0: {} + path-browserify@1.0.1: {} path-exists@3.0.0: {} @@ -36123,6 +36307,11 @@ snapshots: - supports-color - utf-8-validate + react-reconciler@0.33.0(react@19.2.3): + dependencies: + react: 19.2.3 + scheduler: 0.27.0 + react-refresh@0.14.2: {} react-refresh@0.17.0: {} @@ -36478,6 +36667,11 @@ snapshots: onetime: 5.1.2 signal-exit: 3.0.7 + restore-cursor@4.0.0: + dependencies: + onetime: 5.1.2 + signal-exit: 3.0.7 + restore-cursor@5.1.0: dependencies: onetime: 7.0.0 @@ -37124,6 +37318,11 @@ snapshots: ansi-styles: 6.2.3 is-fullwidth-code-point: 5.1.0 + slice-ansi@9.0.1: + dependencies: + ansi-styles: 6.2.3 + is-fullwidth-code-point: 5.1.0 + slugify@1.6.9: {} smart-buffer@4.2.0: {} @@ -37220,6 +37419,10 @@ snapshots: cpu-features: 0.0.10 nan: 2.27.0 + stack-utils@2.0.6: + dependencies: + escape-string-regexp: 2.0.0 + stackback@0.0.2: {} stackframe@1.3.4: {} @@ -37285,18 +37488,13 @@ snapshots: dependencies: eastasianwidth: 0.2.0 emoji-regex: 9.2.2 - strip-ansi: 7.1.2 + strip-ansi: 7.2.0 string-width@7.2.0: dependencies: emoji-regex: 10.6.0 - get-east-asian-width: 1.6.0 - strip-ansi: 7.1.2 - - string-width@8.2.1: - dependencies: - get-east-asian-width: 1.6.0 - strip-ansi: 7.1.2 + get-east-asian-width: 1.7.0 + strip-ansi: 7.2.0 string-width@8.2.2: dependencies: @@ -37324,10 +37522,6 @@ snapshots: dependencies: ansi-regex: 5.0.1 - strip-ansi@7.1.2: - dependencies: - ansi-regex: 6.2.2 - strip-ansi@7.2.0: dependencies: ansi-regex: 6.3.0 @@ -37524,6 +37718,8 @@ snapshots: ansi-escapes: 4.3.2 supports-hyperlinks: 2.3.0 + terminal-size@4.0.1: {} + terser@5.44.1: dependencies: '@jridgewell/source-map': 0.3.11 @@ -37799,6 +37995,10 @@ snapshots: type-fest@2.19.0: {} + type-fest@5.10.0: + dependencies: + tagged-tag: 1.0.0 + type-fest@5.3.1: dependencies: tagged-tag: 1.0.0 @@ -38851,6 +39051,10 @@ snapshots: dependencies: string-width: 5.1.2 + widest-line@6.0.0: + dependencies: + string-width: 8.2.2 + word-wrap@1.2.5: {} workerd@1.20260609.1: @@ -38886,6 +39090,11 @@ snapshots: - bufferutil - utf-8-validate + wrap-ansi@10.0.2: + dependencies: + ansi-styles: 6.2.3 + string-width: 8.2.2 + wrap-ansi@6.2.0: dependencies: ansi-styles: 4.3.0 @@ -38902,13 +39111,13 @@ snapshots: dependencies: ansi-styles: 6.2.3 string-width: 5.1.2 - strip-ansi: 7.1.2 + strip-ansi: 7.2.0 wrap-ansi@9.0.2: dependencies: ansi-styles: 6.2.3 string-width: 7.2.0 - strip-ansi: 7.1.2 + strip-ansi: 7.2.0 wrappy@1.0.2: {} @@ -39016,6 +39225,8 @@ snapshots: yoctocolors@2.1.2: {} + yoga-layout@3.2.1: {} + youch-core@0.3.3: dependencies: '@poppinss/exception': 1.2.3 diff --git a/testing/e2e/fixtures/harness/basic.json b/testing/e2e/fixtures/harness/basic.json index 252c90de69..edf1a61780 100644 --- a/testing/e2e/fixtures/harness/basic.json +++ b/testing/e2e/fixtures/harness/basic.json @@ -1,20 +1,44 @@ { "fixtures": [ { - "match": { "userMessage": "[harness-turns] first" }, - "response": { "content": "First answer." } + "match": { + "userMessage": "[harness-turns] first" + }, + "response": { + "content": "First answer." + } }, { - "match": { "userMessage": "[harness-turns] second" }, - "response": { "content": "Second answer." } + "match": { + "userMessage": "[harness-turns] second" + }, + "response": { + "content": "Second answer." + } }, { - "match": { "userMessage": "[harness-agent] price vendor a" }, - "response": { "content": "Vendor A costs 12 dollars." } + "match": { + "userMessage": "[harness-agent] price vendor a" + }, + "response": { + "content": "Vendor A costs 12 dollars." + } }, { - "match": { "userMessage": "[harness-agent] what did it cost?" }, - "response": { "content": "It cost 12 dollars." } + "match": { + "userMessage": "[harness-agent] what did it cost?" + }, + "response": { + "content": "It cost 12 dollars." + } + }, + { + "match": { + "userMessage": "[harness-protocol] hello" + }, + "response": { + "content": "Hello over the harness protocol." + } } ] } diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 87b7a9edc4..5c3c381028 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -126,6 +126,7 @@ import { Route as ApiVideoStreamRouteImport } from './routes/api.video.stream' import { Route as ApiTtsStreamRouteImport } from './routes/api.tts.stream' import { Route as ApiTranscriptionStreamRouteImport } from './routes/api.transcription.stream' import { Route as ApiImageStreamRouteImport } from './routes/api.image.stream' +import { Route as ApiHarnessProtocolSplatRouteImport } from './routes/api.harness-protocol.$' import { Route as ApiAudioStreamRouteImport } from './routes/api.audio.stream' const WebsocketAdapterRoute = WebsocketAdapterRouteImport.update({ @@ -739,6 +740,11 @@ const ApiImageStreamRoute = ApiImageStreamRouteImport.update({ path: '/stream', getParentRoute: () => ApiImageRoute, } as any) +const ApiHarnessProtocolSplatRoute = ApiHarnessProtocolSplatRouteImport.update({ + id: '/api/harness-protocol/$', + path: '/api/harness-protocol/$', + getParentRoute: () => rootRouteImport, +} as any) const ApiAudioStreamRoute = ApiAudioStreamRouteImport.update({ id: '/stream', path: '/stream', @@ -860,6 +866,7 @@ export interface FileRoutesByFullPath { '/api/world': typeof ApiWorldRoute '/$provider/': typeof ProviderIndexRoute '/api/audio/stream': typeof ApiAudioStreamRoute + '/api/harness-protocol/$': typeof ApiHarnessProtocolSplatRoute '/api/image/stream': typeof ApiImageStreamRoute '/api/transcription/stream': typeof ApiTranscriptionStreamRoute '/api/tts/stream': typeof ApiTtsStreamRoute @@ -980,6 +987,7 @@ export interface FileRoutesByTo { '/api/world': typeof ApiWorldRoute '/$provider': typeof ProviderIndexRoute '/api/audio/stream': typeof ApiAudioStreamRoute + '/api/harness-protocol/$': typeof ApiHarnessProtocolSplatRoute '/api/image/stream': typeof ApiImageStreamRoute '/api/transcription/stream': typeof ApiTranscriptionStreamRoute '/api/tts/stream': typeof ApiTtsStreamRoute @@ -1101,6 +1109,7 @@ export interface FileRoutesById { '/api/world': typeof ApiWorldRoute '/$provider/': typeof ProviderIndexRoute '/api/audio/stream': typeof ApiAudioStreamRoute + '/api/harness-protocol/$': typeof ApiHarnessProtocolSplatRoute '/api/image/stream': typeof ApiImageStreamRoute '/api/transcription/stream': typeof ApiTranscriptionStreamRoute '/api/tts/stream': typeof ApiTtsStreamRoute @@ -1223,6 +1232,7 @@ export interface FileRouteTypes { | '/api/world' | '/$provider/' | '/api/audio/stream' + | '/api/harness-protocol/$' | '/api/image/stream' | '/api/transcription/stream' | '/api/tts/stream' @@ -1343,6 +1353,7 @@ export interface FileRouteTypes { | '/api/world' | '/$provider' | '/api/audio/stream' + | '/api/harness-protocol/$' | '/api/image/stream' | '/api/transcription/stream' | '/api/tts/stream' @@ -1463,6 +1474,7 @@ export interface FileRouteTypes { | '/api/world' | '/$provider/' | '/api/audio/stream' + | '/api/harness-protocol/$' | '/api/image/stream' | '/api/transcription/stream' | '/api/tts/stream' @@ -1583,6 +1595,7 @@ export interface RootRouteChildren { ApiWebMcpPageToolsRoute: typeof ApiWebMcpPageToolsRoute ApiWorldRoute: typeof ApiWorldRoute ProviderIndexRoute: typeof ProviderIndexRoute + ApiHarnessProtocolSplatRoute: typeof ApiHarnessProtocolSplatRoute } declare module '@tanstack/react-router' { @@ -2406,6 +2419,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiImageStreamRouteImport parentRoute: typeof ApiImageRoute } + '/api/harness-protocol/$': { + id: '/api/harness-protocol/$' + path: '/api/harness-protocol/$' + fullPath: '/api/harness-protocol/$' + preLoaderRoute: typeof ApiHarnessProtocolSplatRouteImport + parentRoute: typeof rootRouteImport + } '/api/audio/stream': { id: '/api/audio/stream' path: '/stream' @@ -2589,6 +2609,7 @@ const rootRouteChildren: RootRouteChildren = { ApiWebMcpPageToolsRoute: ApiWebMcpPageToolsRoute, ApiWorldRoute: ApiWorldRoute, ProviderIndexRoute: ProviderIndexRoute, + ApiHarnessProtocolSplatRoute: ApiHarnessProtocolSplatRoute, } export const routeTree = rootRouteImport ._addFileChildren(rootRouteChildren) diff --git a/testing/e2e/src/routes/api.harness-protocol.$.ts b/testing/e2e/src/routes/api.harness-protocol.$.ts new file mode 100644 index 0000000000..50eddad947 --- /dev/null +++ b/testing/e2e/src/routes/api.harness-protocol.$.ts @@ -0,0 +1,57 @@ +import { createFileRoute } from '@tanstack/react-router' +import { defineAgent } from '@tanstack/ai' +import { + createHarnessHandler, + createHarnessHost, + defineHarness, +} from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { z } from 'zod' +import { createTextAdapter } from '@/lib/providers' + +/** + * The harness session protocol behind `createHarnessHandler`. The main + * model is the OpenAI adapter against aimock. The aimock port and test id + * come from headers, so one handler serves every test. + */ +const hosts = new Map>() + +function handlerFor(request: Request) { + const testId = request.headers.get('x-test-id') ?? 'default' + const port = Number(request.headers.get('x-aimock-port') ?? '4010') + let host = hosts.get(testId) + if (!host) { + host = createHarnessHost({ persistence: memoryPersistence() }) + hosts.set(testId, host) + } + const harness = defineHarness({ + name: 'e2e/protocol', + adapter: createTextAdapter('openai', undefined, port, testId).adapter, + agents: [ + defineAgent({ + name: 'echo', + description: 'Echoes its input', + inputSchema: z.object({ text: z.string() }), + run: async (ctx) => ctx.input.text, + }), + ], + expose: { agents: ['echo'] }, + }) + return createHarnessHandler({ + host, + harness, + authorize: (req) => + req.headers.get('authorization') === 'Bearer e2e-token' + ? { id: 'e2e' } + : null, + }) +} + +export const Route = createFileRoute('/api/harness-protocol/$')({ + server: { + handlers: { + GET: ({ request }) => handlerFor(request)(request), + POST: ({ request }) => handlerFor(request)(request), + }, + }, +}) diff --git a/testing/e2e/tests/harness-protocol.spec.ts b/testing/e2e/tests/harness-protocol.spec.ts new file mode 100644 index 0000000000..43c81f3d47 --- /dev/null +++ b/testing/e2e/tests/harness-protocol.spec.ts @@ -0,0 +1,83 @@ +import { test, expect } from './fixtures' + +test.describe('harness protocol', () => { + const headers = (testId: string, aimockPort: number) => ({ + authorization: 'Bearer e2e-token', + 'x-test-id': testId, + 'x-aimock-port': String(aimockPort), + }) + + test('refuses a request without the token', async ({ request }) => { + const response = await request.get('/api/harness-protocol/capabilities') + expect(response.status()).toBe(401) + }) + + test('lists exposed agents in the capabilities', async ({ + request, + testId, + aimockPort, + }) => { + const response = await request.get('/api/harness-protocol/capabilities', { + headers: headers(testId, aimockPort), + }) + const body = await response.json() + expect(body.identity.name).toBe('e2e/protocol') + expect(body.custom.tanstack.agents[0].name).toBe('echo') + }) + + test('streams a standard AG-UI run through a session', async ({ + request, + testId, + aimockPort, + }) => { + const response = await request.post('/api/harness-protocol/run', { + headers: { + ...headers(testId, aimockPort), + 'content-type': 'application/json', + }, + data: { + threadId: `protocol-${testId}`, + runId: 'client-run', + messages: [ + { id: 'u1', role: 'user', content: '[harness-protocol] hello' }, + ], + tools: [], + context: [], + state: {}, + forwardedProps: {}, + }, + }) + expect(response.headers()['content-type']).toContain('text/event-stream') + const text = await response.text() + const events = text + .split('\n') + .filter((line) => line.startsWith('data: ')) + .map((line) => JSON.parse(line.slice(6))) + const answer = events + .filter((event) => event.type === 'TEXT_MESSAGE_CONTENT') + .map((event) => event.delta) + .join('') + expect(answer).toBe('Hello over the harness protocol.') + expect(events.at(-1).name).toBe('harness.operation.finished') + }) + + test('takes a control input and returns a receipt', async ({ + request, + testId, + aimockPort, + }) => { + const response = await request.post('/api/harness-protocol/control', { + headers: { + ...headers(testId, aimockPort), + 'content-type': 'application/json', + }, + data: { + threadId: `control-${testId}`, + input: { op: 'agent', agent: 'echo', input: { text: 'hi' } }, + }, + }) + const receipt = await response.json() + expect(receipt.status).toBe('accepted') + expect(receipt.operationId).toMatch(/^op-agent-/) + }) +}) From 1eb4b7f584ded69c17b40dd211aba27984baaf6d Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 13:05:33 +0200 Subject: [PATCH 008/286] feat(ai-harness): commands, settings, auth, and first-party plugins Plugins can return commands (defineCommand), settings (configOption), extension point items, and a main-model pick. The setup context adds collect, typed events, persisted plugin state, settings, credentials, and a session API with ask, prompt, transcript, and setConfig. The session adds command, commands, setConfig, config, answer, and inspect, and the protocol accepts command, answer, and config inputs. Auth: oauthConnector adds connect/disconnect commands and gives tools a fresh token. The OAuth runner does PKCE S256 with a single-use loopback on 127.0.0.1, and device-code sign-in. Missing credentials emit harness.auth_required. ai-persistence adds the optional credentials store and compare-and-set on the memory metadata store. First-party plugins at @tanstack/ai-harness/plugins: permissions with modes, workspaceTools, todos, modelPicker, projectInstructions, fileCommands, compact, and usage. The CLI runs plugin commands, /config, /connect, answers questions, and opens sign-in links. --- .changeset/harness-p3-plugins.md | 15 + docs/config.json | 15 + docs/harness/auth.md | 101 ++++ docs/harness/cli.md | 9 +- docs/harness/coding-agent.md | 97 ++++ docs/harness/plugins.md | 156 ++++++ packages/ai-harness-cli/src/commands.ts | 130 ++++- packages/ai-harness-cli/src/interactive.tsx | 23 +- packages/ai-harness-cli/src/lines.ts | 7 + packages/ai-harness-cli/src/session-view.ts | 60 ++- packages/ai-harness-cli/tests/cli.test.ts | 74 ++- packages/ai-harness/package.json | 4 + packages/ai-harness/src/auth.ts | 67 +++ packages/ai-harness/src/commands.ts | 84 ++++ packages/ai-harness/src/config.ts | 88 ++++ packages/ai-harness/src/connectors.ts | 95 ++++ packages/ai-harness/src/extensions.ts | 40 ++ packages/ai-harness/src/first-party/files.ts | 92 ++++ packages/ai-harness/src/first-party/index.ts | 19 + .../src/first-party/model-picker.ts | 62 +++ .../ai-harness/src/first-party/permissions.ts | 150 ++++++ .../src/first-party/session-tools.ts | 115 +++++ packages/ai-harness/src/first-party/todos.ts | 100 ++++ .../ai-harness/src/first-party/workspace.ts | 275 +++++++++++ packages/ai-harness/src/host.ts | 8 +- packages/ai-harness/src/index.ts | 35 ++ packages/ai-harness/src/oauth.ts | 298 +++++++++++ packages/ai-harness/src/plugins.ts | 214 +++++++- packages/ai-harness/src/protocol.ts | 28 ++ packages/ai-harness/src/session.ts | 465 +++++++++++++++++- packages/ai-harness/src/types.ts | 8 + packages/ai-harness/tests/auth.test.ts | 273 ++++++++++ packages/ai-harness/tests/extensions.test.ts | 330 +++++++++++++ packages/ai-harness/tests/first-party.test.ts | 218 ++++++++ packages/ai-harness/tests/helpers.ts | 12 +- packages/ai-harness/vite.config.ts | 2 +- packages/ai-persistence/src/index.ts | 4 + packages/ai-persistence/src/memory.ts | 92 +++- .../ai-persistence/src/testkit/conformance.ts | 71 +++ packages/ai-persistence/src/types.ts | 36 ++ packages/ai-persistence/tests/memory.test.ts | 1 + .../tests/persistence-types.test-d.ts | 2 + .../activities/chat/middleware/metadata.ts | 21 + .../e2e/src/routes/api.harness-protocol.$.ts | 24 + testing/e2e/tests/harness-protocol.spec.ts | 29 ++ 45 files changed, 4012 insertions(+), 37 deletions(-) create mode 100644 .changeset/harness-p3-plugins.md create mode 100644 docs/harness/auth.md create mode 100644 docs/harness/coding-agent.md create mode 100644 docs/harness/plugins.md create mode 100644 packages/ai-harness/src/auth.ts create mode 100644 packages/ai-harness/src/commands.ts create mode 100644 packages/ai-harness/src/config.ts create mode 100644 packages/ai-harness/src/connectors.ts create mode 100644 packages/ai-harness/src/extensions.ts create mode 100644 packages/ai-harness/src/first-party/files.ts create mode 100644 packages/ai-harness/src/first-party/index.ts create mode 100644 packages/ai-harness/src/first-party/model-picker.ts create mode 100644 packages/ai-harness/src/first-party/permissions.ts create mode 100644 packages/ai-harness/src/first-party/session-tools.ts create mode 100644 packages/ai-harness/src/first-party/todos.ts create mode 100644 packages/ai-harness/src/first-party/workspace.ts create mode 100644 packages/ai-harness/src/oauth.ts create mode 100644 packages/ai-harness/tests/auth.test.ts create mode 100644 packages/ai-harness/tests/extensions.test.ts create mode 100644 packages/ai-harness/tests/first-party.test.ts diff --git a/.changeset/harness-p3-plugins.md b/.changeset/harness-p3-plugins.md new file mode 100644 index 0000000000..54c883daa5 --- /dev/null +++ b/.changeset/harness-p3-plugins.md @@ -0,0 +1,15 @@ +--- +'@tanstack/ai-harness': minor +'@tanstack/ai-harness-cli': minor +'@tanstack/ai-persistence': minor +'@tanstack/ai': minor +--- + +Harness plugins can now do everything a Claude Code style agent needs. + +- **Plugin API.** `setup` can return `commands` (`defineCommand`, with typed input), `config` (`configOption.select`, `.boolean`, `.text`, `.number`), `contribute` for extension points (`createExtensionPoint`), and `adapter` to pick the main model. The setup context adds `collect`, typed events (`createPluginEvent`, `emit`, `on`), `state` (saved in the metadata store, published as `STATE_SNAPSHOT`), `config`, `credentials`, and `session` (`ask`, `prompt`, `transcript`, `setConfig`). Prompts can be functions called for each turn. +- **Session.** `session.command`, `session.commands`, `session.setConfig`, `session.config`, `session.answer`, and `session.inspect`. The protocol accepts `command`, `answer`, and `config` inputs. +- **Auth.** `oauthConnector` adds `connect:` and `disconnect:` commands and gives tools a fresh token. The OAuth runner (`loopbackLogin`, `deviceLogin`, `createPkce`, `exchangeCode`, `refreshCredential`) uses PKCE S256 and a single-use loopback on `127.0.0.1`. Tools that need a missing credential emit `harness.auth_required`. +- **First-party plugins** in `@tanstack/ai-harness/plugins`: `permissions` (with `default`, `plan`, `acceptEdits`, and `bypass` modes and the `PermissionRules` extension point), `workspaceTools`, `todos`, `modelPicker`, `projectInstructions`, `fileCommands`, `compact`, and `usage`. +- **`@tanstack/ai-persistence`** adds the optional `credentials` store (`CredentialStore`, `defineCredentialStore`) and compare-and-set on the memory metadata store. **`@tanstack/ai`** adds optional `getVersioned` and `setIf` to `MetadataStore`. +- **CLI.** Plugin commands, `/config`, `/connect`, and `/disconnect`. Questions from plugins are answered on the next line. Sign-in links open in the browser. diff --git a/docs/config.json b/docs/config.json index 4dac967f22..784c711a32 100644 --- a/docs/config.json +++ b/docs/config.json @@ -823,6 +823,21 @@ "label": "Durable sessions", "to": "harness/durable-sessions", "addedAt": "2026-09-26" + }, + { + "label": "Write a plugin", + "to": "harness/plugins", + "addedAt": "2026-09-26" + }, + { + "label": "Build a coding agent", + "to": "harness/coding-agent", + "addedAt": "2026-09-26" + }, + { + "label": "Auth and connectors", + "to": "harness/auth", + "addedAt": "2026-09-26" } ] }, diff --git a/docs/harness/auth.md b/docs/harness/auth.md new file mode 100644 index 0000000000..90e38aab16 --- /dev/null +++ b/docs/harness/auth.md @@ -0,0 +1,101 @@ +--- +title: Auth and connectors +id: harness-auth +order: 7 +description: "Let users sign in to GitHub and other OAuth services from the harness. Tokens stay in your credential store and never reach the model." +keywords: + - tanstack ai + - harness + - oauth + - connectors + - credentials +--- + +Your agent needs to open a pull request, so it needs the user's GitHub token. The user should sign in once, in the browser, and the model should never see the token. `oauthConnector` does that: it adds `/connect github`, stores the token, and hands it to your tools. + +## 1. Add a connector + +```ts group=harness-auth +import { toolDefinition } from '@tanstack/ai' +import { defineHarness, oauthConnector } from '@tanstack/ai-harness' +import { openaiText } from '@tanstack/ai-openai' + +const github = oauthConnector({ + id: 'github', + label: 'GitHub', + oauth: { + authorizationUrl: 'https://github.com/login/oauth/authorize', + tokenUrl: 'https://github.com/login/oauth/access_token', + deviceUrl: 'https://github.com/login/device/code', + clientId: 'your-github-oauth-app-client-id', + scopes: ['repo'], + }, + tools: (token) => [ + toolDefinition({ name: 'list_issues', description: 'List my open issues' }).server( + async () => { + const response = await fetch('https://api.github.com/issues', { + headers: { authorization: `Bearer ${await token()}` }, + }) + return response.json() + }, + ), + ], +}) + +export const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-5.6'), + plugins: () => [github], +}) +``` + +`token()` returns a fresh access token. It refreshes an expired token when the service gave a refresh token. + +## 2. Sign in + +- The user runs `/connect github`. +- The CLI opens the browser. The harness listens on `127.0.0.1` on a random port for one callback, with PKCE and a random `state`. +- The token goes into the credential store. `/disconnect github` deletes it. + +Set `login: 'device'` for SSH sessions, containers, and CI. The CLI then shows a code to enter on the service's page. + +If a tool runs before sign-in, it stops with a `harness.auth_required` event. The CLI shows which service to connect. + +## Keep credentials + +The host reads credentials from `stores.credentials`, keyed by the session's principal. Without it, credentials live in memory until the process stops. + +```ts group=harness-auth +import { defineCredentialStore } from '@tanstack/ai-persistence' +import type { Credential } from '@tanstack/ai-persistence' + +const saved = new Map() +const keyOf = (userId: string | undefined, id: string) => `${userId ?? 'tenant'}:${id}` + +export const credentials = defineCredentialStore({ + get: async (scope, id) => saved.get(keyOf(scope.userId, id)) ?? null, + set: async (scope, id, credential) => { + saved.set(keyOf(scope.userId, id), credential) + }, + delete: async (scope, id) => { + saved.delete(keyOf(scope.userId, id)) + }, + list: async (scope) => + [...saved.entries()] + .filter(([key]) => key.startsWith(`${scope.userId ?? 'tenant'}:`)) + .map(([key, credential]) => ({ id: key.split(':')[1] ?? key, type: credential.type })), +}) +``` + +- Encrypt tokens at rest in a real store. +- `list` returns ids and types only, never the secret values. +- A credential saved without a `userId` belongs to the whole tenant. + +## Read credentials in your own plugin + +`ctx.credentials.require('github')` returns the credential, or stops with `auth_required` when the user has not signed in. Plugins read only the credentials of the session's principal. + +## What you have now + +- Browser and device-code sign-in for any OAuth service. +- Tools that get a fresh token without the model seeing it. diff --git a/docs/harness/cli.md b/docs/harness/cli.md index a572a2bc02..9889069159 100644 --- a/docs/harness/cli.md +++ b/docs/harness/cli.md @@ -68,13 +68,20 @@ When stdin is a pipe, the CLI reads one message or command per line and waits fo ## Commands in the interactive UI +For the session: + - `/agents`: list the agents. - `/agent {"json":"input"}`: run an agent in the background. When it is done, a new turn starts with its result. +- `/config`: show the settings. `/config ` changes one. +- `/connect ` and `/disconnect `: sign in to a connector, or out. + +For the running work: + - `/cancel`: cancel the running turn. - `/status`: show what runs and what waits. - `/exit`: quit. -When a turn stops for an approval, type `y` to approve or `n` to reject. +Plugin commands (for example `/model` or `/todos`) show up in `/help`. When a turn stops for an approval or a plugin asks a question, type your answer. For yes-or-no questions, `y` approves and `n` refuses. ## What you have now diff --git a/docs/harness/coding-agent.md b/docs/harness/coding-agent.md new file mode 100644 index 0000000000..8a45100343 --- /dev/null +++ b/docs/harness/coding-agent.md @@ -0,0 +1,97 @@ +--- +title: Build a coding agent +id: harness-coding-agent +order: 6 +description: "Turn a harness into a coding agent with file tools, permission modes, a todo list, a model picker, project instructions, and /compact." +keywords: + - tanstack ai + - harness + - coding agent + - permissions + - workspace tools +--- + +You want your own coding agent in the terminal: it reads and edits files, runs commands after you approve them, and keeps a todo list. The first-party plugins in `@tanstack/ai-harness/plugins` give you those parts. You pick the model and the rules. + +## 1. Define the agent + +```ts group=harness-coding-agent +import { defineHarness } from '@tanstack/ai-harness' +import { + compact, + fileCommands, + modelPicker, + permissions, + projectInstructions, + todos, + usage, + workspaceTools, +} from '@tanstack/ai-harness/plugins' +import { runCli } from '@tanstack/ai-harness-cli' +import { openaiText } from '@tanstack/ai-openai' + +const root = process.cwd() +const smart = openaiText('gpt-5.6') +const fast = openaiText('gpt-5.6-luna') + +const coder = defineHarness({ + name: 'acme/coder', + adapter: smart, + systemPrompts: ['You are a careful coding agent. Read before you edit.'], + plugins: () => [ + permissions(), + workspaceTools({ root }), + todos(), + modelPicker({ choices: { smart, fast }, default: 'smart' }), + projectInstructions({ root }), + fileCommands({ dir: `${root}/.claude/commands` }), + compact({ adapter: fast }), + usage(), + ], +}) + +process.exitCode = await runCli(coder) +``` + +## 2. Run it + +Run the file with `npx tsx coder.ts`. Ask for a change. The agent reads files freely and asks before it writes a file or runs a command. Type `y` to allow a call or `n` to refuse it. + +## What each plugin adds + +| Plugin | Adds | +|---|---| +| `permissions()` | Asks before risky tool calls. `/mode` switches between `default`, `plan` (read-only), `acceptEdits` (edits run without asking), and `bypass`. | +| `workspaceTools({ root })` | `read_file`, `write_file`, `edit_file`, `list_files`, `grep`, and `bash`, confined to `root`. | +| `todos()` | A `todo_write` tool the model uses for multi-step work, and `/todos`. | +| `modelPicker({ choices })` | `/model ` switches the model at the next turn. | +| `projectInstructions({ root })` | Adds `AGENTS.md` and `CLAUDE.md` to the system prompt. | +| `fileCommands({ dir })` | Each `.md` file becomes a slash command. `$ARGUMENTS` is replaced by what you type after it. | +| `compact({ adapter })` | `/compact` replaces a long conversation with a summary. | +| `usage()` | `/usage` shows the tokens of the session. | + +## Add your own rules + +Tool plugins add permission rules to the `PermissionRules` extension point. Add your own for any tool: + +```ts group=harness-coding-agent +import { PermissionRules } from '@tanstack/ai-harness/plugins' +import { definePlugin } from '@tanstack/ai-harness' + +export const noDeploys = definePlugin({ + name: 'acme/no-deploys', + setup: () => ({ + contribute: [PermissionRules.item({ tool: 'deploy_*', decision: 'deny' })], + }), +}) +``` + +A trailing `*` matches every tool that starts with the text. The last matching rule wins. + +The workspace tools run on your machine with your permissions. Run code you do not trust in a sandbox. + +## What you have now + +- A terminal coding agent with file tools, approvals, modes, a todo list, and a model picker. + +Next: connect the agent to GitHub and other services with [auth and connectors](./auth). diff --git a/docs/harness/plugins.md b/docs/harness/plugins.md new file mode 100644 index 0000000000..4e01346a91 --- /dev/null +++ b/docs/harness/plugins.md @@ -0,0 +1,156 @@ +--- +title: Write a harness plugin +id: harness-plugins +order: 5 +description: "Package tools, prompts, middleware, settings, commands, and state as a plugin that any harness can load." +keywords: + - tanstack ai + - harness + - plugins + - definePlugin + - extension points +--- + +You built a feature for one harness, a todo list or a GitHub integration, and now you want it in every harness you ship. A plugin packages it: tools and prompts for the model, commands for the user, settings, and state that survives restarts. + +## The smallest plugin + +```ts group=harness-plugins +import { definePlugin } from '@tanstack/ai-harness' + +export const today = definePlugin({ + name: 'acme/today', + setup: () => ({ prompts: [`Today is ${new Date().toDateString()}.`] }), +}) +``` + +Add it with `plugins: () => [today]` in `defineHarness`. `setup` runs once per session. + +## What `setup` can return + +- `tools`: tools for the model, made with `toolDefinition`. +- `prompts`: text for the system prompt. A function runs for each turn, so it can show current state. +- `middleware`: chat middleware, the same type as `chat({ middleware })`. +- `generationMiddleware`: middleware for the activities agents call. +- `agents`: agents added to `session.agents`. +- `commands`: user actions, see below. +- `config`: session settings, see below. +- `contribute`: items for another plugin's extension point. + +If two plugins add the same tool, command, setting, or prompt id, `host.open` fails and names both plugins. + +## Add commands and settings + +Commands are actions for the user, not the model: slash commands, buttons, dashboard actions. Settings are typed, saved per session, and apply at the next turn. + +```ts group=harness-plugins +import { configOption, defineCommand } from '@tanstack/ai-harness' +import { z } from 'zod' + +export const greeter = definePlugin({ + name: 'acme/greeter', + setup: (ctx) => ({ + config: { + tone: configOption.select({ options: ['plain', 'warm'], default: 'plain' }), + }, + commands: { + greet: defineCommand({ + description: 'Greet someone', + input: z.object({ name: z.string() }), + run: async ({ name }, command) => { + const sure = await command.session.ask({ + message: `Greet ${name}?`, + schema: z.boolean(), + }) + if (!sure) return 'Skipped.' + return ctx.config.get('tone') === 'warm' ? `Hello, dear ${name}!` : `Hello, ${name}.` + }, + }), + }, + }), +}) +``` + +- `session.command('greet', { name: 'Ada' })` runs the command. The input is checked against its schema. +- `session.setConfig('tone', 'warm')` changes the setting. A value the option does not accept is rejected. +- `ask` waits for `session.answer(questionId, value)`. The question shows up in `session.snapshot().pendingQuestions`, and the CLI prompts for it. + +## Keep state + +`ctx.state(initial)` gives a plugin its own state, saved in the metadata store. Clients see each change as an AG-UI `STATE_SNAPSHOT` event. + +```ts group=harness-plugins +export const counter = definePlugin({ + name: 'acme/counter', + setup: (ctx) => { + const state = ctx.state({ count: 0 }) + return { + commands: { + bump: defineCommand({ + description: 'Add one', + run: async () => (await state.update((current) => ({ count: current.count + 1 }))).count, + }), + }, + } + }, +}) +``` + +When two writers race, `update` runs your function again with fresh state. + +## Let plugins work together + +Three ways, from simple to loose: + +- **Capabilities.** One plugin declares `provides: [cap]` and calls `ctx.provide(cap, value)`. Another declares `requires: [cap]` and reads `ctx.get(cap)`. The provider must come first. +- **Extension points.** One plugin reads a list, and others add to it in any order. +- **Events.** `ctx.emit(event, value)` reaches every `ctx.on(event, handler)`. + +```ts group=harness-plugins +import { createExtensionPoint, createPluginEvent } from '@tanstack/ai-harness' + +export const Checks = createExtensionPoint<{ name: string }>('acme/checks') +export const checked = createPluginEvent<{ name: string }>('acme/checked') + +export const runner = definePlugin({ + name: 'acme/check-runner', + setup: (ctx) => { + const checks = ctx.collect(Checks) + return { + commands: { + check: defineCommand({ + description: 'Run every check', + run: () => { + for (const check of checks) ctx.emit(checked, check) + return checks.map((check) => check.name) + }, + }), + }, + } + }, +}) + +export const lint = definePlugin({ + name: 'acme/lint', + setup: () => ({ contribute: [Checks.item({ name: 'lint' })] }), +}) +``` + +Read `ctx.collect` at run time (in a command, a tool, or a middleware hook). During `setup`, later plugins have not added their items yet. + +## Clean up resources + +Open resources with `ctx.resources.acquire(open, close)`. They close when the session closes, newest first. If a later plugin fails in `setup`, every plugin set up so far is cleaned up, and no model request is sent. + +Set `lifetime: 'run'` to set a plugin up again for each turn. + +## See the plan + +`session.inspect()` lists the plugins in order, and who owns each tool, prompt, command, setting, and extension point item. + +## What you have now + +- A plugin that adds tools, prompts, commands, settings, and state to any harness. +- Plugins that share services, lists, and events without knowing each other. + +Next: see the [first-party plugins](./coding-agent) that turn a harness into a coding agent. diff --git a/packages/ai-harness-cli/src/commands.ts b/packages/ai-harness-cli/src/commands.ts index 4c0ce10d28..3f5ef06fd7 100644 --- a/packages/ai-harness-cli/src/commands.ts +++ b/packages/ai-harness-cli/src/commands.ts @@ -5,14 +5,79 @@ export const HELP_TEXT = [ 'Commands:', ' /agents list the agents', ' /agent [json] run an agent in the background', + ' /config [key value] show or change a setting', + ' /connect sign in to a connector (and /disconnect )', ' /cancel cancel the running turn (or press Esc)', ' /status show the session status', ' /exit quit', ].join('\n') +/** A command argument: JSON when it parses, else the text itself. */ +function parseArg(text: string): unknown { + if (text === '') return undefined + try { + return JSON.parse(text) + } catch { + return text + } +} + +function show(value: unknown): string { + if (value === undefined) return 'Done.' + return typeof value === 'string' ? value : JSON.stringify(value, null, 2) +} + +/** + * Turn the user's reply into an answer: `y`/`n` for yes-or-no questions, + * JSON when it parses, else the text. + */ +export function parseAnswer(reply: string, schema: unknown): unknown { + const text = reply.trim() + const isBoolean = + typeof schema === 'object' && + schema !== null && + 'type' in schema && + schema.type === 'boolean' + if (isBoolean && /^(y|yes)$/i.test(text)) return true + if (isBoolean && /^(n|no)$/i.test(text)) return false + return parseArg(text) +} + +async function runCommand( + session: HarnessSession, + name: string, + input: unknown, +): Promise { + const describe = ( + outcome: { ok: true; value: unknown } | { ok: false; error: unknown }, + ) => + outcome.ok + ? show(outcome.value) + : `/${name} failed: ${outcome.error instanceof Error ? outcome.error.message : String(outcome.error)}` + const settled = Promise.resolve( + session.command(name, input).then( + (value) => ({ ok: true as const, value }), + (error: unknown) => ({ ok: false as const, error }), + ), + ) + // A command that asks a question waits for the next line, so return now + // and print its result when it is done. + let done = false + void settled.then(() => { + done = true + }) + while (!done) { + if (session.snapshot().pendingQuestions.length > 0) { + return { type: 'notice', text: '', later: settled.then(describe) } + } + await new Promise((resolve) => setTimeout(resolve, 10)) + } + return { type: 'notice', text: describe(await settled) } +} + /** What the UI does after a line of input. */ export type LineResult = - | { type: 'notice'; text: string } + | { type: 'notice'; text: string; later?: Promise } | { type: 'exit' } | { type: 'sent' } @@ -26,6 +91,20 @@ export async function handleLine( ): Promise { const text = line.trim() if (text === '') return { type: 'notice', text: '' } + const [question] = session.snapshot().pendingQuestions + if (question) { + const receipt = await session.answer( + question.questionId, + parseAnswer(text, question.schema), + ) + return { + type: 'notice', + text: + receipt.status === 'rejected' + ? `Answer again: ${receipt.reason ?? ''}` + : '', + } + } if (!text.startsWith('/')) { if (session.snapshot().status === 'running') await session.steer(text) else void session.prompt(text) @@ -34,8 +113,50 @@ export async function handleLine( const [name = '', ...rest] = text.slice(1).split(' ') const argText = rest.join(' ').trim() switch (name) { - case 'help': - return { type: 'notice', text: HELP_TEXT } + case 'help': { + const commands = session.commands() + return { + type: 'notice', + text: + commands.length === 0 + ? HELP_TEXT + : `${HELP_TEXT}\nPlugin commands:\n${commands + .map((command) => ` /${command.name} ${command.description}`) + .join('\n')}`, + } + } + case 'config': { + const [key = '', ...valueParts] = argText.split(' ') + if (!key) { + const entries = Object.entries(session.config()) + return { + type: 'notice', + text: + entries.length === 0 + ? 'This harness has no settings.' + : entries + .map( + ([name, entry]) => + ` ${name} = ${JSON.stringify(entry.value)}`, + ) + .join('\n'), + } + } + const receipt = await session.setConfig( + key, + parseArg(valueParts.join(' ').trim()), + ) + return { + type: 'notice', + text: + receipt.status === 'rejected' + ? `Not changed: ${receipt.reason ?? ''}` + : `${key} changed.`, + } + } + case 'connect': + case 'disconnect': + return runCommand(session, `${name}:${argText}`, undefined) case 'exit': case 'quit': return { type: 'exit' } @@ -95,6 +216,9 @@ export async function handleLine( return { type: 'notice', text: `Started ${agentName} in the background.` } } default: + if (session.commands().some((command) => command.name === name)) { + return runCommand(session, name, parseArg(argText)) + } return { type: 'notice', text: `Unknown command: /${name}. Type /help.` } } } diff --git a/packages/ai-harness-cli/src/interactive.tsx b/packages/ai-harness-cli/src/interactive.tsx index d05540ed51..fbfc35bb45 100644 --- a/packages/ai-harness-cli/src/interactive.tsx +++ b/packages/ai-harness-cli/src/interactive.tsx @@ -3,7 +3,12 @@ import { Box, Text, render, useApp, useInput } from 'ink' import { HARNESS_EVENTS } from '@tanstack/ai-harness' import { EventType } from '@tanstack/ai' import { handleLine } from './commands' -import { applyEvent, approvalQuestion, resolveAll } from './session-view' +import { + applyEvent, + approvalQuestion, + openUrl, + resolveAll, +} from './session-view' import type { AnyHarness, HarnessSession, @@ -51,11 +56,11 @@ function App({ })) { setEntries((current) => applyEvent(current, entry).slice(-MAX_ENTRIES)) const event = entry.event - if ( - event.type === EventType.CUSTOM && - (event.name === HARNESS_EVENTS.operationStarted || - event.name === HARNESS_EVENTS.operationFinished) - ) { + if (event.type === EventType.CUSTOM) { + if (event.name === HARNESS_EVENTS.authRequired) { + const url = (event.value as { url?: unknown }).url + if (typeof url === 'string') openUrl(url) + } // Let the operation settle before reading the status. setTimeout(() => setSnapshot(session.snapshot()), 0) } @@ -119,8 +124,10 @@ function App({ })() }) - const status = - snapshot.status === 'running' + const question = snapshot.pendingQuestions[0] + const status = question + ? `${question.message} (type your answer)` + : snapshot.status === 'running' ? 'working (Esc to cancel, Enter steers)' : snapshot.status === 'requires_action' ? approvalQuestion(snapshot.pendingInterrupts) diff --git a/packages/ai-harness-cli/src/lines.ts b/packages/ai-harness-cli/src/lines.ts index c16ff9e0ba..26cc914f8e 100644 --- a/packages/ai-harness-cli/src/lines.ts +++ b/packages/ai-harness-cli/src/lines.ts @@ -52,6 +52,7 @@ export async function runLines( })() const lines = createInterface({ input, crlfDelay: Infinity }) + const pendingLater: Array> = [] for await (const line of lines) { const snapshot = session.snapshot() if (snapshot.status === 'requires_action') { @@ -65,6 +66,11 @@ export async function runLines( if (result.type === 'exit') break if (result.type === 'notice' && result.text) stdout.write(`${result.text}\n`) + if (result.type === 'notice' && result.later) { + pendingLater.push( + result.later.then((text) => void stdout.write(`\n${text}\n`)), + ) + } } await waitIdle(session) const after = session.snapshot() @@ -73,6 +79,7 @@ export async function runLines( } } lines.close() + await Promise.allSettled(pendingLater) await waitIdle(session) reader.abort() await printing diff --git a/packages/ai-harness-cli/src/session-view.ts b/packages/ai-harness-cli/src/session-view.ts index 24325f072d..43df00f9bb 100644 --- a/packages/ai-harness-cli/src/session-view.ts +++ b/packages/ai-harness-cli/src/session-view.ts @@ -53,6 +53,41 @@ export function applyEvent( { kind: 'notice', text: 'Resumed a turn that a crash stopped.' }, ] } + if ( + event.type === EventType.CUSTOM && + event.name === HARNESS_EVENTS.question + ) { + const value = event.value as { message?: unknown } + return [ + ...entries, + { kind: 'notice', text: `? ${String(value.message ?? '')}` }, + ] + } + if ( + event.type === EventType.CUSTOM && + event.name === HARNESS_EVENTS.authRequired + ) { + const value = event.value as { + connector?: unknown + url?: unknown + userCode?: unknown + } + const code = + typeof value.userCode === 'string' + ? ` and enter the code ${value.userCode}` + : '' + const where = + typeof value.url === 'string' + ? ` Open ${value.url}${code}.` + : ` Run /connect ${String(value.connector)}.` + return [ + ...entries, + { + kind: 'notice', + text: `Sign in to ${String(value.connector)}.${where}`, + }, + ] + } return entries } @@ -70,6 +105,25 @@ export function resolveAll( return session.resolve(resume) } +/** Open an http(s) URL in the default browser. No shell is involved. */ +export function openUrl(url: string): void { + if (!/^https?:\/\//.test(url)) return + void import('node:child_process').then(({ spawn }) => { + const [command, args] = + process.platform === 'win32' + ? // Not `cmd /c start`: cmd would read `&` in the URL as a command separator. + ['rundll32', ['url.dll,FileProtocolHandler', url]] + : process.platform === 'darwin' + ? ['open', [url]] + : ['xdg-open', [url]] + try { + spawn(command, args, { stdio: 'ignore', detached: true }).unref() + } catch { + // The notice still shows the URL. + } + }) +} + /** A short question for the open interrupts. */ export function approvalQuestion(interrupts: ReadonlyArray): string { const names = interrupts.map( @@ -78,10 +132,14 @@ export function approvalQuestion(interrupts: ReadonlyArray): string { return `Approve ${names.join(', ')}? [y/n]` } -/** Resolves when no chat turn runs or waits in the queue. */ +/** + * Resolves when no chat turn runs or waits in the queue, or when the + * session waits for an answer to a question. + */ export async function waitIdle(session: HarnessSession): Promise { while (true) { const snapshot = session.snapshot() + if (snapshot.pendingQuestions.length > 0) return const chatActive = snapshot.activeOperations.some( (operation) => operation.kind === 'chat', ) diff --git a/packages/ai-harness-cli/tests/cli.test.ts b/packages/ai-harness-cli/tests/cli.test.ts index 645e5b9537..7b2b8649b5 100644 --- a/packages/ai-harness-cli/tests/cli.test.ts +++ b/packages/ai-harness-cli/tests/cli.test.ts @@ -1,7 +1,13 @@ import { Readable } from 'node:stream' import { describe, expect, it, vi } from 'vitest' import { EventType, defineAgent, toolDefinition } from '@tanstack/ai' -import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { + configOption, + createHarnessHost, + defineCommand, + defineHarness, + definePlugin, +} from '@tanstack/ai-harness' import { memoryPersistence } from '@tanstack/ai-persistence' import { EXIT, parseCliArgs, runCli } from '../src' import { serve } from '../src/serve' @@ -222,6 +228,72 @@ describe('line mode', () => { }) }) +describe('plugin commands in line mode', () => { + it('runs plugin commands, answers questions, and changes settings', async () => { + const tools = definePlugin({ + name: 'test/tools', + setup: (ctx) => ({ + config: { + tone: configOption.select({ + options: ['plain', 'warm'], + default: 'plain', + }), + }, + commands: { + greet: defineCommand({ + description: 'Greet someone', + run: (input: unknown) => + `Hello, ${typeof input === 'object' && input !== null && 'name' in input ? String(input.name) : 'you'}!`, + }), + confirm: defineCommand({ + description: 'Ask first', + run: async () => { + const sure = await ctx.session.ask({ + message: 'Really?', + schema: { type: 'boolean' }, + }) + return sure === true ? 'Confirmed.' : 'Stopped.' + }, + }), + }, + }), + }) + const { adapter } = scripted([]) + const stdout = capture() + const input = Readable.from([ + '/greet {"name":"Ada"}\n', + '/confirm\n', + 'yes\n', + '/config tone warm\n', + '/config\n', + '/help\n', + '/exit\n', + ]) + await runCli( + defineHarness({ + name: 'test/plugin-lines', + adapter, + plugins: () => [tools], + }), + { + argv: [], + stdin: Object.assign(input, { + isTTY: false, + }) as unknown as NodeJS.ReadStream, + stdout, + stderr: capture(), + persistence: memoryPersistence(), + }, + ) + expect(stdout.text).toContain('Hello, Ada!') + expect(stdout.text).toContain('? Really?') + expect(stdout.text).toContain('Confirmed.') + expect(stdout.text).toContain('tone changed.') + expect(stdout.text).toContain('tone = "warm"') + expect(stdout.text).toContain('/greet Greet someone') + }) +}) + describe('serve mode', () => { it('serves the session protocol and requires the token', async () => { const { adapter } = scripted([textTurn('served')]) diff --git a/packages/ai-harness/package.json b/packages/ai-harness/package.json index fd59be8e31..01b208a614 100644 --- a/packages/ai-harness/package.json +++ b/packages/ai-harness/package.json @@ -33,6 +33,10 @@ "./client": { "types": "./dist/esm/client.d.ts", "import": "./dist/esm/client.js" + }, + "./plugins": { + "types": "./dist/esm/first-party/index.d.ts", + "import": "./dist/esm/first-party/index.js" } }, "scripts": { diff --git a/packages/ai-harness/src/auth.ts b/packages/ai-harness/src/auth.ts new file mode 100644 index 0000000000..a41e8a01a2 --- /dev/null +++ b/packages/ai-harness/src/auth.ts @@ -0,0 +1,67 @@ +import type { Scope } from '@tanstack/ai' +import type { Credential, CredentialStore } from '@tanstack/ai-persistence' + +/** + * Thrown when a tool or command needs a credential the user has not saved. + * The session publishes `harness.auth_required`, so a host can show the + * sign-in link or run `/connect`. + */ +export class AuthRequiredError extends Error { + readonly connector: string + readonly url: string | undefined + constructor(connector: string, url?: string) { + super( + url + ? `Sign in to ${connector} first: ${url}` + : `Sign in to ${connector} first. Run /connect ${connector}.`, + ) + this.name = 'AuthRequiredError' + this.connector = connector + this.url = url + } +} + +/** Credentials of the session's principal, as plugins see them. */ +export interface CredentialsAccess { + get: (id: string) => Promise + /** Like `get`, but throws {@link AuthRequiredError} when the credential is missing. */ + require: (id: string) => Promise + set: (id: string, credential: Credential) => Promise + delete: (id: string) => Promise + list: () => Promise< + Array<{ id: string; type: Credential['type']; expiresAt?: number }> + > +} + +/** Scope credentials to one principal and thread. */ +export function credentialsFor( + store: CredentialStore, + scope: Scope, + onMissing: (error: AuthRequiredError) => void, +): CredentialsAccess { + return { + get: (id) => store.get(scope, id), + require: async (id) => { + const credential = await store.get(scope, id) + if (credential) return credential + const error = new AuthRequiredError(id) + onMissing(error) + throw error + }, + set: (id, credential) => store.set(scope, id, credential), + delete: (id) => store.delete(scope, id), + list: () => store.list(scope), + } +} + +/** Remove secret-looking values from a message before it reaches a model or a log. */ +export function scrubSecrets( + text: string, + secrets: ReadonlyArray, +): string { + let scrubbed = text + for (const secret of secrets) { + if (secret.length >= 6) scrubbed = scrubbed.split(secret).join('[redacted]') + } + return scrubbed +} diff --git a/packages/ai-harness/src/commands.ts b/packages/ai-harness/src/commands.ts new file mode 100644 index 0000000000..86a48a872e --- /dev/null +++ b/packages/ai-harness/src/commands.ts @@ -0,0 +1,84 @@ +import type { InferSchemaType, ModelMessage, SchemaInput } from '@tanstack/ai' +import type { SessionSnapshot } from './session' +import type { Principal, Receipt } from './types' + +/** A question a command asks the user. The host renders it from the schema. */ +export interface Question { + message: string + /** The answer's shape. Without one, the answer is any text. */ + schema?: TSchema +} + +export type AnswerOf = TSchema extends SchemaInput + ? InferSchemaType + : string + +/** The parts of the session a plugin can use. */ +export interface PluginSessionApi { + threadId: string + principal: Principal | undefined + snapshot: () => SessionSnapshot + /** Start a chat turn, as if the user typed `text`. */ + prompt: (text: string) => void + /** The saved transcript. */ + transcript: () => Promise> + /** Replace the saved transcript (for example after a summary). */ + replaceTranscript: (messages: Array) => Promise + /** + * Ask the user and wait for the answer. Hosts show it (the CLI prompts, + * ACP shows an elicitation). Rejects when the session closes. + */ + ask: ( + question: Question, + ) => Promise> + /** + * Tell clients that the user must sign in. Hosts show the link (the CLI + * opens the browser) or the device code. + */ + authRequired: (info: { + connector: string + url?: string + userCode?: string + }) => void + /** Change a session setting, the same as a client. Applies at the next turn. */ + setConfig: (key: string, value: unknown) => Promise +} + +/** What a command handler receives besides its input. */ +export interface CommandContext { + signal: AbortSignal + session: PluginSessionApi +} + +/** A user action: a slash command, a button, a dashboard action. Not a model tool. */ +export interface CommandDefinition< + TSchema extends SchemaInput | undefined = SchemaInput | undefined, +> { + description: string + /** The input shape, checked before `run`. */ + input?: TSchema + run: ( + input: TSchema extends SchemaInput ? InferSchemaType : undefined, + ctx: CommandContext, + ) => unknown +} + +/** A command with any input type. */ +export type AnyCommand = CommandDefinition + +/** + * Define a command with a typed input. + * + * @example + * ```ts + * const clear = defineCommand({ + * description: 'Remove every todo', + * run: async () => ({ removed: await db.clear() }), + * }) + * ``` + */ +export function defineCommand< + const TSchema extends SchemaInput | undefined = undefined, +>(command: CommandDefinition): CommandDefinition { + return command +} diff --git a/packages/ai-harness/src/config.ts b/packages/ai-harness/src/config.ts new file mode 100644 index 0000000000..1ae460642f --- /dev/null +++ b/packages/ai-harness/src/config.ts @@ -0,0 +1,88 @@ +/** + * A session setting a plugin declares, for example the model or the + * thinking level. Hosts render it from this description (a picker in the + * CLI, `configOptions` in ACP). A change applies at the next turn. + */ +export type ConfigOption = + | { + type: 'select' + options: ReadonlyArray + default: string + description?: string + /** ACP category, for example `'model'` or `'thought_level'`. */ + category?: string + } + | { type: 'boolean'; default: boolean; description?: string } + | { type: 'text'; default: string; description?: string } + | { + type: 'number' + default: number + min?: number + max?: number + description?: string + } + +type Without = Omit +type Option = Extract< + ConfigOption, + { type: TType } +> + +/** Helpers to declare {@link ConfigOption}s. */ +export const configOption = { + select: (option: Without, 'type'>): Option<'select'> => { + if (!option.options.includes(option.default)) { + throw new Error( + `configOption.select: the default "${option.default}" is not an option.`, + ) + } + return { type: 'select', ...option } + }, + boolean: (option: Without, 'type'>): Option<'boolean'> => ({ + type: 'boolean', + ...option, + }), + text: (option: Without, 'type'>): Option<'text'> => ({ + type: 'text', + ...option, + }), + number: (option: Without, 'type'>): Option<'number'> => ({ + type: 'number', + ...option, + }), +} + +/** Check a value for an option. Returns the value, or throws with a reason. */ +export function checkConfigValue( + key: string, + option: ConfigOption, + value: unknown, +): unknown { + const fail = (expected: string): never => { + throw new Error( + `Config "${key}" expects ${expected}, got ${JSON.stringify(value)}.`, + ) + } + switch (option.type) { + case 'select': + if (typeof value !== 'string' || !option.options.includes(value)) { + fail(`one of ${option.options.join(', ')}`) + } + return value + case 'boolean': + if (typeof value !== 'boolean') fail('true or false') + return value + case 'text': + if (typeof value !== 'string') fail('text') + return value + case 'number': { + if (typeof value !== 'number' || !Number.isFinite(value)) + return fail('a number') + if (option.min !== undefined && value < option.min) + fail(`a number >= ${option.min}`) + if (option.max !== undefined && value > option.max) + fail(`a number <= ${option.max}`) + return value + } + } +} diff --git a/packages/ai-harness/src/connectors.ts b/packages/ai-harness/src/connectors.ts new file mode 100644 index 0000000000..7b5aad47cd --- /dev/null +++ b/packages/ai-harness/src/connectors.ts @@ -0,0 +1,95 @@ +import { defineCommand } from './commands' +import { definePlugin } from './plugins' +import { + deviceLogin, + isExpired, + loopbackLogin, + refreshCredential, +} from './oauth' +import type { AnyTool } from '@tanstack/ai' +import type { OAuthConfig } from './oauth' + +export interface OAuthConnectorOptions { + /** A short id, for example `'github'`. Commands are `connect:` and `disconnect:`. */ + id: string + label: string + oauth: OAuthConfig + /** `'loopback'` (default) opens a browser here. `'device'` shows a code to enter elsewhere. */ + login?: 'loopback' | 'device' + /** + * The tools of this service. Call `token()` inside a tool: it returns a + * fresh access token, or pauses with `auth_required` before sign-in. + */ + tools?: (token: () => Promise) => ReadonlyArray + /** Test hook for the token endpoint. */ + fetch?: typeof fetch +} + +/** + * A plugin that signs the user in to an OAuth service and gives the model + * that service's tools. Adds `connect:` and `disconnect:` commands. + * Tokens stay in the credential store. The model never sees them. + * + * @example + * ```ts + * const github = oauthConnector({ + * id: 'github', + * label: 'GitHub', + * oauth: { authorizationUrl, tokenUrl, deviceUrl, clientId, scopes: ['repo'] }, + * tools: (token) => [listIssues(token)], + * }) + * ``` + */ +export function oauthConnector(options: OAuthConnectorOptions) { + const { id, label, oauth } = options + return definePlugin({ + name: `connector/${id}`, + setup: (ctx) => { + const token = async (): Promise => { + let credential = await ctx.credentials.require(id) + if (credential.type === 'api_key') return credential.value + if (isExpired(credential) && credential.refreshToken) { + credential = await refreshCredential(oauth, credential, options.fetch) + await ctx.credentials.set(id, credential) + } + if (credential.type === 'api_key') return credential.value + return credential.accessToken + } + return { + tools: options.tools?.(token) ?? [], + commands: { + [`connect:${id}`]: defineCommand({ + description: `Sign in to ${label}`, + run: async () => { + const credential = + options.login === 'device' + ? await deviceLogin(oauth, { + onCode: ({ userCode, verificationUri }) => + ctx.session.authRequired({ + connector: id, + url: verificationUri, + userCode, + }), + ...(options.fetch ? { fetch: options.fetch } : {}), + }) + : await loopbackLogin(oauth, { + onUrl: (url) => + ctx.session.authRequired({ connector: id, url }), + ...(options.fetch ? { fetch: options.fetch } : {}), + }) + await ctx.credentials.set(id, credential) + return `Connected to ${label}.` + }, + }), + [`disconnect:${id}`]: defineCommand({ + description: `Sign out of ${label}`, + run: async () => { + await ctx.credentials.delete(id) + return `Disconnected from ${label}.` + }, + }), + }, + } + }, + }) +} diff --git a/packages/ai-harness/src/extensions.ts b/packages/ai-harness/src/extensions.ts new file mode 100644 index 0000000000..3181822750 --- /dev/null +++ b/packages/ai-harness/src/extensions.ts @@ -0,0 +1,40 @@ +/** + * A list one plugin reads and other plugins add to, in any order. For + * example, a permissions plugin reads rules that tool plugins contribute. + * + * @example + * ```ts + * export const PermissionRules = createExtensionPoint('acme/permission-rules') + * // a tool plugin: contribute: [PermissionRules.item({ tool: 'write_file', decision: 'ask' })] + * // the permissions plugin, at run time: ctx.collect(PermissionRules) + * ``` + */ +export interface ExtensionPoint { + readonly name: string + item: (value: T) => ExtensionItem +} + +/** One contribution to an extension point. */ +export interface ExtensionItem { + readonly point: string + readonly value: T +} + +export function createExtensionPoint(name: string): ExtensionPoint { + return { name, item: (value) => ({ point: name, value }) } +} + +/** + * A typed plugin event. `ctx.emit(event, value)` sends it to every + * `ctx.on(event, handler)` in the session, and clients see it as a + * `harness.plugin.event` CUSTOM event. + */ +export interface PluginEvent { + readonly name: string + /** Type only. Never read at runtime. */ + readonly __type?: T +} + +export function createPluginEvent(name: string): PluginEvent { + return { name } +} diff --git a/packages/ai-harness/src/first-party/files.ts b/packages/ai-harness/src/first-party/files.ts new file mode 100644 index 0000000000..2b9fa2e0d2 --- /dev/null +++ b/packages/ai-harness/src/first-party/files.ts @@ -0,0 +1,92 @@ +import { readFile, readdir } from 'node:fs/promises' +import { basename, extname, resolve } from 'node:path' +import { defineCommand } from '../commands' +import { definePlugin } from '../plugins' +import type { AnyCommand } from '../commands' + +async function readIfPresent(path: string): Promise { + try { + return await readFile(path, 'utf8') + } catch { + return undefined + } +} + +/** + * Add project instruction files (AGENTS.md, CLAUDE.md) to the system prompt, + * read once when the session opens. + */ +export function projectInstructions(options: { + root: string + files?: ReadonlyArray +}) { + const files = options.files ?? ['AGENTS.md', 'CLAUDE.md'] + return definePlugin({ + name: 'tanstack/project-instructions', + setup: async () => { + const found: Array = [] + for (const file of files) { + const text = await readIfPresent(resolve(options.root, file)) + if (text?.trim()) + found.push(`Project instructions from ${file}:\n${text.trim()}`) + } + return { prompts: found } + }, + }) +} + +/** Split `---` frontmatter from a Markdown file. */ +function splitFrontmatter(text: string): { + description?: string + body: string +} { + const match = /^---\r?\n([\s\S]*?)\r?\n---\r?\n?/.exec(text) + if (!match) return { body: text } + const description = /^description:\s*(.+)$/m.exec(match[1] ?? '')?.[1]?.trim() + return { + ...(description ? { description } : {}), + body: text.slice(match[0].length), + } +} + +/** + * One command per Markdown file in `dir` (like `.claude/commands/*.md`). + * Running `/name args` sends the file as a prompt, with `$ARGUMENTS` + * replaced by the args. + */ +export function fileCommands(options: { dir: string }) { + return definePlugin({ + name: 'tanstack/file-commands', + setup: async (ctx) => { + let names: Array = [] + try { + names = (await readdir(options.dir)).filter( + (file) => extname(file) === '.md', + ) + } catch { + return {} + } + const commands: Record = {} + for (const file of names) { + const name = basename(file, '.md') + const { description, body } = splitFrontmatter( + await readFile(resolve(options.dir, file), 'utf8'), + ) + commands[name] = defineCommand({ + description: description ?? `Run ${file}`, + run: (input: unknown) => { + const args = + typeof input === 'string' + ? input + : input === undefined + ? '' + : JSON.stringify(input) + ctx.session.prompt(body.split('$ARGUMENTS').join(args).trim()) + return `Sent /${name}.` + }, + }) + } + return { commands } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/index.ts b/packages/ai-harness/src/first-party/index.ts new file mode 100644 index 0000000000..0d49df50cc --- /dev/null +++ b/packages/ai-harness/src/first-party/index.ts @@ -0,0 +1,19 @@ +// First-party plugins. This entry is Node only: some plugins use the file +// system and the shell. +export { modelPicker } from './model-picker' +export { + PERMISSION_MODES, + PermissionRules, + decidePermission, + permissions, +} from './permissions' +export type { + PermissionDecision, + PermissionMode, + PermissionRule, +} from './permissions' +export { globToRegExp, workspaceTools } from './workspace' +export { formatTodos, todos } from './todos' +export type { Todo } from './todos' +export { fileCommands, projectInstructions } from './files' +export { compact, usage } from './session-tools' diff --git a/packages/ai-harness/src/first-party/model-picker.ts b/packages/ai-harness/src/first-party/model-picker.ts new file mode 100644 index 0000000000..18807debfa --- /dev/null +++ b/packages/ai-harness/src/first-party/model-picker.ts @@ -0,0 +1,62 @@ +import { defineCommand } from '../commands' +import { configOption } from '../config' +import { definePlugin } from '../plugins' +import type { AnyTextAdapter } from '@tanstack/ai' + +/** + * Switch the main model at the next turn with the `model` setting or the + * `/model ` command. + * + * @example + * ```ts + * modelPicker({ + * choices: { fast: openaiText('gpt-5.6-luna'), smart: openaiText('gpt-5.6') }, + * default: 'smart', + * }) + * ``` + */ +export function modelPicker(options: { + choices: Record + default?: string +}) { + const names = Object.keys(options.choices) + const fallback = options.default ?? names[0] + if (fallback === undefined) + throw new Error('modelPicker needs at least one choice.') + return definePlugin({ + name: 'tanstack/model-picker', + setup: (ctx) => ({ + config: { + model: configOption.select({ + options: names, + default: fallback, + category: 'model', + description: 'The main model', + }), + }, + adapter: () => { + const name = ctx.config.get('model') + return typeof name === 'string' ? options.choices[name] : undefined + }, + commands: { + model: defineCommand({ + description: `Show or switch the model (${names.join(', ')})`, + run: async (input: unknown) => { + const name = + typeof input === 'string' + ? input + : typeof input === 'object' && input !== null && 'name' in input + ? String(input.name) + : undefined + if (!name) + return `Model: ${String(ctx.config.get('model'))}. Choices: ${names.join(', ')}.` + if (!names.includes(name)) + return `Unknown model "${name}". Choices: ${names.join(', ')}.` + await ctx.session.setConfig('model', name) + return `Model: ${name}. It applies at the next turn.` + }, + }), + }, + }), + }) +} diff --git a/packages/ai-harness/src/first-party/permissions.ts b/packages/ai-harness/src/first-party/permissions.ts new file mode 100644 index 0000000000..8f77eee0ed --- /dev/null +++ b/packages/ai-harness/src/first-party/permissions.ts @@ -0,0 +1,150 @@ +import { defineCommand } from '../commands' +import { configOption } from '../config' +import { createExtensionPoint } from '../extensions' +import { definePlugin } from '../plugins' + +export type PermissionDecision = 'allow' | 'ask' | 'deny' + +/** A rule for one tool. `kind` lets the modes treat edits and commands differently. */ +export interface PermissionRule { + /** A tool name. A trailing `*` matches a prefix, for example `git_*`. */ + tool: string + decision: PermissionDecision + kind?: 'read' | 'edit' | 'execute' +} + +/** Tool plugins add their rules here. `permissions()` reads them. */ +export const PermissionRules = createExtensionPoint( + 'tanstack/permission-rules', +) + +export const PERMISSION_MODES = [ + 'default', + 'plan', + 'acceptEdits', + 'bypass', +] as const +export type PermissionMode = (typeof PERMISSION_MODES)[number] + +function matches(rule: PermissionRule, tool: string): boolean { + return rule.tool.endsWith('*') + ? tool.startsWith(rule.tool.slice(0, -1)) + : rule.tool === tool +} + +/** The decision for a tool in a mode. The last matching rule wins. */ +export function decidePermission( + rules: ReadonlyArray, + tool: string, + mode: PermissionMode, + fallback: PermissionDecision = 'allow', +): PermissionDecision { + const rule = rules.findLast((candidate) => matches(candidate, tool)) + const decision = rule?.decision ?? fallback + if (mode === 'bypass') return 'allow' + if (mode === 'plan') { + return rule?.kind === 'edit' || + rule?.kind === 'execute' || + decision !== 'allow' + ? 'deny' + : 'allow' + } + if (mode === 'acceptEdits' && rule?.kind === 'edit' && decision === 'ask') + return 'allow' + return decision +} + +const PLAN_PROMPT = + 'You are in plan mode. Do not change files or run commands. Read what you need, then describe your plan.' + +/** + * Check every tool call against permission rules, with a `mode` setting and a + * `/mode` command: + * + * - `default`: rules apply as written. `ask` asks the user. + * - `plan`: read-only. Edits and commands are denied. + * - `acceptEdits`: edits run without asking. Commands still ask. + * - `bypass`: everything runs. + */ +export function permissions( + options: { + rules?: ReadonlyArray + default?: PermissionDecision + } = {}, +) { + return definePlugin({ + name: 'tanstack/permissions', + setup: (ctx) => { + const contributed = ctx.collect(PermissionRules) + const mode = (): PermissionMode => { + const value = ctx.config.get('mode') + return ( + PERMISSION_MODES.find((candidate) => candidate === value) ?? 'default' + ) + } + return { + config: { + mode: configOption.select({ + options: PERMISSION_MODES, + default: 'default', + category: 'mode', + description: 'How tool calls are approved', + }), + }, + prompts: [ + { + id: 'tanstack/permissions:plan', + text: () => (mode() === 'plan' ? PLAN_PROMPT : ''), + }, + ], + commands: { + mode: defineCommand({ + description: `Show or switch the mode (${PERMISSION_MODES.join(', ')})`, + run: async (input: unknown) => { + const next = typeof input === 'string' ? input.trim() : '' + if (!next) return `Mode: ${mode()}.` + const receipt = await ctx.session.setConfig('mode', next) + return receipt.status === 'rejected' + ? `Unknown mode "${next}". Modes: ${PERMISSION_MODES.join(', ')}.` + : `Mode: ${next}.` + }, + }), + }, + middleware: [ + { + name: 'tanstack/permissions', + onBeforeToolCall: async (_run, hook) => { + const rules = [...(options.rules ?? []), ...contributed] + const decision = decidePermission( + rules, + hook.toolName, + mode(), + options.default, + ) + if (decision === 'allow') return undefined + if (decision === 'ask') { + const preview = JSON.stringify(hook.args ?? {}).slice(0, 300) + const answer: unknown = await ctx.session.ask({ + message: `Allow ${hook.toolName} ${preview}? (y/n)`, + }) + const allowed = + answer === true || /^y(es)?$/i.test(String(answer).trim()) + if (allowed) return undefined + } + return { + type: 'skip', + result: { + denied: true, + reason: + decision === 'ask' + ? 'The user denied this tool call.' + : `This tool is not allowed in ${mode()} mode.`, + }, + } + }, + }, + ], + } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/session-tools.ts b/packages/ai-harness/src/first-party/session-tools.ts new file mode 100644 index 0000000000..02dbe6fb6c --- /dev/null +++ b/packages/ai-harness/src/first-party/session-tools.ts @@ -0,0 +1,115 @@ +import { chat } from '@tanstack/ai' +import { defineCommand } from '../commands' +import { definePlugin } from '../plugins' +import type { AnyTextAdapter, ModelMessage } from '@tanstack/ai' + +const SUMMARY_PROMPT = + 'Summarize the conversation so far for yourself. Keep decisions, open tasks, file names, and facts you still need. Leave out small talk.' + +function textOf(message: ModelMessage): string { + if (typeof message.content === 'string') return message.content + if (Array.isArray(message.content)) { + return message.content + .map((part) => + 'text' in part && typeof part.text === 'string' ? part.text : '', + ) + .join('') + } + return '' +} + +/** + * `/compact`: replace a long transcript with a summary, so later turns send + * fewer tokens. The summary is written by `adapter`. + */ +export function compact(options: { + adapter: AnyTextAdapter + keepLast?: number +}) { + return definePlugin({ + name: 'tanstack/compact', + setup: (ctx) => ({ + commands: { + compact: defineCommand({ + description: 'Summarize the conversation to save tokens', + run: async () => { + const messages = await ctx.session.transcript() + const keep = options.keepLast ?? 0 + if (messages.length <= keep + 2) + return 'The conversation is already short.' + const older = messages.slice(0, messages.length - keep) + const transcript = older + .map((message) => `${message.role}: ${textOf(message)}`) + .filter((line) => !line.endsWith(': ')) + .join('\n') + const summary = await chat({ + adapter: options.adapter, + messages: [ + { role: 'user', content: `${SUMMARY_PROMPT}\n\n${transcript}` }, + ], + stream: false, + }) + await ctx.session.replaceTranscript([ + { + role: 'user', + content: `Summary of our conversation so far:\n${summary}`, + }, + { + role: 'assistant', + content: 'Understood. I will continue from this summary.', + }, + ...messages.slice(messages.length - keep), + ]) + return `Compacted ${older.length} messages into a summary.` + }, + }), + }, + }), + }) +} + +interface UsageTotals { + turns: number + promptTokens: number + completionTokens: number + totalTokens: number +} + +/** Count tokens across the session. `/usage` shows the totals. */ +export function usage() { + return definePlugin({ + name: 'tanstack/usage', + setup: (ctx) => { + const state = ctx.state({ + turns: 0, + promptTokens: 0, + completionTokens: 0, + totalTokens: 0, + }) + const show = (totals: UsageTotals) => + `${totals.turns} model calls, ${totals.promptTokens} input tokens, ${totals.completionTokens} output tokens, ${totals.totalTokens} total.` + return { + middleware: [ + { + name: 'tanstack/usage', + onUsage: async (_run, info) => { + await state.update((totals) => ({ + turns: totals.turns + 1, + promptTokens: totals.promptTokens + (info.promptTokens ?? 0), + completionTokens: + totals.completionTokens + (info.completionTokens ?? 0), + totalTokens: totals.totalTokens + (info.totalTokens ?? 0), + })) + }, + }, + ], + commands: { + usage: defineCommand({ + description: 'Show token usage for this session', + run: async () => show(await state.get()), + }), + }, + } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/todos.ts b/packages/ai-harness/src/first-party/todos.ts new file mode 100644 index 0000000000..eae2c2cce8 --- /dev/null +++ b/packages/ai-harness/src/first-party/todos.ts @@ -0,0 +1,100 @@ +import { toolDefinition } from '@tanstack/ai' +import { defineCommand } from '../commands' +import { definePlugin } from '../plugins' + +export interface Todo { + text: string + status: 'pending' | 'in_progress' | 'done' +} + +function isTodo(value: unknown): value is Todo { + return ( + typeof value === 'object' && + value !== null && + 'text' in value && + typeof value.text === 'string' && + 'status' in value && + (value.status === 'pending' || + value.status === 'in_progress' || + value.status === 'done') + ) +} + +const MARK = { pending: '[ ]', in_progress: '[~]', done: '[x]' } as const + +export function formatTodos(items: ReadonlyArray): string { + return items.map((item) => `${MARK[item.status]} ${item.text}`).join('\n') +} + +/** + * A todo list the model keeps for multi-step work (like Claude Code's + * TodoWrite). The list is plugin state, so clients see it change and it + * survives restarts. `/todos` shows it. + */ +export function todos() { + return definePlugin({ + name: 'tanstack/todos', + setup: async (ctx) => { + const state = ctx.state<{ items: Array }>({ items: [] }) + // The prompt runs for each turn, so keep a copy of the list at hand. + let current = (await state.get()).items + return { + tools: [ + toolDefinition({ + name: 'todo_write', + description: + 'Replace the todo list. Use it for work with several steps. Mark one item in_progress at a time.', + inputSchema: { + type: 'object', + properties: { + todos: { + type: 'array', + items: { + type: 'object', + properties: { + text: { type: 'string' }, + status: { + type: 'string', + enum: ['pending', 'in_progress', 'done'], + }, + }, + required: ['text', 'status'], + }, + }, + }, + required: ['todos'], + }, + replay: 'safe', + }).server(async (args: unknown) => { + const list = + typeof args === 'object' && + args !== null && + 'todos' in args && + Array.isArray(args.todos) + ? args.todos + : [] + if (!list.every(isTodo)) + throw new Error('Each todo needs text and a status.') + current = (await state.update(() => ({ items: list }))).items + return formatTodos(current) || 'The todo list is empty.' + }), + ], + prompts: [ + { + id: 'tanstack/todos:list', + text: () => + current.length > 0 + ? `Current todo list:\n${formatTodos(current)}` + : '', + }, + ], + commands: { + todos: defineCommand({ + description: 'Show the todo list', + run: () => formatTodos(current) || 'The todo list is empty.', + }), + }, + } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/workspace.ts b/packages/ai-harness/src/first-party/workspace.ts new file mode 100644 index 0000000000..98d9cb13aa --- /dev/null +++ b/packages/ai-harness/src/first-party/workspace.ts @@ -0,0 +1,275 @@ +import { exec } from 'node:child_process' +import { readFile, readdir, stat, writeFile, mkdir } from 'node:fs/promises' +import { dirname, isAbsolute, relative, resolve, sep } from 'node:path' +import { toolDefinition } from '@tanstack/ai' +import { definePlugin } from '../plugins' +import { PermissionRules } from './permissions' + +const MAX_OUTPUT = 20_000 +const MAX_READ_LINES = 2000 +const SKIP_DIRS = new Set(['node_modules', '.git', 'dist', '.nx', '.turbo']) + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null +} + +function stringArg(args: unknown, key: string): string { + const value = isRecord(args) ? args[key] : undefined + if (typeof value === 'string') return value + throw new Error(`Argument "${key}" must be a string.`) +} + +function optionalString(args: unknown, key: string): string | undefined { + const value = isRecord(args) ? args[key] : undefined + return typeof value === 'string' ? value : undefined +} + +function clip(text: string): string { + return text.length > MAX_OUTPUT + ? `${text.slice(0, MAX_OUTPUT)}\n[${text.length - MAX_OUTPUT} more characters]` + : text +} + +/** A simple glob to RegExp: `**` any path, `*` any name part, `?` one character. */ +export function globToRegExp(glob: string): RegExp { + let pattern = '' + for (let index = 0; index < glob.length; index += 1) { + const char = glob.charAt(index) + if (char === '*' && glob[index + 1] === '*') { + pattern += '.*' + index += glob[index + 2] === '/' ? 2 : 1 + } else if (char === '*') pattern += '[^/]*' + else if (char === '?') pattern += '[^/]' + else pattern += char.replace(/[.+^${}()|[\]\\]/g, '\\$&') + } + return new RegExp(`^${pattern}$`) +} + +/** + * File and shell tools for a coding agent, confined to `root`: `read_file`, + * `write_file`, `edit_file`, `list_files`, `grep`, and `bash`. Edits and + * `bash` ask for approval through `permissions()`. + * + * These tools run on this machine with the host's authority. Use a sandbox + * for code you do not trust. + */ +export function workspaceTools(options: { + root: string + bashTimeoutMs?: number +}) { + const root = resolve(options.root) + + const inRoot = (path: string): string => { + const full = isAbsolute(path) ? resolve(path) : resolve(root, path) + if (full !== root && !full.startsWith(root + sep)) { + throw new Error(`Path "${path}" is outside the workspace.`) + } + return full + } + const shown = (full: string) => + relative(root, full).split(sep).join('/') || '.' + + async function walk(dir: string, out: Array): Promise { + for (const entry of await readdir(dir, { withFileTypes: true })) { + if (entry.isDirectory()) { + if (!SKIP_DIRS.has(entry.name)) + await walk(resolve(dir, entry.name), out) + } else out.push(resolve(dir, entry.name)) + if (out.length > 5000) return + } + } + + const tools = [ + toolDefinition({ + name: 'read_file', + description: + 'Read a text file in the workspace. Lines are numbered from 1.', + inputSchema: { + type: 'object', + properties: { + path: { type: 'string' }, + offset: { type: 'number', description: 'First line, from 1' }, + limit: { type: 'number' }, + }, + required: ['path'], + }, + replay: 'safe', + }).server(async (args: unknown) => { + const full = inRoot(stringArg(args, 'path')) + const lines = (await readFile(full, 'utf8')).split('\n') + const offset = + isRecord(args) && typeof args.offset === 'number' + ? Math.max(1, args.offset) + : 1 + const limit = + isRecord(args) && typeof args.limit === 'number' + ? args.limit + : MAX_READ_LINES + return clip( + lines + .slice(offset - 1, offset - 1 + limit) + .map((line, index) => `${offset + index}\t${line}`) + .join('\n'), + ) + }), + toolDefinition({ + name: 'write_file', + description: 'Create or replace a file in the workspace.', + inputSchema: { + type: 'object', + properties: { path: { type: 'string' }, content: { type: 'string' } }, + required: ['path', 'content'], + }, + }).server(async (args: unknown) => { + const full = inRoot(stringArg(args, 'path')) + await mkdir(dirname(full), { recursive: true }) + await writeFile(full, stringArg(args, 'content'), 'utf8') + return `Wrote ${shown(full)}.` + }), + toolDefinition({ + name: 'edit_file', + description: + 'Replace exact text in a file. `old` must appear once, unless `replaceAll` is true.', + inputSchema: { + type: 'object', + properties: { + path: { type: 'string' }, + old: { type: 'string' }, + new: { type: 'string' }, + replaceAll: { type: 'boolean' }, + }, + required: ['path', 'old', 'new'], + }, + }).server(async (args: unknown) => { + const full = inRoot(stringArg(args, 'path')) + const before = await readFile(full, 'utf8') + const oldText = stringArg(args, 'old') + const newText = stringArg(args, 'new') + const count = before.split(oldText).length - 1 + const replaceAll = isRecord(args) && args.replaceAll === true + if (count === 0) throw new Error('The old text is not in the file.') + if (count > 1 && !replaceAll) { + throw new Error( + `The old text appears ${count} times. Add context, or set replaceAll.`, + ) + } + await writeFile(full, before.split(oldText).join(newText), 'utf8') + return `Edited ${shown(full)} (${replaceAll ? count : 1} change${count > 1 && replaceAll ? 's' : ''}).` + }), + toolDefinition({ + name: 'list_files', + description: + 'List files in the workspace, optionally matching a glob like `src/**/*.ts`.', + inputSchema: { + type: 'object', + properties: { pattern: { type: 'string' } }, + }, + replay: 'safe', + }).server(async (args: unknown) => { + const pattern = optionalString(args, 'pattern') + const files: Array = [] + await walk(root, files) + const regex = pattern ? globToRegExp(pattern) : undefined + const listed = files + .map(shown) + .filter((file) => !regex || regex.test(file)) + return clip(listed.slice(0, 1000).join('\n') || 'No files.') + }), + toolDefinition({ + name: 'grep', + description: + 'Search file contents with a regular expression. Returns `file:line: text`.', + inputSchema: { + type: 'object', + properties: { pattern: { type: 'string' }, glob: { type: 'string' } }, + required: ['pattern'], + }, + replay: 'safe', + }).server(async (args: unknown) => { + const regex = new RegExp(stringArg(args, 'pattern')) + const glob = optionalString(args, 'glob') + const fileFilter = glob ? globToRegExp(glob) : undefined + const files: Array = [] + await walk(root, files) + const hits: Array = [] + for (const file of files) { + if (fileFilter && !fileFilter.test(shown(file))) continue + if ((await stat(file)).size > 1_000_000) continue + const lines = (await readFile(file, 'utf8')).split('\n') + lines.forEach((line, index) => { + if (hits.length < 200 && regex.test(line)) + hits.push(`${shown(file)}:${index + 1}: ${line.trim()}`) + }) + if (hits.length >= 200) break + } + return clip(hits.join('\n') || 'No matches.') + }), + toolDefinition({ + name: 'bash', + description: + 'Run a shell command in the workspace folder. Output is cut at 20,000 characters.', + inputSchema: { + type: 'object', + properties: { command: { type: 'string' } }, + required: ['command'], + }, + replay: 'never', + }).server( + (args: unknown) => + new Promise((done) => { + exec( + stringArg(args, 'command'), + { + cwd: root, + timeout: options.bashTimeoutMs ?? 120_000, + maxBuffer: 10 * 1024 * 1024, + }, + (error, stdout, stderr) => { + const code = error && 'code' in error ? error.code : 0 + done( + clip( + `exit code: ${String(code ?? 0)}\n${stdout}${stderr ? `\nstderr:\n${stderr}` : ''}`, + ), + ) + }, + ) + }), + ), + ] + + return definePlugin({ + name: 'tanstack/workspace-tools', + setup: () => ({ + tools, + prompts: [`Your workspace is ${root}. Paths are relative to it.`], + contribute: [ + PermissionRules.item({ + tool: 'read_file', + decision: 'allow', + kind: 'read', + }), + PermissionRules.item({ + tool: 'list_files', + decision: 'allow', + kind: 'read', + }), + PermissionRules.item({ tool: 'grep', decision: 'allow', kind: 'read' }), + PermissionRules.item({ + tool: 'write_file', + decision: 'ask', + kind: 'edit', + }), + PermissionRules.item({ + tool: 'edit_file', + decision: 'ask', + kind: 'edit', + }), + PermissionRules.item({ + tool: 'bash', + decision: 'ask', + kind: 'execute', + }), + ], + }), + }) +} diff --git a/packages/ai-harness/src/host.ts b/packages/ai-harness/src/host.ts index dd8dafb6c7..9fb84539b7 100644 --- a/packages/ai-harness/src/host.ts +++ b/packages/ai-harness/src/host.ts @@ -3,6 +3,7 @@ import { HarnessSession } from './session' import type { AIPersistence, ChatTranscriptStores, + CredentialStore, InboxStore, } from '@tanstack/ai-persistence' import type { AnyHarness } from './define' @@ -14,7 +15,7 @@ import type { Principal } from './types' * restart loses the ones not yet applied. */ export type HarnessPersistence = AIPersistence< - ChatTranscriptStores & { inbox?: InboxStore } + ChatTranscriptStores & { inbox?: InboxStore; credentials?: CredentialStore } > export interface HarnessHostOptions { @@ -68,6 +69,10 @@ export function createHarnessHost( // ponytail: a memory inbox when the stores have none. Pass `stores.inbox` // to keep accepted inputs across restarts. const inbox = persistence.stores.inbox ?? memoryPersistence().stores.inbox + // ponytail: memory credentials when the stores have none. Pass + // `stores.credentials` to keep sign-ins across restarts. + const credentials = + persistence.stores.credentials ?? memoryPersistence().stores.credentials const sessions = new Map>() const hostId = `host-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 9)}` @@ -81,6 +86,7 @@ export function createHarnessHost( threadId, persistence, inbox, + credentials, hostId, ...(principal ? { principal } : {}), onClose: () => sessions.delete(key), diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts index e1de83cca8..548768916e 100644 --- a/packages/ai-harness/src/index.ts +++ b/packages/ai-harness/src/index.ts @@ -15,8 +15,28 @@ export type { PluginLifetime, PluginPrompt, PluginSetupContext, + PluginState, } from './plugins' +export { createExtensionPoint, createPluginEvent } from './extensions' +export type { ExtensionItem, ExtensionPoint, PluginEvent } from './extensions' + +export { checkConfigValue, configOption } from './config' +export type { ConfigOption } from './config' + +export { defineCommand } from './commands' +export type { + AnswerOf, + AnyCommand, + CommandContext, + CommandDefinition, + PluginSessionApi, + Question, +} from './commands' + +export { AuthRequiredError, scrubSecrets } from './auth' +export type { CredentialsAccess } from './auth' + export type { AgentInputOf, AgentRegistryView, @@ -39,6 +59,7 @@ export type { AgentRunOptions, AgentStartOptions, DynamicAgentHandle, + SessionInspection, SessionSnapshot, } from './session' @@ -75,3 +96,17 @@ export type { export { harnessText } from './harness-text' export type { HarnessTextOptions } from './harness-text' + +export { + buildAuthorizationUrl, + createPkce, + deviceLogin, + exchangeCode, + isExpired, + loopbackLogin, + refreshCredential, +} from './oauth' +export type { OAuthConfig } from './oauth' + +export { oauthConnector } from './connectors' +export type { OAuthConnectorOptions } from './connectors' diff --git a/packages/ai-harness/src/oauth.ts b/packages/ai-harness/src/oauth.ts new file mode 100644 index 0000000000..8928322a76 --- /dev/null +++ b/packages/ai-harness/src/oauth.ts @@ -0,0 +1,298 @@ +import type { Credential } from '@tanstack/ai-persistence' + +/** An OAuth 2 app: where to send the user, and where to get tokens. */ +export interface OAuthConfig { + authorizationUrl: string + tokenUrl: string + /** The device authorization endpoint (RFC 8628), for logins without a browser here. */ + deviceUrl?: string + clientId: string + /** Only for confidential clients. A CLI is a public client and has none. */ + clientSecret?: string + scopes?: ReadonlyArray +} + +type Fetch = typeof fetch + +const encoder = new TextEncoder() + +function base64url(bytes: Uint8Array): string { + let binary = '' + for (const byte of bytes) binary += String.fromCharCode(byte) + return btoa(binary).replace(/\+/g, '-').replace(/\//g, '_').replace(/=+$/, '') +} + +function randomString(bytes = 32): string { + return base64url(crypto.getRandomValues(new Uint8Array(bytes))) +} + +/** A PKCE pair with the S256 method (RFC 7636). */ +export async function createPkce(): Promise<{ + verifier: string + challenge: string +}> { + const verifier = randomString(32) + const digest = await crypto.subtle.digest('SHA-256', encoder.encode(verifier)) + return { verifier, challenge: base64url(new Uint8Array(digest)) } +} + +/** The URL that starts a browser sign-in. */ +export function buildAuthorizationUrl( + config: OAuthConfig, + options: { redirectUri: string; state: string; challenge: string }, +): string { + const url = new URL(config.authorizationUrl) + url.searchParams.set('response_type', 'code') + url.searchParams.set('client_id', config.clientId) + url.searchParams.set('redirect_uri', options.redirectUri) + url.searchParams.set('state', options.state) + url.searchParams.set('code_challenge', options.challenge) + url.searchParams.set('code_challenge_method', 'S256') + if (config.scopes?.length) + url.searchParams.set('scope', config.scopes.join(' ')) + return url.toString() +} + +async function tokenRequest( + config: OAuthConfig, + body: Record, + doFetch: Fetch, +): Promise> { + const response = await doFetch(config.tokenUrl, { + method: 'POST', + headers: { + 'Content-Type': 'application/x-www-form-urlencoded', + Accept: 'application/json', + }, + body: new URLSearchParams({ + client_id: config.clientId, + ...(config.clientSecret ? { client_secret: config.clientSecret } : {}), + ...body, + }), + }) + const parsed: unknown = await response.json().catch(() => ({})) + if (typeof parsed !== 'object' || parsed === null) { + throw new Error(`OAuth token request failed (${response.status}).`) + } + // An OAuth token endpoint answers with a flat JSON object. + return parsed as Record +} + +function credentialFrom( + response: Record, + previous?: Credential, +): Credential { + if (typeof response.access_token !== 'string') { + const reason = + typeof response.error === 'string' ? response.error : 'no access_token' + throw new Error(`OAuth token request failed: ${reason}`) + } + const refreshToken = + typeof response.refresh_token === 'string' + ? response.refresh_token + : previous?.type === 'oauth' + ? previous.refreshToken + : undefined + return { + type: 'oauth', + accessToken: response.access_token, + ...(refreshToken ? { refreshToken } : {}), + ...(typeof response.expires_in === 'number' + ? { expiresAt: Date.now() + response.expires_in * 1000 } + : {}), + ...(typeof response.scope === 'string' + ? { scopes: response.scope.split(/[ ,]/).filter(Boolean) } + : {}), + } +} + +/** Trade an authorization code for tokens. */ +export async function exchangeCode( + config: OAuthConfig, + options: { + code: string + verifier: string + redirectUri: string + fetch?: Fetch + }, +): Promise { + const response = await tokenRequest( + config, + { + grant_type: 'authorization_code', + code: options.code, + code_verifier: options.verifier, + redirect_uri: options.redirectUri, + }, + options.fetch ?? fetch, + ) + return credentialFrom(response) +} + +/** True when an OAuth credential expires within `skewMs`. */ +export function isExpired(credential: Credential, skewMs = 60_000): boolean { + return ( + credential.type === 'oauth' && + credential.expiresAt !== undefined && + credential.expiresAt - skewMs <= Date.now() + ) +} + +/** Get a new access token with the refresh token. */ +export async function refreshCredential( + config: OAuthConfig, + credential: Credential, + doFetch: Fetch = fetch, +): Promise { + if (credential.type !== 'oauth' || !credential.refreshToken) { + throw new Error('This credential has no refresh token.') + } + const response = await tokenRequest( + config, + { grant_type: 'refresh_token', refresh_token: credential.refreshToken }, + doFetch, + ) + return credentialFrom(response, credential) +} + +const DONE_PAGE = + 'Signed in

You are signed in. You can close this tab.

' + +/** + * Sign in through the browser with a loopback redirect (RFC 8252 + PKCE). + * Listens on `127.0.0.1` on a random port, for one callback only. Calls + * `onUrl` with the URL to open. Resolves with the tokens. + */ +export async function loopbackLogin( + config: OAuthConfig, + options: { onUrl: (url: string) => void; fetch?: Fetch; timeoutMs?: number }, +): Promise { + const { createServer } = await import('node:http') + const { verifier, challenge } = await createPkce() + const state = randomString(16) + const server = createServer() + await new Promise((resolve, reject) => { + server.once('error', reject) + server.listen(0, '127.0.0.1', () => resolve()) + }) + const address = server.address() + const port = typeof address === 'object' && address ? address.port : 0 + const redirectUri = `http://127.0.0.1:${port}/callback` + + try { + const code = await new Promise((resolve, reject) => { + const timer = setTimeout( + () => reject(new Error('Sign-in timed out.')), + options.timeoutMs ?? 10 * 60_000, + ) + server.on('request', (req, res) => { + const url = new URL(req.url ?? '/', redirectUri) + if (url.pathname !== '/callback') { + res.writeHead(404).end() + return + } + clearTimeout(timer) + const returnedState = url.searchParams.get('state') + const returnedCode = url.searchParams.get('code') + res.writeHead(200, { 'Content-Type': 'text/html' }).end(DONE_PAGE) + if (returnedState !== state) + reject(new Error('Sign-in failed: the state does not match.')) + else if (!returnedCode) { + reject( + new Error( + `Sign-in failed: ${url.searchParams.get('error') ?? 'no code'}`, + ), + ) + } else resolve(returnedCode) + }) + options.onUrl( + buildAuthorizationUrl(config, { redirectUri, state, challenge }), + ) + }) + return await exchangeCode(config, { + code, + verifier, + redirectUri, + ...(options.fetch ? { fetch: options.fetch } : {}), + }) + } finally { + server.closeAllConnections() + server.close() + } +} + +/** + * Sign in with a device code (RFC 8628), for SSH sessions, containers, and + * CI. Calls `onCode` with the code and the page to enter it on, then polls. + */ +export async function deviceLogin( + config: OAuthConfig, + options: { + onCode: (info: { userCode: string; verificationUri: string }) => void + fetch?: Fetch + /** Test hook. Default waits the interval the server asks for. */ + sleep?: (ms: number) => Promise + }, +): Promise { + if (!config.deviceUrl) + throw new Error('This OAuth app has no device endpoint.') + const doFetch = options.fetch ?? fetch + const sleep = + options.sleep ?? + ((ms: number) => new Promise((resolve) => setTimeout(resolve, ms))) + const started = await doFetch(config.deviceUrl, { + method: 'POST', + headers: { + 'Content-Type': 'application/x-www-form-urlencoded', + Accept: 'application/json', + }, + body: new URLSearchParams({ + client_id: config.clientId, + ...(config.scopes?.length ? { scope: config.scopes.join(' ') } : {}), + }), + }) + const device: unknown = await started.json() + if ( + typeof device !== 'object' || + device === null || + !('device_code' in device) || + typeof device.device_code !== 'string' || + !('user_code' in device) || + typeof device.user_code !== 'string' + ) { + throw new Error('Device sign-in failed: the server sent no device code.') + } + const verificationUri = + 'verification_uri' in device && typeof device.verification_uri === 'string' + ? device.verification_uri + : config.authorizationUrl + let interval = + 'interval' in device && typeof device.interval === 'number' + ? device.interval * 1000 + : 5000 + const expiresAt = + Date.now() + + ('expires_in' in device && typeof device.expires_in === 'number' + ? device.expires_in * 1000 + : 15 * 60_000) + options.onCode({ userCode: device.user_code, verificationUri }) + + while (Date.now() < expiresAt) { + await sleep(interval) + const response = await tokenRequest( + config, + { + grant_type: 'urn:ietf:params:oauth:grant-type:device_code', + device_code: device.device_code, + }, + doFetch, + ) + if (response.error === 'authorization_pending') continue + if (response.error === 'slow_down') { + interval += 5000 + continue + } + return credentialFrom(response) + } + throw new Error('Device sign-in expired.') +} diff --git a/packages/ai-harness/src/plugins.ts b/packages/ai-harness/src/plugins.ts index 550ec092fa..347dad4ac9 100644 --- a/packages/ai-harness/src/plugins.ts +++ b/packages/ai-harness/src/plugins.ts @@ -4,12 +4,16 @@ import type { AgentProduces, AnyChatMiddleware, AnyGenerationMiddleware, + AnyTextAdapter, AnyTool, Capability, CapabilityHandle, - SystemPrompt, } from '@tanstack/ai' import type { AgentRegistry, AgentRegistryView, AnyAgent } from './agents' +import type { CredentialsAccess } from './auth' +import type { AnyCommand, PluginSessionApi } from './commands' +import type { ConfigOption } from './config' +import type { ExtensionItem, ExtensionPoint, PluginEvent } from './extensions' /** * How long a plugin's `setup` result and resources live: @@ -19,24 +23,55 @@ import type { AgentRegistry, AgentRegistryView, AnyAgent } from './agents' */ export type PluginLifetime = 'session' | 'run' -/** A prompt a plugin contributes. `id` must be unique in the session. */ +/** + * A prompt a plugin contributes. `id` must be unique in the session. A + * function `text` is called for each turn, so it can show current state. + */ export interface PluginPrompt { id: string - text: string + text: string | (() => string) } /** What a plugin's `setup` returns. Every field is optional. */ export interface PluginContributions { /** Tools for the main model. */ tools?: ReadonlyArray - /** System prompt text for every chat turn. */ - prompts?: ReadonlyArray + /** System prompt text for every chat turn. A function runs for each turn. */ + prompts?: ReadonlyArray string) | PluginPrompt> /** Chat middleware, the same type as `chat({ middleware })`. */ middleware?: ReadonlyArray /** Middleware for the activities agents call (`ctx.generateImage`, ...). */ generationMiddleware?: ReadonlyArray /** Agents added to `session.agents`. */ agents?: ReadonlyArray + /** User actions, keyed by name. Run with `session.command(name, input)`. */ + commands?: Record + /** Session settings, keyed by name. Read with `ctx.config.get(name)`. */ + config?: Record + /** Items for extension points that other plugins read. */ + contribute?: ReadonlyArray + /** + * Pick the main-loop adapter for the next turn, or return `undefined` to + * keep the harness adapter. The last plugin that returns one wins. + */ + adapter?: () => AnyTextAdapter | undefined +} + +/** Plugin state that survives restarts, stored in the metadata store. */ +export interface PluginState { + get: () => Promise + /** Change the state. On a write conflict, `change` runs again with fresh state. */ + update: (change: (current: T) => T) => Promise +} + +/** What the session gives plugins. */ +export interface PluginServices { + emit: (plugin: string, name: string, value: unknown) => void + on: (name: string, handler: (value: unknown) => void) => () => void + config: { get: (key: string) => unknown } + state: (plugin: string, initial: T) => PluginState + credentials: CredentialsAccess + session: PluginSessionApi } /** What a plugin's `setup` receives. */ @@ -57,7 +92,23 @@ export interface PluginSetupContext { provide: (capability: Capability, value: T) => void /** The agents of this session (harness agents plus earlier plugins' agents). */ agents: AgentRegistryView - session: { threadId: string } + /** + * The items other plugins contributed to `point`. The list fills while + * plugins set up, so read it at run time (in a tool, a command, or a + * middleware hook), not during `setup`. + */ + collect: (point: ExtensionPoint) => ReadonlyArray + /** Send a typed event to every plugin that listens, and to clients. */ + emit: (event: PluginEvent, value: T) => void + /** Listen for a typed event. Returns a function that stops listening. */ + on: (event: PluginEvent, handler: (value: T) => void) => () => void + /** This plugin's state, created with `initial` the first time. */ + state: (initial: T) => PluginState + /** Session settings. A change applies at the next turn. */ + config: { get: (key: string) => unknown } + /** Credentials of the session's principal. */ + credentials: CredentialsAccess + session: PluginSessionApi } export interface PluginDefinition { @@ -114,7 +165,24 @@ interface Owned { /** Everything a set of mounted plugins contributes, plus its cleanup. */ export interface MountedPlugins { tools: Array - prompts: Array + /** Strings, or functions to call for each turn. */ + prompts: Array string)> + commands: Map + config: Map + /** Contributions to extension points, by point name. */ + extensions: Map> + adapters: Array<() => AnyTextAdapter | undefined> + /** Who contributed what, for `session.inspect()`. */ + owners: { + plugins: Array<{ + name: string + lifetime: PluginLifetime + requires: Array + provides: Array + }> + tools: Array<{ name: string; owner: string }> + prompts: Array<{ id: string; owner: string }> + } middleware: Array generationMiddleware: Array /** @@ -137,6 +205,45 @@ export interface MountEnvironment { harnessProvides: ReadonlyArray /** Capability values from an outer mount (session plugins, for a run mount). */ inherited?: CapabilityValues + /** Names taken by an outer mount (session plugins, for a run mount). */ + takenCommands?: ReadonlyMap + takenConfig?: ReadonlyMap + /** Extension items from an outer mount, visible to `collect`. */ + inheritedExtensions?: ReadonlyMap< + string, + ReadonlyArray<{ value: unknown; owner: string }> + > + /** Session services. Tests of the mount alone can leave them out. */ + services?: PluginServices +} + +const unavailable = (what: string) => () => { + throw new Error(`${what} is only available inside a harness session.`) +} + +const NO_SERVICES: PluginServices = { + emit: () => {}, + on: () => () => {}, + config: { get: () => undefined }, + state: unavailable('ctx.state'), + credentials: { + get: unavailable('ctx.credentials'), + require: unavailable('ctx.credentials'), + set: unavailable('ctx.credentials'), + delete: unavailable('ctx.credentials'), + list: unavailable('ctx.credentials'), + }, + session: { + threadId: '', + principal: undefined, + snapshot: unavailable('ctx.session'), + prompt: unavailable('ctx.session'), + transcript: unavailable('ctx.session'), + replaceTranscript: unavailable('ctx.session'), + ask: unavailable('ctx.session'), + authRequired: unavailable('ctx.session'), + setConfig: unavailable('ctx.session'), + }, } /** Capability values provided by plugins, keyed by handle. */ @@ -198,8 +305,12 @@ function checkCapabilityOrder( } } -function promptOf(prompt: string | PluginPrompt, owner: string, index: number) { - return typeof prompt === 'string' +function promptOf( + prompt: string | (() => string) | PluginPrompt, + owner: string, + index: number, +): PluginPrompt { + return typeof prompt === 'string' || typeof prompt === 'function' ? { id: `${owner}#${index}`, text: prompt } : prompt } @@ -231,6 +342,14 @@ export async function mountPlugins( const prompts: Array> = [] const middleware: Array = [] const generationMiddleware: Array = [] + const commands = new Map() + const config = new Map() + const extensions = new Map>() + for (const [point, items] of env.inheritedExtensions ?? []) { + extensions.set(point, [...items]) + } + const adapters: Array<() => AnyTextAdapter | undefined> = [] + const services = env.services ?? NO_SERVICES try { for (const plugin of plugins) { @@ -271,7 +390,29 @@ export async function mountPlugins( provided.add(handle) }, agents: env.registry, - session: { threadId: env.threadId }, + collect: (point: ExtensionPoint): ReadonlyArray => { + let items = extensions.get(point.name) + if (!items) { + items = [] + extensions.set(point.name, items) + } + const source = items + // A live view: items contributed after this call appear too. + return new Proxy>([], { + get: (_target, key) => + Reflect.get( + source.map((item) => item.value as T), + key, + ), + }) + }, + emit: (event, value) => services.emit(plugin.name, event.name, value), + on: (event, handler) => + services.on(event.name, (value) => handler(value as never)), + state: (initial) => services.state(plugin.name, initial), + config: services.config, + credentials: services.credentials, + session: services.session, }) for (const handle of plugin.provides ?? []) { if (!provided.has(handle)) { @@ -302,6 +443,35 @@ export async function mountPlugins( }) middleware.push(...(contributions.middleware ?? [])) generationMiddleware.push(...(contributions.generationMiddleware ?? [])) + for (const [name, command] of Object.entries( + contributions.commands ?? {}, + )) { + const clash = commands.get(name) ?? env.takenCommands?.get(name) + if (clash) { + throw new Error( + `Duplicate command "${name}": first owner ${clash.owner}, second owner ${plugin.name}.`, + ) + } + commands.set(name, { command, owner: plugin.name }) + } + for (const [key, option] of Object.entries(contributions.config ?? {})) { + const clash = config.get(key) ?? env.takenConfig?.get(key) + if (clash) { + throw new Error( + `Duplicate config key "${key}": first owner ${clash.owner}, second owner ${plugin.name}.`, + ) + } + config.set(key, { option, owner: plugin.name }) + } + for (const item of contributions.contribute ?? []) { + let items = extensions.get(item.point) + if (!items) { + items = [] + extensions.set(item.point, items) + } + items.push({ value: item.value, owner: plugin.name }) + } + if (contributions.adapter) adapters.push(contributions.adapter) for (const agent of contributions.agents ?? []) { env.registry.add(agent, plugin.name) } @@ -350,6 +520,30 @@ export async function mountPlugins( prompts: prompts.map((entry) => entry.value.text), middleware, generationMiddleware, + commands, + config, + extensions, + adapters, + owners: { + plugins: plugins.map((plugin) => ({ + name: plugin.name, + lifetime: plugin.lifetime ?? 'session', + requires: (plugin.requires ?? []).map( + (handle) => handle.capabilityName, + ), + provides: (plugin.provides ?? []).map( + (handle) => handle.capabilityName, + ), + })), + tools: tools.map((entry) => ({ + name: entry.value.name, + owner: entry.owner, + })), + prompts: prompts.map((entry) => ({ + id: entry.value.id, + owner: entry.owner, + })), + }, capabilityBridge, values, dispose: () => (disposed ??= disposeAll(scopes)), diff --git a/packages/ai-harness/src/protocol.ts b/packages/ai-harness/src/protocol.ts index 3532454792..77a30fcf3d 100644 --- a/packages/ai-harness/src/protocol.ts +++ b/packages/ai-harness/src/protocol.ts @@ -41,6 +41,9 @@ const INPUT_OPS = new Set([ 'resolve', 'agent', 'cancel', + 'command', + 'answer', + 'config', ]) function isRecord(value: unknown): value is Record { @@ -71,6 +74,15 @@ export function parseHarnessInput(value: unknown): HarnessInput { if (value.op === 'agent' && typeof value.agent !== 'string') { throw new Error('Invalid input: agent needs an agent name.') } + if (value.op === 'command' && typeof value.name !== 'string') { + throw new Error('Invalid input: command needs a name.') + } + if (value.op === 'answer' && typeof value.questionId !== 'string') { + throw new Error('Invalid input: answer needs a questionId.') + } + if (value.op === 'config' && typeof value.key !== 'string') { + throw new Error('Invalid input: config needs a key.') + } // The checks above cover every field the session reads. return value as HarnessInput } @@ -134,6 +146,22 @@ export async function applyInput( return session.resolve(input.resume) case 'cancel': return session.cancel(input.operationId) + case 'command': { + const operation = session.command(input.name, input.input) + operation.then( + () => {}, + () => {}, + ) + return { + inputId: operation.id, + status: 'accepted', + operationId: operation.id, + } + } + case 'answer': + return session.answer(input.questionId, input.value) + case 'config': + return session.setConfig(input.key, input.value) case 'agent': { const exposed = (harness.expose?.agents ?? []).includes(input.agent) if (!exposed) { diff --git a/packages/ai-harness/src/session.ts b/packages/ai-harness/src/session.ts index 3e638d0f31..19c899c1e3 100644 --- a/packages/ai-harness/src/session.ts +++ b/packages/ai-harness/src/session.ts @@ -3,10 +3,13 @@ import { RUN_CANCEL_REASON, chat, compactForModel, + convertSchemaToJsonSchema, createSubagentId, runAgentStream, validateWithStandardSchema, } from '@tanstack/ai' +import { credentialsFor } from './auth' +import { checkConfigValue } from './config' import { withPersistence } from '@tanstack/ai-persistence' import { AgentRegistry } from './agents' import { SessionFeed } from './feed' @@ -23,10 +26,18 @@ import type { Interrupt, ModelMessage, RunAgentResumeItem, + SchemaInput, StreamChunk, SubagentBinding, } from '@tanstack/ai' -import type { InboxEntry, InboxStore } from '@tanstack/ai-persistence' +import type { + CredentialStore, + InboxEntry, + InboxStore, +} from '@tanstack/ai-persistence' +import type { CredentialsAccess } from './auth' +import type { PluginSessionApi, Question } from './commands' +import type { ConfigOption } from './config' import type { AgentInputOf, AgentRegistryView, @@ -35,7 +46,12 @@ import type { } from './agents' import type { AnyHarness, HarnessAgentsOf } from './define' import type { HarnessPersistence } from './host' -import type { HarnessPlugin, MountedPlugins } from './plugins' +import type { + HarnessPlugin, + MountedPlugins, + PluginServices, + PluginState, +} from './plugins' import type { BusyPolicy, ChatTurnResult, @@ -101,16 +117,34 @@ export interface SessionSnapshot { activeOperations: Array<{ id: string; kind: string; agent?: string }> queuedTurns: number pendingInterrupts: Array + /** Questions a command or a plugin asked, waiting for `session.answer`. */ + pendingQuestions: Array<{ + questionId: string + message: string + schema?: unknown + }> /** The cursor of the newest event. */ cursor: Cursor } +/** The resolved plugin plan of a session, for debugging and tooling. */ +export interface SessionInspection { + plugins: MountedPlugins['owners']['plugins'] + tools: MountedPlugins['owners']['tools'] + prompts: MountedPlugins['owners']['prompts'] + commands: Array<{ name: string; owner: string }> + config: Array<{ key: string; owner: string }> + extensionPoints: Record> + agents: Array +} + /** What the host hands a new session. */ export interface SessionDependencies { harness: AnyHarness threadId: string persistence: HarnessPersistence inbox: InboxStore + credentials: CredentialStore principal?: Principal /** Identifies this host on run leases. */ hostId: string @@ -187,6 +221,21 @@ export class HarnessSession { private closing: Promise | undefined private readonly onClose: () => void private readonly checkpoint: AnyChatMiddleware + private readonly listeners = new Map void>>() + private readonly configValues = new Map() + private readonly questions = new Map< + string, + { + message: string + schema: SchemaInput | undefined + resolve: (value: unknown) => void + reject: (error: unknown) => void + } + >() + private readonly stateDoc: Record = {} + private readonly localState = new Map() + private readonly credentialAccess: CredentialsAccess + private readonly services: PluginServices constructor(deps: SessionDependencies) { this.harness = deps.harness as THarness @@ -196,6 +245,37 @@ export class HarnessSession { this.principal = deps.principal this.onClose = deps.onClose this.checkpoint = checkpointMiddleware(deps.persistence, deps.hostId) + this.credentialAccess = credentialsFor( + deps.credentials, + { + threadId: deps.threadId, + ...(deps.principal ? { userId: deps.principal.id } : {}), + }, + (error) => + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.authRequired, { + connector: error.connector, + ...(error.url ? { url: error.url } : {}), + }), + ), + ) + this.services = { + emit: (plugin, name, value) => this.emitPluginEvent(plugin, name, value), + on: (name, handler) => { + let set = this.listeners.get(name) + if (!set) { + set = new Set() + this.listeners.set(name, set) + } + set.add(handler) + return () => set.delete(handler) + }, + config: { get: (key) => this.configValue(key) }, + state: (plugin, initial) => this.pluginState(plugin, initial), + credentials: this.credentialAccess, + session: this.pluginApi(), + } this.registry = this.agentRegistry for (const agent of this.harness.agents ?? []) { this.agentRegistry.add(agent, 'the harness') @@ -248,8 +328,10 @@ export class HarnessSession { harnessProvides: (this.harness.middleware ?? []).flatMap( (middleware) => middleware.provides ?? [], ), + services: this.services, }, ) + await this.loadConfig() await this.recoverCrashedTurn() await this.recoverInbox() } @@ -380,10 +462,361 @@ export class HarnessSession { })), queuedTurns: this.queue.length, pendingInterrupts: this.interrupted?.interrupts ?? [], + pendingQuestions: [...this.questions.entries()].map( + ([questionId, question]) => ({ + questionId, + message: question.message, + ...(question.schema + ? { schema: convertSchemaToJsonSchema(question.schema) } + : {}), + }), + ), cursor: this.feed.head(), } } + // =========================== + // Config, commands, questions + // =========================== + + /** Every session setting, with its option and current value. */ + config(): Record< + string, + { option: ConfigOption; value: unknown; owner: string } + > { + const result: Record< + string, + { option: ConfigOption; value: unknown; owner: string } + > = {} + for (const [key, entry] of this.sessionPlugins?.config ?? []) { + result[key] = { + option: entry.option, + owner: entry.owner, + value: this.configValue(key), + } + } + return result + } + + /** Change a session setting. It applies at the next turn. */ + async setConfig(key: string, value: unknown): Promise { + const inputId = createInputId() + await this.accept(inputId, { op: 'config', key, value }) + const entry = this.sessionPlugins?.config.get(key) + if (!entry) { + this.reject(inputId, 'unknown_config') + return { inputId, status: 'rejected', reason: 'unknown_config' } + } + let checked: unknown + try { + checked = checkConfigValue(key, entry.option, value) + } catch (error) { + const reason = error instanceof Error ? error.message : String(error) + this.reject(inputId, reason) + return { inputId, status: 'rejected', reason } + } + this.configValues.set(key, checked) + await this.persistence.stores.metadata?.set( + 'harness:config', + this.threadId, + Object.fromEntries(this.configValues), + ) + await this.applied(inputId, 'session') + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.configChanged, { key, value: checked }), + ) + return { inputId, status: 'accepted' } + } + + /** The commands of this session, for hosts to list. */ + commands(): Array<{ + name: string + description: string + owner: string + input?: unknown + }> { + return [...(this.sessionPlugins?.commands ?? [])].map(([name, entry]) => ({ + name, + description: entry.command.description, + owner: entry.owner, + ...(entry.command.input + ? { input: convertSchemaToJsonSchema(entry.command.input) } + : {}), + })) + } + + /** Run a plugin command. Its input is checked against the command's schema. */ + command(name: string, input?: unknown): Operation { + const operation = new OperationImpl( + 'command', + this.feed, + (target) => this.cancel(target.id), + ) + this.operations.set(operation.id, operation) + void this.executeCommand(operation, name, input) + return operation + } + + /** Answer a question from `ctx.session.ask`. */ + async answer(questionId: string, value: unknown): Promise { + const inputId = createInputId() + await this.accept(inputId, { op: 'answer', questionId, value }) + const question = this.questions.get(questionId) + if (!question) { + this.reject(inputId, 'unknown_question') + return { inputId, status: 'rejected', reason: 'unknown_question' } + } + let checked: unknown = value + if (question.schema !== undefined) { + const result = await validateWithStandardSchema(question.schema, value) + if (!result.success) { + const reason = `Invalid answer: ${result.issues.map((issue) => issue.message).join(', ')}` + this.reject(inputId, reason) + return { inputId, status: 'rejected', reason } + } + checked = result.data + } + this.questions.delete(questionId) + await this.applied(inputId, 'session') + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.questionAnswered, { questionId }), + ) + question.resolve(checked) + return { inputId, status: 'accepted' } + } + + /** The resolved plugin plan: order, owners, and extension contributors. */ + inspect(): SessionInspection { + const mounted = this.sessionPlugins + const extensionPoints: Record> = {} + for (const [point, items] of mounted?.extensions ?? []) { + extensionPoints[point] = [...new Set(items.map((item) => item.owner))] + } + return { + plugins: mounted?.owners.plugins ?? [], + tools: mounted?.owners.tools ?? [], + prompts: mounted?.owners.prompts ?? [], + commands: [...(mounted?.commands ?? [])].map(([name, entry]) => ({ + name, + owner: entry.owner, + })), + config: [...(mounted?.config ?? [])].map(([key, entry]) => ({ + key, + owner: entry.owner, + })), + extensionPoints, + agents: this.agentRegistry.list().map((agent) => agent.name), + } + } + + private configValue(key: string): unknown { + if (this.configValues.has(key)) return this.configValues.get(key) + return this.sessionPlugins?.config.get(key)?.option.default + } + + private async loadConfig(): Promise { + const stored = await this.persistence.stores.metadata?.get( + 'harness:config', + this.threadId, + ) + if (typeof stored !== 'object' || stored === null) return + for (const [key, value] of Object.entries(stored)) { + const entry = this.sessionPlugins?.config.get(key) + if (!entry) continue + try { + this.configValues.set(key, checkConfigValue(key, entry.option, value)) + } catch { + // A stored value an option no longer accepts falls back to the default. + } + } + } + + private ask(question: Question): Promise { + if (this.closing) return Promise.reject(new Error('Session closed.')) + const questionId = `q-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 9)}` + return new Promise((resolve, reject) => { + this.questions.set(questionId, { + message: question.message, + schema: question.schema, + resolve, + reject, + }) + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.question, { + questionId, + message: question.message, + ...(question.schema + ? { schema: convertSchemaToJsonSchema(question.schema) } + : {}), + }), + ) + }) + } + + private pluginApi(): PluginSessionApi { + return { + threadId: this.threadId, + principal: this.principal, + snapshot: () => this.snapshot(), + prompt: (text) => { + this.prompt(text).then( + () => {}, + () => {}, + ) + }, + transcript: async () => [ + ...(await this.persistence.stores.messages.loadThread(this.threadId)), + ], + replaceTranscript: (messages) => + this.persistence.stores.messages.saveThread(this.threadId, messages), + // The public type narrows the answer from the schema. + ask: ((question: Question) => + this.ask(question)) as PluginSessionApi['ask'], + authRequired: (info) => + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.authRequired, { ...info }), + ), + setConfig: (key, value) => this.setConfig(key, value), + } + } + + private emitPluginEvent(plugin: string, name: string, value: unknown): void { + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.pluginEvent, { plugin, name, value }), + ) + for (const handler of this.listeners.get(name) ?? []) { + try { + handler(value) + } catch { + // One broken listener must not stop the others. + } + } + } + + private pluginState(plugin: string, initial: T): PluginState { + const metadata = this.persistence.stores.metadata + const namespace = `plugin:${plugin}` + const key = this.threadId + const read = async (): Promise<{ value: T; revision: string | null }> => { + if (metadata?.getVersioned) { + const stored = await metadata.getVersioned(namespace, key) + // The store holds what this plugin wrote. + return stored + ? { value: stored.value as T, revision: stored.revision } + : { value: initial, revision: null } + } + if (metadata) { + const stored = await metadata.get(namespace, key) + return { + value: stored === null ? initial : (stored as T), + revision: null, + } + } + return { + value: this.localState.has(namespace) + ? (this.localState.get(namespace) as T) + : initial, + revision: null, + } + } + const publish = (value: T) => { + this.stateDoc[plugin] = value + this.feed.publish('session', { + type: EventType.STATE_SNAPSHOT, + snapshot: { plugins: { ...this.stateDoc } }, + timestamp: Date.now(), + }) + } + return { + get: async () => (await read()).value, + update: async (change) => { + for (let attempt = 0; attempt < 5; attempt += 1) { + const current = await read() + const next = change(structuredClone(current.value)) + if (metadata?.setIf) { + const written = await metadata.setIf( + namespace, + key, + next, + current.revision, + ) + if (!written.ok) continue + } else if (metadata) { + await metadata.set(namespace, key, next) + } else { + this.localState.set(namespace, next) + } + publish(next) + return next + } + throw new Error(`Plugin ${plugin}: state update conflicted 5 times.`) + }, + } + } + + private async executeCommand( + operation: OperationImpl, + name: string, + input: unknown, + ): Promise { + const inputId = createInputId() + await this.accept(inputId, { op: 'command', name, input }) + const entry = this.sessionPlugins?.commands.get(name) + if (!entry) { + this.reject(inputId, 'unknown_command') + operation.fail('failed', new Error(`Unknown command: ${name}`)) + return + } + let checked: unknown = input + if (entry.command.input !== undefined) { + const result = await validateWithStandardSchema( + entry.command.input, + input ?? {}, + ) + if (!result.success) { + const reason = `Input validation failed for command ${name}: ${result.issues + .map((issue) => issue.message) + .join(', ')}` + this.reject(inputId, 'invalid_input') + operation.fail('failed', new Error(reason)) + return + } + checked = result.data + } + operation.setStatus('running') + await this.applied(inputId, operation.id) + this.publishStarted(operation) + try { + const result: unknown = await entry.command.run(checked, { + signal: operation.abortController.signal, + session: this.services.session, + }) + operation.publish( + customEvent('harness.command.result', { + name, + result: compactForModel(result), + }), + ) + operation.finish('completed', result) + } catch (error) { + operation.publish({ + type: EventType.RUN_ERROR, + message: error instanceof Error ? error.message : String(error), + timestamp: Date.now(), + }) + operation.fail( + operation.abortController.signal.aborted ? 'cancelled' : 'failed', + error, + ) + } + this.publishFinished(operation) + } + /** * Stop every running operation, wait for them, then dispose session plugins. * Safe to call twice. @@ -399,6 +832,10 @@ export class HarnessSession { for (const operation of running) { operation.abortController.abort(RUN_CANCEL_REASON) } + for (const question of this.questions.values()) { + question.reject(new Error('Session closed.')) + } + this.questions.clear() await Promise.allSettled( running.map((operation) => Promise.resolve(operation)), ) @@ -518,8 +955,14 @@ export class HarnessSession { (middleware) => middleware.provides ?? [], ), ...(this.sessionPlugins - ? { inherited: this.sessionPlugins.values } + ? { + inherited: this.sessionPlugins.values, + takenCommands: this.sessionPlugins.commands, + takenConfig: this.sessionPlugins.config, + inheritedExtensions: this.sessionPlugins.extensions, + } : {}), + services: this.services, }) } const session = this.sessionPlugins @@ -528,16 +971,26 @@ export class HarnessSession { runPlugins?.capabilityBridge, ].filter((bridge): bridge is AnyChatMiddleware => bridge !== undefined) const subagents = this.harness.subagents + const picked = [ + ...(session?.adapters ?? []), + ...(runPlugins?.adapters ?? []), + ] + .map((pick) => pick()) + .filter((adapter) => adapter !== undefined) + .at(-1) + const resolvePrompt = (prompt: string | (() => string)) => + typeof prompt === 'function' ? prompt() : prompt const stream = chat({ - adapter: this.harness.adapter, + adapter: picked ?? this.harness.adapter, messages: turn.message !== undefined ? [{ id: createMessageId(), role: 'user', content: turn.message }] : [], systemPrompts: [ ...(this.harness.systemPrompts ?? []), - ...(session?.prompts ?? []), - ...(runPlugins?.prompts ?? []), + ...[...(session?.prompts ?? []), ...(runPlugins?.prompts ?? [])] + .map(resolvePrompt) + .filter((prompt) => prompt !== ''), ], tools: [ ...(this.harness.tools ?? []), diff --git a/packages/ai-harness/src/types.ts b/packages/ai-harness/src/types.ts index 43c65515b7..ffcf2efc89 100644 --- a/packages/ai-harness/src/types.ts +++ b/packages/ai-harness/src/types.ts @@ -54,6 +54,9 @@ export type HarnessInput = | { op: 'resolve'; resume: Array } | { op: 'agent'; agent: string; input?: unknown; detached?: boolean } | { op: 'cancel'; operationId?: string } + | { op: 'command'; name: string; input?: unknown } + | { op: 'answer'; questionId: string; value: unknown } + | { op: 'config'; key: string; value: unknown } /** Who sent an input, from the host's `authorize`. */ export interface Principal { @@ -94,6 +97,11 @@ export const HARNESS_EVENTS = { operationStarted: 'harness.operation.started', operationFinished: 'harness.operation.finished', operationResumed: 'harness.operation.resumed', + configChanged: 'harness.config.changed', + question: 'harness.question', + questionAnswered: 'harness.question.answered', + pluginEvent: 'harness.plugin.event', + authRequired: 'harness.auth_required', inputAccepted: 'harness.input.accepted', inputApplied: 'harness.input.applied', inputRejected: 'harness.input.rejected', diff --git a/packages/ai-harness/tests/auth.test.ts b/packages/ai-harness/tests/auth.test.ts new file mode 100644 index 0000000000..1a3509ea75 --- /dev/null +++ b/packages/ai-harness/tests/auth.test.ts @@ -0,0 +1,273 @@ +import { describe, expect, it, vi } from 'vitest' +import { toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + HARNESS_EVENTS, + buildAuthorizationUrl, + createHarnessHost, + createPkce, + defineHarness, + deviceLogin, + exchangeCode, + loopbackLogin, + oauthConnector, + refreshCredential, + scrubSecrets, +} from '../src' +import { mockAdapter, text, toolCall } from './helpers' +import type { OAuthConfig, SessionEvent } from '../src' + +const oauth: OAuthConfig = { + authorizationUrl: 'https://auth.example/authorize', + tokenUrl: 'https://auth.example/token', + deviceUrl: 'https://auth.example/device', + clientId: 'client-1', + scopes: ['repo', 'read:user'], +} + +/** A fake token endpoint that records each form body. */ +function tokenEndpoint(responses: Array>) { + const bodies: Array = [] + const fake: typeof fetch = async (_input, init) => { + bodies.push(new URLSearchParams(String(init?.body))) + return new Response(JSON.stringify(responses.shift() ?? {}), { + headers: { 'content-type': 'application/json' }, + }) + } + return { fake, bodies } +} + +describe('PKCE and URLs', () => { + it('makes an S256 challenge from the verifier', async () => { + const { verifier, challenge } = await createPkce() + const digest = await crypto.subtle.digest( + 'SHA-256', + new TextEncoder().encode(verifier), + ) + const expected = Buffer.from(digest).toString('base64url') + expect(challenge).toBe(expected) + expect(verifier.length).toBeGreaterThanOrEqual(43) + }) + + it('builds the authorization URL with state, PKCE, and scopes', () => { + const url = new URL( + buildAuthorizationUrl(oauth, { + redirectUri: 'http://127.0.0.1:5000/callback', + state: 's1', + challenge: 'c1', + }), + ) + expect(url.searchParams.get('client_id')).toBe('client-1') + expect(url.searchParams.get('code_challenge_method')).toBe('S256') + expect(url.searchParams.get('scope')).toBe('repo read:user') + expect(url.searchParams.get('state')).toBe('s1') + }) +}) + +describe('token exchange', () => { + it('trades a code for tokens and refreshes them', async () => { + const { fake, bodies } = tokenEndpoint([ + { + access_token: 'a1', + refresh_token: 'r1', + expires_in: 3600, + scope: 'repo', + }, + { access_token: 'a2', expires_in: 3600 }, + ]) + const first = await exchangeCode(oauth, { + code: 'code-1', + verifier: 'v1', + redirectUri: 'http://127.0.0.1/callback', + fetch: fake, + }) + expect(first).toMatchObject({ + type: 'oauth', + accessToken: 'a1', + refreshToken: 'r1', + scopes: ['repo'], + }) + expect(bodies[0]?.get('code_verifier')).toBe('v1') + + const refreshed = await refreshCredential(oauth, first, fake) + // The old refresh token stays when the server sends no new one. + expect(refreshed).toMatchObject({ accessToken: 'a2', refreshToken: 'r1' }) + expect(bodies[1]?.get('grant_type')).toBe('refresh_token') + }) + + it('reports an OAuth error', async () => { + const { fake } = tokenEndpoint([{ error: 'bad_verification_code' }]) + await expect( + exchangeCode(oauth, { + code: 'x', + verifier: 'y', + redirectUri: 'z', + fetch: fake, + }), + ).rejects.toThrow('bad_verification_code') + }) +}) + +describe('loopback login', () => { + it('receives the code on 127.0.0.1 and rejects a wrong state', async () => { + const { fake } = tokenEndpoint([{ access_token: 'loop-token' }]) + const signedIn = loopbackLogin(oauth, { + fetch: fake, + onUrl: (url) => { + const parsed = new URL(url) + const redirect = new URL(parsed.searchParams.get('redirect_uri')!) + expect(redirect.hostname).toBe('127.0.0.1') + redirect.searchParams.set('code', 'the-code') + redirect.searchParams.set('state', parsed.searchParams.get('state')!) + void fetch(redirect) + }, + }) + await expect(signedIn).resolves.toMatchObject({ accessToken: 'loop-token' }) + + const forged = loopbackLogin(oauth, { + fetch: fake, + onUrl: (url) => { + const redirect = new URL(new URL(url).searchParams.get('redirect_uri')!) + redirect.searchParams.set('code', 'x') + redirect.searchParams.set('state', 'forged') + void fetch(redirect) + }, + }) + await expect(forged).rejects.toThrow('state does not match') + }) +}) + +describe('device login', () => { + it('shows the code and polls until the user approves', async () => { + const responses: Array> = [ + { + device_code: 'd1', + user_code: 'ABCD-1234', + verification_uri: 'https://auth.example/device', + interval: 1, + }, + { error: 'authorization_pending' }, + { access_token: 'device-token' }, + ] + const fake: typeof fetch = async () => + new Response(JSON.stringify(responses.shift()), { + headers: { 'content-type': 'application/json' }, + }) + const onCode = vi.fn() + const credential = await deviceLogin(oauth, { + fetch: fake, + onCode, + sleep: async () => {}, + }) + expect(onCode).toHaveBeenCalledWith({ + userCode: 'ABCD-1234', + verificationUri: 'https://auth.example/device', + }) + expect(credential).toMatchObject({ accessToken: 'device-token' }) + }) +}) + +describe('oauthConnector', () => { + it('asks for sign-in, connects, and gives tools a token', async () => { + const seenTokens: Array = [] + const responses: Array> = [ + { + device_code: 'd1', + user_code: 'WXYZ', + verification_uri: 'https://gh.example/device', + interval: 1, + }, + { access_token: 'gh-token' }, + ] + const fake: typeof fetch = async () => + new Response(JSON.stringify(responses.shift()), { + headers: { 'content-type': 'application/json' }, + }) + const github = oauthConnector({ + id: 'github', + label: 'GitHub', + oauth, + login: 'device', + fetch: fake, + tools: (token) => [ + toolDefinition({ + name: 'list_issues', + description: 'List issues', + }).server(async () => { + seenTokens.push(await token()) + return ['#1'] + }), + ], + }) + const persistence = memoryPersistence() + const { adapter } = mockAdapter([ + () => toolCall('list_issues', {}), + () => text('Sign in first.'), + () => toolCall('list_issues', {}), + () => text('One issue.'), + ]) + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ + name: 'test/connector', + adapter, + plugins: () => [github], + }), + { threadId: 't', principal: { id: 'user-1' } }, + ) + const events: Array = [] + const reader = new AbortController() + const reading = (async () => { + for await (const entry of session.events({ signal: reader.signal })) + events.push(entry) + })() + + // Before sign-in, the tool asks for it. + await session.prompt('list my issues') + expect(seenTokens).toEqual([]) + expect( + events.some( + (entry) => + entry.event.type === 'CUSTOM' && + entry.event.name === HARNESS_EVENTS.authRequired && + (entry.event.value as { connector: string }).connector === 'github', + ), + ).toBe(true) + + await expect(session.command('connect:github')).resolves.toBe( + 'Connected to GitHub.', + ) + const saved = await persistence.stores.credentials.get( + { threadId: 't', userId: 'user-1' }, + 'github', + ) + expect(saved).toMatchObject({ accessToken: 'gh-token' }) + + await session.prompt('list my issues again') + expect(seenTokens).toEqual(['gh-token']) + // The model never sees the token. + const transcript = JSON.stringify( + await persistence.stores.messages.loadThread('t'), + ) + expect(transcript).not.toContain('gh-token') + + await session.command('disconnect:github') + expect( + await persistence.stores.credentials.list({ + threadId: 't', + userId: 'user-1', + }), + ).toEqual([]) + reader.abort() + await reading + await host.close() + }) +}) + +describe('scrubSecrets', () => { + it('removes known secrets from text', () => { + expect(scrubSecrets('token gh-abcdef1234 failed', ['gh-abcdef1234'])).toBe( + 'token [redacted] failed', + ) + }) +}) diff --git a/packages/ai-harness/tests/extensions.test.ts b/packages/ai-harness/tests/extensions.test.ts new file mode 100644 index 0000000000..4628ff8d6b --- /dev/null +++ b/packages/ai-harness/tests/extensions.test.ts @@ -0,0 +1,330 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { EventType } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + HARNESS_EVENTS, + configOption, + createExtensionPoint, + createHarnessHost, + createPluginEvent, + defineCommand, + defineHarness, + definePlugin, +} from '../src' +import { mockAdapter, text } from './helpers' +import type { StreamChunk } from '@tanstack/ai' + +describe('extension points', () => { + it('lets one plugin read what others contribute, in any order', async () => { + const Rules = createExtensionPoint<{ tool: string }>('test/rules') + let read: ReadonlyArray<{ tool: string }> = [] + const reader = definePlugin({ + name: 'test/reader', + setup: (ctx) => { + const rules = ctx.collect(Rules) + return { + commands: { + rules: defineCommand({ + description: 'List rules', + run: () => (read = [...rules]), + }), + }, + } + }, + }) + const writerA = definePlugin({ + name: 'test/a', + setup: () => ({ contribute: [Rules.item({ tool: 'write_file' })] }), + }) + const writerB = definePlugin({ + name: 'test/b', + setup: () => ({ contribute: [Rules.item({ tool: 'bash' })] }), + }) + const { adapter } = mockAdapter([]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/ext', + adapter, + plugins: () => [reader, writerA, writerB], + }), + { threadId: 't' }, + ) + await session.command('rules') + expect(read).toEqual([{ tool: 'write_file' }, { tool: 'bash' }]) + expect(session.inspect().extensionPoints['test/rules']).toEqual([ + 'test/a', + 'test/b', + ]) + await host.close() + }) +}) + +describe('events and state', () => { + it('delivers typed events and keeps plugin state', async () => { + const added = createPluginEvent<{ id: string }>('todo.added') + const heard: Array = [] + const todos = definePlugin({ + name: 'test/todos', + setup: (ctx) => { + const state = ctx.state<{ items: Array }>({ items: [] }) + return { + commands: { + add: defineCommand({ + description: 'Add a todo', + input: z.object({ text: z.string() }), + run: async ({ text: item }) => { + const next = await state.update((current) => ({ + items: [...current.items, item], + })) + ctx.emit(added, { id: item }) + return next.items.length + }, + }), + }, + } + }, + }) + const listener = definePlugin({ + name: 'test/listener', + setup: (ctx) => { + ctx.on(added, (value) => heard.push(value.id)) + }, + }) + const persistence = memoryPersistence() + const { adapter } = mockAdapter([]) + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ + name: 'test/state', + adapter, + plugins: () => [todos, listener], + }), + { threadId: 't' }, + ) + + const seen: Array = [] + const reader = new AbortController() + const reading = (async () => { + for await (const entry of session.events({ signal: reader.signal })) + seen.push(entry.event) + })() + + expect(await session.command('add', { text: 'milk' })).toBe(1) + expect(await session.command('add', { text: 'eggs' })).toBe(2) + await expect(session.command('add', { text: 3 })).rejects.toThrow( + 'Input validation failed', + ) + expect(heard).toEqual(['milk', 'eggs']) + expect( + await persistence.stores.metadata.get('plugin:test/todos', 't'), + ).toEqual({ + items: ['milk', 'eggs'], + }) + + await vi.waitFor(() => + expect( + seen.some((event) => event.type === EventType.STATE_SNAPSHOT), + ).toBe(true), + ) + reader.abort() + await reading + expect( + seen.some( + (event) => + event.type === EventType.CUSTOM && + event.name === HARNESS_EVENTS.pluginEvent, + ), + ).toBe(true) + await host.close() + }) +}) + +describe('session config', () => { + it('validates, persists, and applies settings at the next turn', async () => { + const seenThinking: Array = [] + const thinking = definePlugin({ + name: 'test/thinking', + setup: (ctx) => ({ + config: { + thinking: configOption.select({ + options: ['off', 'low', 'high'], + default: 'low', + }), + }, + middleware: [ + { + name: 'test/read-thinking', + onStart: () => void seenThinking.push(ctx.config.get('thinking')), + }, + ], + }), + }) + const persistence = memoryPersistence() + const { adapter } = mockAdapter([ + () => text('a'), + () => text('b'), + () => text('c'), + ]) + const harness = defineHarness({ + name: 'test/config', + adapter, + plugins: () => [thinking], + }) + const host = createHarnessHost({ persistence }) + const session = await host.open(harness, { threadId: 't' }) + + await session.prompt('one') + expect(await session.setConfig('thinking', 'extreme')).toMatchObject({ + status: 'rejected', + }) + expect(await session.setConfig('missing', 1)).toMatchObject({ + reason: 'unknown_config', + }) + expect(await session.setConfig('thinking', 'high')).toMatchObject({ + status: 'accepted', + }) + await session.prompt('two') + expect(seenThinking).toEqual(['low', 'high']) + expect(session.config().thinking).toMatchObject({ + value: 'high', + owner: 'test/thinking', + }) + await host.close() + + // A new host reads the saved setting. + const again = await createHarnessHost({ persistence }).open(harness, { + threadId: 't', + }) + expect(again.config().thinking?.value).toBe('high') + await again.close() + }) +}) + +describe('commands and questions', () => { + it('asks the user inside a command and continues with the answer', async () => { + const deploy = definePlugin({ + name: 'test/deploy', + setup: () => ({ + commands: { + deploy: defineCommand({ + description: 'Deploy after a confirmation', + run: async (_input, ctx) => { + const sure = await ctx.session.ask({ + message: 'Deploy to production?', + schema: z.boolean(), + }) + return sure ? 'deployed' : 'skipped' + }, + }), + }, + }), + }) + const { adapter } = mockAdapter([]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ name: 'test/ask', adapter, plugins: () => [deploy] }), + { threadId: 't' }, + ) + + const running = session.command('deploy') + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + const [question] = session.snapshot().pendingQuestions + expect(question?.message).toBe('Deploy to production?') + expect( + await session.answer(question!.questionId, 'yes please'), + ).toMatchObject({ + status: 'rejected', + }) + expect(await session.answer(question!.questionId, true)).toMatchObject({ + status: 'accepted', + }) + await expect(running).resolves.toBe('deployed') + expect(session.commands().map((command) => command.name)).toEqual([ + 'deploy', + ]) + await host.close() + }) + + it('rejects two plugins with the same command or config key', async () => { + const make = (name: string) => + definePlugin({ + name, + setup: () => ({ + commands: { same: defineCommand({ description: 'x', run: () => 1 }) }, + }), + }) + const { adapter } = mockAdapter([]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + await expect( + host.open( + defineHarness({ + name: 'test/dup-cmd', + adapter, + plugins: () => [make('test/one'), make('test/two')], + }), + { + threadId: 't', + }, + ), + ).rejects.toThrow( + 'Duplicate command "same": first owner test/one, second owner test/two', + ) + }) +}) + +describe('adapter override and credentials', () => { + it('lets a plugin pick the main adapter and read credentials', async () => { + const main = mockAdapter([() => text('from main')]) + const other = mockAdapter([() => text('from other')]) + let token: string | undefined + const picker = definePlugin({ + name: 'test/picker', + setup: (ctx) => ({ + adapter: () => + ctx.config.get('model') === 'other' ? other.adapter : undefined, + config: { + model: configOption.select({ + options: ['main', 'other'], + default: 'main', + }), + }, + commands: { + token: defineCommand({ + description: 'Read the GitHub token', + run: async () => { + const credential = await ctx.credentials.require('github') + token = + credential.type === 'oauth' + ? credential.accessToken + : credential.value + return 'ok' + }, + }), + }, + }), + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/pick', + adapter: main.adapter, + plugins: () => [picker], + }), + { threadId: 't', principal: { id: 'user-1' } }, + ) + + await session.setConfig('model', 'other') + expect(await session.prompt('hi')).toEqual({ text: 'from other' }) + expect(main.calls).toHaveLength(0) + + await expect(session.command('token')).rejects.toThrow( + 'Sign in to github first', + ) + expect(token).toBeUndefined() + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/first-party.test.ts b/packages/ai-harness/tests/first-party.test.ts new file mode 100644 index 0000000000..2718939d76 --- /dev/null +++ b/packages/ai-harness/tests/first-party.test.ts @@ -0,0 +1,218 @@ +import { mkdtemp, mkdir, readFile, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { + compact, + decidePermission, + fileCommands, + globToRegExp, + modelPicker, + permissions, + projectInstructions, + todos, + usage, + workspaceTools, +} from '../src/first-party' +import { messageTexts, mockAdapter, text, toolCall } from './helpers' +import type { HarnessPlugin } from '../src' +import type { AnyTextAdapter } from '@tanstack/ai' + +let root = '' +beforeEach(async () => { + root = await mkdtemp(join(tmpdir(), 'harness-')) +}) +afterEach(async () => { + await rm(root, { recursive: true, force: true }) +}) + +async function open(adapter: AnyTextAdapter, plugins: Array) { + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ + name: 'test/first-party', + adapter, + plugins: () => plugins, + }), + { + threadId: 't', + }, + ) + return { host, session, persistence } +} + +describe('decidePermission', () => { + const rules = [ + { tool: 'read_file', decision: 'allow' as const, kind: 'read' as const }, + { tool: 'write_file', decision: 'ask' as const, kind: 'edit' as const }, + { tool: 'bash', decision: 'ask' as const, kind: 'execute' as const }, + ] + it('applies modes on top of the rules', () => { + expect(decidePermission(rules, 'write_file', 'default')).toBe('ask') + expect(decidePermission(rules, 'write_file', 'acceptEdits')).toBe('allow') + expect(decidePermission(rules, 'bash', 'acceptEdits')).toBe('ask') + expect(decidePermission(rules, 'write_file', 'plan')).toBe('deny') + expect(decidePermission(rules, 'read_file', 'plan')).toBe('allow') + expect(decidePermission(rules, 'bash', 'bypass')).toBe('allow') + }) +}) + +describe('workspace tools with permissions', () => { + it('reads freely, asks before writing, and stays inside the root', async () => { + await writeFile(join(root, 'notes.txt'), 'hello\nworld\n') + const { adapter, calls } = mockAdapter([ + () => toolCall('read_file', { path: 'notes.txt' }, 'c1'), + () => + toolCall('write_file', { path: 'out.txt', content: 'written' }, 'c2'), + () => toolCall('read_file', { path: '../escape.txt' }, 'c3'), + () => text('done'), + ]) + const { host, session } = await open(adapter, [ + permissions(), + workspaceTools({ root }), + ]) + + const turn = session.prompt('work') + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + const [question] = session.snapshot().pendingQuestions + expect(question?.message).toContain('write_file') + await session.answer(question!.questionId, 'y') + await turn + + expect(await readFile(join(root, 'out.txt'), 'utf8')).toBe('written') + const toolResults = JSON.stringify(calls[3].messages) + expect(toolResults).toContain('1\\thello') + expect(toolResults).toContain('outside the workspace') + await host.close() + }) + + it('denies edits in plan mode without asking', async () => { + const { adapter } = mockAdapter([ + () => toolCall('write_file', { path: 'x.txt', content: 'x' }, 'c1'), + () => text('ok'), + ]) + const { host, session } = await open(adapter, [ + permissions(), + workspaceTools({ root }), + ]) + expect(await session.command('mode', 'plan')).toBe('Mode: plan.') + await session.prompt('write') + await expect(readFile(join(root, 'x.txt'), 'utf8')).rejects.toThrow() + expect(session.snapshot().pendingQuestions).toHaveLength(0) + await host.close() + }) + + it('matches globs like the list tool', () => { + expect(globToRegExp('src/**/*.ts').test('src/a/b.ts')).toBe(true) + expect(globToRegExp('src/**/*.ts').test('src/b.ts')).toBe(true) + expect(globToRegExp('*.md').test('docs/a.md')).toBe(false) + }) +}) + +describe('modelPicker', () => { + it('switches the main model with /model', async () => { + const fast = mockAdapter([() => text('fast answer')]) + const smart = mockAdapter([() => text('smart answer')]) + const { host, session } = await open(smart.adapter, [ + modelPicker({ + choices: { smart: smart.adapter, fast: fast.adapter }, + default: 'smart', + }), + ]) + expect(await session.command('model', 'fast')).toContain('Model: fast') + expect(await session.prompt('hi')).toEqual({ text: 'fast answer' }) + expect(await session.command('model', 'nope')).toContain('Unknown model') + await host.close() + }) +}) + +describe('todos', () => { + it('keeps the list in state and shows it in the next prompt', async () => { + const { adapter, calls } = mockAdapter([ + () => + toolCall('todo_write', { + todos: [{ text: 'write tests', status: 'in_progress' }], + }), + () => text('planned'), + () => text('next'), + ]) + const { host, session, persistence } = await open(adapter, [todos()]) + await session.prompt('plan it') + expect(await session.command('todos')).toBe('[~] write tests') + await session.prompt('continue') + expect(JSON.stringify(calls[2].systemPrompts)).toContain('[~] write tests') + expect( + await persistence.stores.metadata.get('plugin:tanstack/todos', 't'), + ).toEqual({ + items: [{ text: 'write tests', status: 'in_progress' }], + }) + await host.close() + }) +}) + +describe('project files', () => { + it('adds AGENTS.md to the prompt and turns command files into commands', async () => { + await writeFile(join(root, 'AGENTS.md'), 'Use tabs.') + await mkdir(join(root, 'commands')) + await writeFile( + join(root, 'commands', 'review.md'), + '---\ndescription: Review a file\n---\nReview $ARGUMENTS carefully.', + ) + const { adapter, calls } = mockAdapter([() => text('reviewed')]) + const { host, session } = await open(adapter, [ + projectInstructions({ root }), + fileCommands({ dir: join(root, 'commands') }), + ]) + expect(session.commands()).toEqual([ + { + name: 'review', + description: 'Review a file', + owner: 'tanstack/file-commands', + }, + ]) + await session.command('review', 'src/app.ts') + await vi.waitFor(() => expect(calls).toHaveLength(1)) + expect(messageTexts(calls[0])).toEqual(['Review src/app.ts carefully.']) + expect(JSON.stringify(calls[0].systemPrompts)).toContain('Use tabs.') + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + await host.close() + }) +}) + +describe('compact and usage', () => { + it('replaces the transcript with a summary and counts tokens', async () => { + const summarizer = mockAdapter([() => text('We planned a trip.')]) + // Real adapters report usage on RUN_FINISHED. + const withUsage = (answer: string) => () => + text(answer).map((chunk) => + chunk.type === 'RUN_FINISHED' + ? { + ...chunk, + usage: [{ inputTokens: 10, outputTokens: 5, totalTokens: 15 }], + } + : chunk, + ) + const main = mockAdapter([withUsage('one'), withUsage('two')]) + const { host, session, persistence } = await open(main.adapter, [ + compact({ adapter: summarizer.adapter }), + usage(), + ]) + await session.prompt('first') + await session.prompt('second') + expect(await session.command('compact')).toBe( + 'Compacted 4 messages into a summary.', + ) + const transcript = await persistence.stores.messages.loadThread('t') + expect(transcript).toHaveLength(2) + expect(transcript[0]?.content).toContain('We planned a trip.') + expect(await session.command('usage')).toBe( + '2 model calls, 20 input tokens, 10 output tokens, 30 total.', + ) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/helpers.ts b/packages/ai-harness/tests/helpers.ts index e8e3876082..ae2e3ed929 100644 --- a/packages/ai-harness/tests/helpers.ts +++ b/packages/ai-harness/tests/helpers.ts @@ -39,8 +39,12 @@ export function mockAdapter(replies: Array | Reply) { const now = () => Date.now() -/** One model call that answers with `content`. */ +let messageCount = 0 + +/** One model call that answers with `content`. Each call gets a new message id. */ export function text(content: string): Array { + messageCount += 1 + const m = `m-${messageCount}` return [ { type: EventType.RUN_STARTED, @@ -50,17 +54,17 @@ export function text(content: string): Array { }, { type: EventType.TEXT_MESSAGE_START, - messageId: 'm', + messageId: m, role: 'assistant', timestamp: now(), }, { type: EventType.TEXT_MESSAGE_CONTENT, - messageId: 'm', + messageId: m, delta: content, timestamp: now(), }, - { type: EventType.TEXT_MESSAGE_END, messageId: 'm', timestamp: now() }, + { type: EventType.TEXT_MESSAGE_END, messageId: m, timestamp: now() }, { type: EventType.RUN_FINISHED, runId: 'r', diff --git a/packages/ai-harness/vite.config.ts b/packages/ai-harness/vite.config.ts index 2fabdfea04..c99a69e684 100644 --- a/packages/ai-harness/vite.config.ts +++ b/packages/ai-harness/vite.config.ts @@ -28,7 +28,7 @@ const config = defineConfig({ export default mergeConfig( config, tanstackViteConfig({ - entry: ['./src/index.ts', './src/client.ts'], + entry: ['./src/index.ts', './src/client.ts', './src/first-party/index.ts'], srcDir: './src', cjs: false, }), diff --git a/packages/ai-persistence/src/index.ts b/packages/ai-persistence/src/index.ts index eac9972b34..23b7903f28 100644 --- a/packages/ai-persistence/src/index.ts +++ b/packages/ai-persistence/src/index.ts @@ -10,6 +10,7 @@ export { defineArtifactStore, defineBlobStore, defineInboxStore, + defineCredentialStore, // Run lifecycle helpers owned by @tanstack/ai, transiting through ./types so // this package's public surface stays a single import for backend authors. isTerminalRunStatus, @@ -30,6 +31,9 @@ export type { InboxEntry, InboxStatus, InboxStore, + // Harness credentials + Credential, + CredentialStore, // Named product shapes (prefer these over a sparse bag) ChatTranscriptStores, ChatPersistenceStores, diff --git a/packages/ai-persistence/src/memory.ts b/packages/ai-persistence/src/memory.ts index fd4df0b8b6..3565cbeb2a 100644 --- a/packages/ai-persistence/src/memory.ts +++ b/packages/ai-persistence/src/memory.ts @@ -14,8 +14,11 @@ import type { BlobStore, GenerationRunRecord, GenerationRunStore, + Credential, + CredentialStore, InboxEntry, InboxStore, + Scope, InterruptCommitEntry, InterruptRecord, InterruptStore, @@ -300,6 +303,47 @@ class MemoryMetadataStore implements MetadataStore { // (This parameter is an app-defined metadata namespace string — not the // shared `Scope` identity type from `@tanstack/ai`.) private readonly values = new Map>() + private readonly revisions = new Map>() + private revisionOf(namespace: string, key: string): number | undefined { + return this.revisions.get(namespace)?.get(key) + } + private bump(namespace: string, key: string): string { + let bucket = this.revisions.get(namespace) + if (!bucket) { + bucket = new Map() + this.revisions.set(namespace, bucket) + } + const next = (bucket.get(key) ?? 0) + 1 + bucket.set(key, next) + return String(next) + } + getVersioned( + namespace: string, + key: string, + ): Promise<{ value: unknown; revision: string } | null> { + const bucket = this.values.get(namespace) + if (!bucket || !bucket.has(key)) return Promise.resolve(null) + return Promise.resolve({ + value: bucket.get(key), + revision: String(this.revisionOf(namespace, key) ?? 0), + }) + } + async setIf( + namespace: string, + key: string, + value: unknown, + expectedRevision: string | null, + ): Promise< + { ok: true; revision: string } | { ok: false; reason: 'conflict' } + > { + const present = this.values.get(namespace)?.has(key) ?? false + const current = present + ? String(this.revisionOf(namespace, key) ?? 0) + : null + if (current !== expectedRevision) return { ok: false, reason: 'conflict' } + await this.set(namespace, key, value) + return { ok: true, revision: String(this.revisionOf(namespace, key)) } + } get(namespace: string, key: string): Promise { const bucket = this.values.get(namespace) if (!bucket || !bucket.has(key)) return Promise.resolve(null) @@ -312,11 +356,13 @@ class MemoryMetadataStore implements MetadataStore { this.values.set(namespace, bucket) } bucket.set(key, value) + this.bump(namespace, key) return Promise.resolve() } delete(namespace: string, key: string): Promise { const bucket = this.values.get(namespace) if (!bucket) return Promise.resolve() + this.revisions.get(namespace)?.delete(key) bucket.delete(key) if (bucket.size === 0) this.values.delete(namespace) return Promise.resolve() @@ -583,6 +629,48 @@ class MemoryInboxStore implements InboxStore { } } +class MemoryCredentialStore implements CredentialStore { + private readonly values = new Map>() + // A credential without a userId belongs to the tenant. + private owner(scope: Scope): string { + return JSON.stringify([scope.tenantId ?? null, scope.userId ?? null]) + } + get(scope: Scope, id: string): Promise { + const stored = this.values.get(this.owner(scope))?.get(id) + return Promise.resolve(stored ? { ...stored } : null) + } + set(scope: Scope, id: string, credential: Credential): Promise { + const key = this.owner(scope) + let bucket = this.values.get(key) + if (!bucket) { + bucket = new Map() + this.values.set(key, bucket) + } + bucket.set(id, { ...credential }) + return Promise.resolve() + } + delete(scope: Scope, id: string): Promise { + this.values.get(this.owner(scope))?.delete(id) + return Promise.resolve() + } + list( + scope: Scope, + ): Promise< + Array<{ id: string; type: Credential['type']; expiresAt?: number }> + > { + const bucket = this.values.get(this.owner(scope)) + return Promise.resolve( + [...(bucket?.entries() ?? [])].map(([id, credential]) => ({ + id, + type: credential.type, + ...(credential.type === 'oauth' && credential.expiresAt !== undefined + ? { expiresAt: credential.expiresAt } + : {}), + })), + ) + } +} + interface MemoryPersistenceStores { messages: MessageStore runs: RunStore @@ -592,13 +680,14 @@ interface MemoryPersistenceStores { artifacts: ArtifactStore blobs: BlobStore inbox: InboxStore + credentials: CredentialStore } /** * In-process reference backend for the full state + generation store set. * * Returns messages + runs + generationRuns + interrupts + metadata + artifacts - * + blobs + inbox. Locks are not included — use `InMemoryLockStore` + + * + blobs + inbox + credentials. Locks are not included — use `InMemoryLockStore` + * `withLocks` from `@tanstack/ai` when a test or single-process app needs * coordination. */ @@ -610,6 +699,7 @@ export function memoryPersistence() { interrupts: new MemoryInterruptStore(), metadata: new MemoryMetadataStore(), inbox: new MemoryInboxStore(), + credentials: new MemoryCredentialStore(), artifacts: new MemoryArtifactStore(), blobs: new MemoryBlobStore(), } diff --git a/packages/ai-persistence/src/testkit/conformance.ts b/packages/ai-persistence/src/testkit/conformance.ts index f621c7f6a2..27b6d0e874 100644 --- a/packages/ai-persistence/src/testkit/conformance.ts +++ b/packages/ai-persistence/src/testkit/conformance.ts @@ -1528,6 +1528,77 @@ export function runPersistenceConformance( }) }) + // Credentials are opt-in: only harness hosts read them. + describe('credentials', () => { + it('keys by user and tenant, lists without values, and deletes', async (ctx) => { + const store = persistence.stores.credentials + if (!store) return ctx.skip('credentials store not provided') + + const user = { threadId: 't', userId: 'u1', tenantId: 'org' } + const tenant = { threadId: 't', tenantId: 'org' } + await store.set(user, 'github', { + type: 'oauth', + accessToken: 'secret-access', + refreshToken: 'secret-refresh', + expiresAt: 100, + }) + await store.set(tenant, 'notion', { + type: 'api_key', + value: 'secret-key', + }) + + expect(await store.get(user, 'github')).toMatchObject({ + accessToken: 'secret-access', + }) + // A user credential is not visible at tenant scope, and the other way around. + expect(await store.get(tenant, 'github')).toBeNull() + expect(await store.get(user, 'notion')).toBeNull() + + const listed = await store.list(user) + expect(listed).toEqual([ + { id: 'github', type: 'oauth', expiresAt: 100 }, + ]) + // `list` never returns secret values. + expect(JSON.stringify(listed)).not.toContain('secret') + + await store.delete(user, 'github') + expect(await store.get(user, 'github')).toBeNull() + expect(await store.list(user)).toEqual([]) + }) + }) + + // Compare-and-set on metadata is opt-in. + describe('metadata setIf', () => { + it('writes only when the revision matches', async (ctx) => { + const store = persistence.stores.metadata + if (!store?.setIf || !store.getVersioned) { + return ctx.skip('metadata setIf not provided') + } + const created = await store.setIf('cas', 'k', { n: 1 }, null) + expect(created.ok).toBe(true) + expect(await store.setIf('cas', 'k', { n: 9 }, null)).toEqual({ + ok: false, + reason: 'conflict', + }) + const current = await store.getVersioned('cas', 'k') + expect(current?.value).toEqual({ n: 1 }) + const revision = current?.revision ?? null + const updated = await store.setIf( + 'cas', + 'k', + { n: 2 }, + revision, + ) + expect(updated.ok).toBe(true) + expect( + await store.setIf('cas', 'k', { n: 3 }, revision), + ).toMatchObject({ + ok: false, + }) + expect(await store.get('cas', 'k')).toEqual({ n: 2 }) + }) + }) + // The inbox is opt-in: only harness hosts read it. A backend without it // skips these cases and needs no `skip` entry. describe('inbox', () => { diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index 986c384bf0..3785eb4240 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -415,6 +415,39 @@ export interface InboxStore { export function defineInboxStore(store: InboxStore): InboxStore { return store } + +/** A secret a user or an organization saved: an API key or OAuth tokens. */ +export type Credential = + | { type: 'api_key'; value: string } + | { + type: 'oauth' + accessToken: string + refreshToken?: string + expiresAt?: number + scopes?: Array + } + +/** + * Durable store for credentials, keyed by scope and credential id (for + * example `'github'`). A credential saved without `scope.userId` belongs to + * the tenant. Encrypt at rest in your implementation. + */ +export interface CredentialStore { + get: (scope: Scope, id: string) => Promise + set: (scope: Scope, id: string, credential: Credential) => Promise + delete: (scope: Scope, id: string) => Promise + /** Ids and types only. `list` never returns secret values. */ + list: ( + scope: Scope, + ) => Promise< + Array<{ id: string; type: Credential['type']; expiresAt?: number }> + > +} + +/** Type a {@link CredentialStore} implementation inline. */ +export function defineCredentialStore(store: CredentialStore): CredentialStore { + return store +} /** Type a {@link GenerationRunStore} implementation inline. */ export function defineGenerationRunStore( store: GenerationRunStore, @@ -669,6 +702,8 @@ export interface AIPersistenceStores { blobs?: BlobStore /** Harness session inputs. Optional: only harness hosts read it. */ inbox?: InboxStore + /** User and tenant credentials. Optional: only harness hosts read it. */ + credentials?: CredentialStore } /** @@ -845,6 +880,7 @@ const storeKeys = [ 'artifacts', 'blobs', 'inbox', + 'credentials', ] satisfies Array const storeKeySet = new Set(storeKeys) diff --git a/packages/ai-persistence/tests/memory.test.ts b/packages/ai-persistence/tests/memory.test.ts index db479ec83b..bf9aa03e61 100644 --- a/packages/ai-persistence/tests/memory.test.ts +++ b/packages/ai-persistence/tests/memory.test.ts @@ -15,6 +15,7 @@ describe('memoryPersistence', () => { expect(Object.keys(memoryPersistence().stores).sort()).toEqual([ 'artifacts', 'blobs', + 'credentials', 'generationRuns', 'inbox', 'interrupts', diff --git a/packages/ai-persistence/tests/persistence-types.test-d.ts b/packages/ai-persistence/tests/persistence-types.test-d.ts index 467e17a3f0..d77673e7e6 100644 --- a/packages/ai-persistence/tests/persistence-types.test-d.ts +++ b/packages/ai-persistence/tests/persistence-types.test-d.ts @@ -22,6 +22,7 @@ import type { ChatTranscriptStores, GenerationRunStore, InboxStore, + CredentialStore, InterruptStore, MessagePage, MessageStore, @@ -141,6 +142,7 @@ expectTypeOf(memoryPersistence()).toEqualTypeOf< artifacts: ArtifactStore blobs: BlobStore inbox: InboxStore + credentials: CredentialStore }> >() const transcript: ChatTranscriptPersistence = messagesOnly diff --git a/packages/ai/src/activities/chat/middleware/metadata.ts b/packages/ai/src/activities/chat/middleware/metadata.ts index 08d71987f1..eb37219235 100644 --- a/packages/ai/src/activities/chat/middleware/metadata.ts +++ b/packages/ai/src/activities/chat/middleware/metadata.ts @@ -13,6 +13,27 @@ export interface MetadataStore { set: (namespace: string, key: string, value: unknown) => Promise /** Delete `(namespace, key)`. Do nothing when it is absent. */ delete: (namespace: string, key: string) => Promise + /** + * Optional. The value with a revision, or `null` when it is absent. Pair it + * with `setIf` for compare-and-set writes. + */ + getVersioned?: ( + namespace: string, + key: string, + ) => Promise<{ value: unknown; revision: string } | null> + /** + * Optional. Write only when the current revision is `expectedRevision` + * (`null` means "only when absent"). Returns the new revision, or + * `{ ok: false, reason: 'conflict' }` when someone else wrote first. + */ + setIf?: ( + namespace: string, + key: string, + value: unknown, + expectedRevision: string | null, + ) => Promise< + { ok: true; revision: string } | { ok: false; reason: 'conflict' } + > } export const MetadataCapability = createCapability()('metadata') diff --git a/testing/e2e/src/routes/api.harness-protocol.$.ts b/testing/e2e/src/routes/api.harness-protocol.$.ts index 50eddad947..727e618a92 100644 --- a/testing/e2e/src/routes/api.harness-protocol.$.ts +++ b/testing/e2e/src/routes/api.harness-protocol.$.ts @@ -1,10 +1,14 @@ import { createFileRoute } from '@tanstack/react-router' import { defineAgent } from '@tanstack/ai' import { + configOption, createHarnessHandler, createHarnessHost, + defineCommand, defineHarness, + definePlugin, } from '@tanstack/ai-harness' +import { todos } from '@tanstack/ai-harness/plugins' import { memoryPersistence } from '@tanstack/ai-persistence' import { z } from 'zod' import { createTextAdapter } from '@/lib/providers' @@ -36,6 +40,26 @@ function handlerFor(request: Request) { }), ], expose: { agents: ['echo'] }, + plugins: () => [ + todos(), + definePlugin({ + name: 'e2e/settings', + setup: () => ({ + config: { + tone: configOption.select({ + options: ['plain', 'warm'], + default: 'plain', + }), + }, + commands: { + greet: defineCommand({ + description: 'Say hello', + run: () => 'hello', + }), + }, + }), + }), + ], }) return createHarnessHandler({ host, diff --git a/testing/e2e/tests/harness-protocol.spec.ts b/testing/e2e/tests/harness-protocol.spec.ts index 43c81f3d47..7078085e49 100644 --- a/testing/e2e/tests/harness-protocol.spec.ts +++ b/testing/e2e/tests/harness-protocol.spec.ts @@ -61,6 +61,35 @@ test.describe('harness protocol', () => { expect(events.at(-1).name).toBe('harness.operation.finished') }) + test('changes a plugin setting and runs a plugin command', async ({ + request, + testId, + aimockPort, + }) => { + const control = async (input: unknown) => + ( + await request.post('/api/harness-protocol/control', { + headers: { + ...headers(testId, aimockPort), + 'content-type': 'application/json', + }, + data: { threadId: `plugins-${testId}`, input }, + }) + ).json() + + expect( + await control({ op: 'config', key: 'tone', value: 'warm' }), + ).toMatchObject({ + status: 'accepted', + }) + const rejected = await control({ op: 'config', key: 'tone', value: 'loud' }) + expect(rejected.status).toBe('rejected') + expect(rejected.reason).toContain('plain, warm') + const command = await control({ op: 'command', name: 'greet' }) + expect(command.status).toBe('accepted') + expect(command.operationId).toMatch(/^op-command-/) + }) + test('takes a control input and returns a receipt', async ({ request, testId, From c0f5f9a8422dba25c6862be74997a36ddb1d9de8 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 11:07:34 +0000 Subject: [PATCH 009/286] ci: apply automated fixes --- .../ai-persistence/src/testkit/conformance.ts | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/packages/ai-persistence/src/testkit/conformance.ts b/packages/ai-persistence/src/testkit/conformance.ts index 27b6d0e874..3b8ab62ebf 100644 --- a/packages/ai-persistence/src/testkit/conformance.ts +++ b/packages/ai-persistence/src/testkit/conformance.ts @@ -1583,18 +1583,13 @@ export function runPersistenceConformance( const current = await store.getVersioned('cas', 'k') expect(current?.value).toEqual({ n: 1 }) const revision = current?.revision ?? null - const updated = await store.setIf( - 'cas', - 'k', - { n: 2 }, - revision, - ) + const updated = await store.setIf('cas', 'k', { n: 2 }, revision) expect(updated.ok).toBe(true) - expect( - await store.setIf('cas', 'k', { n: 3 }, revision), - ).toMatchObject({ - ok: false, - }) + expect(await store.setIf('cas', 'k', { n: 3 }, revision)).toMatchObject( + { + ok: false, + }, + ) expect(await store.get('cas', 'k')).toEqual({ n: 2 }) }) }) From 372784798f2df2d338db7e0af2de9dfff6e7ef94 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 13:15:21 +0200 Subject: [PATCH 010/286] feat(ai): subagent tree limits; harness ctx.agents and harness children @tanstack/ai: subagents.limits (maxDepth, maxConcurrent, maxCalls, timeoutMs) for the tree of children the model starts through tools. One SubagentBudget per root run is shared by every child and passed on through ctx.chat({ subagents }), so a child cannot reset it. A refused start reaches the model as a tool error. @tanstack/ai-harness: plugins get ctx.agents.run, start, and group (with cancel-siblings or collect). harnessAgent(harness) turns a harness into a child agent for subagents.agents, and defineHarness takes a description. A harness applies default limits (depth 2, 3 at once, 12 per tree), and agents started from code count against them. --- .changeset/harness-p4-subagents.md | 11 ++ docs/chat/subagents.md | 22 +++ docs/config.json | 5 + docs/harness/subagents.md | 108 ++++++++++++ packages/ai-harness/src/define.ts | 2 + packages/ai-harness/src/harness-agent.ts | 36 ++++ packages/ai-harness/src/index.ts | 5 + packages/ai-harness/src/plugins.ts | 86 ++++++++- packages/ai-harness/src/session.ts | 105 ++++++++++- packages/ai-harness/tests/subagents.test.ts | 165 ++++++++++++++++++ .../ai/src/activities/chat/agents/bound.ts | 18 ++ .../ai/src/activities/chat/agents/limits.ts | 75 ++++++++ .../ai/src/activities/chat/agents/spawn.ts | 44 ++++- packages/ai/src/activities/index.ts | 1 + packages/ai/src/index.ts | 4 + packages/ai/tests/subagent-limits.test.ts | 159 +++++++++++++++++ testing/e2e/fixtures/harness/basic.json | 37 ++++ testing/e2e/src/routes/api.harness-test.ts | 24 +++ testing/e2e/tests/harness.spec.ts | 14 ++ 19 files changed, 909 insertions(+), 12 deletions(-) create mode 100644 .changeset/harness-p4-subagents.md create mode 100644 docs/harness/subagents.md create mode 100644 packages/ai-harness/src/harness-agent.ts create mode 100644 packages/ai-harness/tests/subagents.test.ts create mode 100644 packages/ai/src/activities/chat/agents/limits.ts create mode 100644 packages/ai/tests/subagent-limits.test.ts diff --git a/.changeset/harness-p4-subagents.md b/.changeset/harness-p4-subagents.md new file mode 100644 index 0000000000..3f64256a15 --- /dev/null +++ b/.changeset/harness-p4-subagents.md @@ -0,0 +1,11 @@ +--- +'@tanstack/ai': minor +'@tanstack/ai-harness': minor +--- + +Subagent trees now have limits, and a harness can run agents from code and call other harnesses. + +- **`subagents.limits`** in `@tanstack/ai`: `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.agents`** in harness plugins: `run`, `start` (with `wake`), and `group` with `onFailure: 'cancel-siblings' | 'collect'`. Every child of a group settles before the group returns. +- **`harnessAgent(harness)`** turns a harness into an agent for `subagents.agents`. `defineHarness` takes an optional `description`. +- A harness applies default limits (`DEFAULT_SUBAGENT_LIMITS`: depth 2, 3 at once, 12 per tree) when `subagents.limits` is not set, and agents started from code count against them. diff --git a/docs/chat/subagents.md b/docs/chat/subagents.md index 25a05c6ad2..7916787a8f 100644 --- a/docs/chat/subagents.md +++ b/docs/chat/subagents.md @@ -257,6 +257,28 @@ const researcher = defineAgent({ }) ``` +## 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-5.6'), + 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 may grow. The first chat is depth 0. +- `maxCalls`: how many children the whole tree may start. +- `maxConcurrent`: how many children one run may have running at once. +- `timeoutMs`: how long one child may 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)`. + ## 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 784c711a32..c1c32488d3 100644 --- a/docs/config.json +++ b/docs/config.json @@ -838,6 +838,11 @@ "label": "Auth and connectors", "to": "harness/auth", "addedAt": "2026-09-26" + }, + { + "label": "Run agents from a harness", + "to": "harness/subagents", + "addedAt": "2026-09-26" } ] }, diff --git a/docs/harness/subagents.md b/docs/harness/subagents.md new file mode 100644 index 0000000000..71b89b77ea --- /dev/null +++ b/docs/harness/subagents.md @@ -0,0 +1,108 @@ +--- +title: Run agents from a harness +id: harness-subagents +order: 8 +description: "Start typed agents from commands and plugins, run them in groups, call a whole harness as a child, and keep the tree within limits." +keywords: + - tanstack ai + - harness + - subagents + - agents + - limits +--- + +A review command needs an explorer to collect facts, then a reviewer to judge them. A research task needs three explorers at once. In a harness, plugins and commands start agents from code, and a whole harness can be a child of another one. + +## Run agents from a plugin + +`ctx.agents` runs any agent of the session, or an agent you pass in. Input and result types come from the agent. + +```ts group=harness-subagents +import { defineAgent } from '@tanstack/ai' +import { defineCommand, definePlugin } from '@tanstack/ai-harness' +import { openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' + +const explorer = defineAgent({ + name: 'explorer', + description: 'Collects facts about a topic', + inputSchema: z.object({ task: z.string() }), + run: (ctx) => + ctx.chat({ + adapter: openaiText('gpt-5.6-luna'), + messages: [{ role: 'user', content: ctx.input.task }], + stream: false, + }), +}) + +export const review = definePlugin({ + name: 'acme/review', + setup: (ctx) => ({ + agents: [explorer], + commands: { + review: defineCommand({ + description: 'Review a topic', + input: z.object({ topic: z.string() }), + run: async ({ topic }) => { + const [auth, sessions] = await ctx.agents.group({ onFailure: 'cancel-siblings' }, (group) => + Promise.all([ + group.run(explorer, { task: `${topic}: auth checks` }), + group.run(explorer, { task: `${topic}: session expiry` }), + ]), + ) + return `${auth}\n\n${sessions}` + }, + }), + }, + }), +}) +``` + +- `ctx.agents.run(agent, input)` runs one agent and resolves with its result. +- `ctx.agents.start(agent, input, { wake: true })` runs it in the background and starts a turn when it is done. +- `ctx.agents.group(options, body)` runs several. With `onFailure: 'cancel-siblings'`, one failure cancels the others. With `'collect'`, use `group.runSettled` to get every result or error. Every child settles before `group` returns. + +## Call a harness as a child + +`harnessAgent` turns a harness into an agent. Put it in `subagents.agents`, and the main model calls it as a tool. The child harness keeps its own tools, plugins, and history. + +```ts group=harness-subagents +import { defineHarness, harnessAgent } from '@tanstack/ai-harness' + +const reviewer = defineHarness({ + name: 'acme/reviewer', + description: 'Reviews a change and lists the risks', + adapter: openaiText('gpt-5.6'), +}) + +export const lead = defineHarness({ + name: 'acme/lead', + adapter: openaiText('gpt-5.6'), + subagents: { agents: [harnessAgent(reviewer)] }, +}) +``` + +The tool name comes from the harness name, with characters other than letters, digits, `_`, and `-` changed to `_`. Pass `name` to pick another. + +## Stay within limits + +A harness always limits its children. Without `subagents.limits`, it uses a depth of 2, 3 children at once, and 12 children per tree. Agents started from code count against the same limits: + +```ts group=harness-subagents +export const careful = defineHarness({ + name: 'acme/careful', + adapter: openaiText('gpt-5.6'), + subagents: { + agents: [explorer], + limits: { maxDepth: 1, maxConcurrent: 2, maxCalls: 6, timeoutMs: 60_000 }, + }, +}) +``` + +See [subagent limits](../chat/subagents) for what each limit does. + +## What you have now + +- Commands and plugins that start typed agents, alone or in groups. +- Harnesses that call other harnesses as tools. +- A tree of children that stays within limits. diff --git a/packages/ai-harness/src/define.ts b/packages/ai-harness/src/define.ts index ca22eafa31..aef7787bd7 100644 --- a/packages/ai-harness/src/define.ts +++ b/packages/ai-harness/src/define.ts @@ -28,6 +28,8 @@ export interface HarnessConfig< > { /** A stable name, for example `'acme/studio'`. */ name: string + /** What the harness does. Shown when another agent can call it. */ + description?: string /** The main agent-loop model, the same as `chat({ adapter })`. */ adapter: TAdapter systemPrompts?: Array diff --git a/packages/ai-harness/src/harness-agent.ts b/packages/ai-harness/src/harness-agent.ts new file mode 100644 index 0000000000..fbb7dddc58 --- /dev/null +++ b/packages/ai-harness/src/harness-agent.ts @@ -0,0 +1,36 @@ +import { defineAgent } from '@tanstack/ai' +import { harnessText } from './harness-text' +import type { AnyHarness } from './define' +import type { HarnessHost } from './host' + +/** + * Use a whole harness as a child agent: put it in `subagents.agents` and the + * main model can call it as a tool, with its own tools, plugins, and history. + * + * @example + * ```ts + * defineHarness({ + * name: 'acme/lead', + * adapter, + * subagents: { agents: [harnessAgent(reviewer)] }, + * }) + * ``` + */ +export function harnessAgent( + harness: AnyHarness, + options: { host?: HarnessHost; name?: string; description?: string } = {}, +) { + const adapter = harnessText( + harness, + options.host ? { host: options.host } : {}, + ) + return defineAgent({ + // Tool names allow letters, digits, `_`, and `-`. + name: options.name ?? harness.name.replace(/[^a-zA-Z0-9_-]/g, '_'), + description: + options.description ?? + harness.description ?? + `The ${harness.name} harness`, + run: (ctx) => ctx.chat({ adapter }), + }) +} diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts index 548768916e..3fda086bd6 100644 --- a/packages/ai-harness/src/index.ts +++ b/packages/ai-harness/src/index.ts @@ -14,6 +14,8 @@ export type { PluginDefinition, PluginLifetime, PluginPrompt, + AgentGroup, + PluginAgentActions, PluginSetupContext, PluginState, } from './plugins' @@ -110,3 +112,6 @@ export type { OAuthConfig } from './oauth' export { oauthConnector } from './connectors' export type { OAuthConnectorOptions } from './connectors' + +export { harnessAgent } from './harness-agent' +export { DEFAULT_SUBAGENT_LIMITS } from './session' diff --git a/packages/ai-harness/src/plugins.ts b/packages/ai-harness/src/plugins.ts index 347dad4ac9..d708655a19 100644 --- a/packages/ai-harness/src/plugins.ts +++ b/packages/ai-harness/src/plugins.ts @@ -9,7 +9,14 @@ import type { Capability, CapabilityHandle, } from '@tanstack/ai' -import type { AgentRegistry, AgentRegistryView, AnyAgent } from './agents' +import type { + AgentInputOf, + AgentRegistry, + AgentRegistryView, + AgentResultOf, + AnyAgent, +} from './agents' +import type { Operation } from './types' import type { CredentialsAccess } from './auth' import type { AnyCommand, PluginSessionApi } from './commands' import type { ConfigOption } from './config' @@ -64,6 +71,61 @@ export interface PluginState { update: (change: (current: T) => T) => Promise } +/** Run agents from a plugin: in the foreground, in the background, or as a group. */ +export interface PluginAgentActions { + run: { + ( + agent: TAgent, + input?: AgentInputOf, + ): Operation> + (name: string, input?: unknown): Operation + } + start: { + ( + agent: TAgent, + input?: AgentInputOf, + options?: { wake?: boolean }, + ): Operation> + ( + name: string, + input?: unknown, + options?: { wake?: boolean }, + ): Operation + } + /** + * Run children together. With `onFailure: 'cancel-siblings'` (default), one + * failure cancels the others. Every child settles before `group` returns. + */ + group: ( + options: { onFailure?: 'cancel-siblings' | 'collect' }, + body: (group: AgentGroup) => Promise, + ) => Promise +} + +/** The children of one `ctx.agents.group` call. */ +export interface AgentGroup { + run: { + ( + agent: TAgent, + input?: AgentInputOf, + ): Promise> + (name: string, input?: unknown): Promise + } + /** Like `run`, but never rejects: resolves with the result or the error. */ + runSettled: { + ( + agent: TAgent, + input?: AgentInputOf, + ): Promise< + { ok: true; value: AgentResultOf } | { ok: false; error: unknown } + > + ( + name: string, + input?: unknown, + ): Promise<{ ok: true; value: unknown } | { ok: false; error: unknown }> + } +} + /** What the session gives plugins. */ export interface PluginServices { emit: (plugin: string, name: string, value: unknown) => void @@ -72,6 +134,7 @@ export interface PluginServices { state: (plugin: string, initial: T) => PluginState credentials: CredentialsAccess session: PluginSessionApi + agents: PluginAgentActions } /** What a plugin's `setup` receives. */ @@ -90,8 +153,11 @@ export interface PluginSetupContext { getOptional: (capability: Capability) => T | undefined /** Provide a capability this plugin declared in `provides`. */ provide: (capability: Capability, value: T) => void - /** The agents of this session (harness agents plus earlier plugins' agents). */ - agents: AgentRegistryView + /** + * The agents of this session (harness agents plus earlier plugins' + * agents): find them, and run them from commands, tools, and hooks. + */ + agents: AgentRegistryView & PluginAgentActions /** * The items other plugins contributed to `point`. The list fills while * plugins set up, so read it at run time (in a tool, a command, or a @@ -244,6 +310,11 @@ const NO_SERVICES: PluginServices = { authRequired: unavailable('ctx.session'), setConfig: unavailable('ctx.session'), }, + agents: { + run: unavailable('ctx.agents.run') as PluginAgentActions['run'], + start: unavailable('ctx.agents.start') as PluginAgentActions['start'], + group: unavailable('ctx.agents.group'), + }, } /** Capability values provided by plugins, keyed by handle. */ @@ -389,7 +460,14 @@ export async function mountPlugins( values.provide(handle, value) provided.add(handle) }, - agents: env.registry, + agents: { + list: () => env.registry.list(), + get: (name) => env.registry.get(name), + find: (query) => env.registry.find(query), + run: services.agents.run, + start: services.agents.start, + group: services.agents.group, + }, collect: (point: ExtensionPoint): ReadonlyArray => { let items = extensions.get(point.name) if (!items) { diff --git a/packages/ai-harness/src/session.ts b/packages/ai-harness/src/session.ts index 19c899c1e3..df0e852c08 100644 --- a/packages/ai-harness/src/session.ts +++ b/packages/ai-harness/src/session.ts @@ -1,6 +1,7 @@ import { EventType, RUN_CANCEL_REASON, + SubagentBudget, chat, compactForModel, convertSchemaToJsonSchema, @@ -47,8 +48,10 @@ import type { import type { AnyHarness, HarnessAgentsOf } from './define' import type { HarnessPersistence } from './host' import type { + AgentGroup, HarnessPlugin, MountedPlugins, + PluginAgentActions, PluginServices, PluginState, } from './plugins' @@ -159,6 +162,13 @@ interface QueuedTurn { inputId?: string } +/** Limits for a harness's children when `subagents.limits` is not set. */ +export const DEFAULT_SUBAGENT_LIMITS = { + maxDepth: 2, + maxConcurrent: 3, + maxCalls: 12, +} + function createInputId(): string { return `in-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 9)}` } @@ -275,6 +285,23 @@ export class HarnessSession { state: (plugin, initial) => this.pluginState(plugin, initial), credentials: this.credentialAccess, session: this.pluginApi(), + agents: { + run: ((target: string | AnyAgent, input?: unknown) => + this.runAgent(target, input, { + wake: false, + })) as PluginAgentActions['run'], + start: (( + target: string | AnyAgent, + input?: unknown, + options?: AgentStartOptions, + ) => + this.runAgent( + target, + input, + options ?? {}, + )) as PluginAgentActions['start'], + group: (options, body) => this.agentGroup(options, body), + }, } this.registry = this.agentRegistry for (const agent of this.harness.agents ?? []) { @@ -924,6 +951,16 @@ export class HarnessSession { } } + /** + * A tree budget for a child started from code. The child counts as the + * first call, and its own children count against the same limits. + */ + private codeBudget(): SubagentBudget { + const root = SubagentBudget.root(this.limits()) + root.reserve(0) + return root.child() + } + private binding(): SubagentBinding { return { generationMiddleware: this.sessionPlugins?.generationMiddleware ?? [], @@ -1007,7 +1044,13 @@ export class HarnessSession { this.steering(), ], ...(subagents - ? { subagents: { ...subagents, binding: this.binding() } } + ? { + subagents: { + ...subagents, + limits: this.limits(), + binding: this.binding(), + }, + } : {}), ...(this.harness.agentLoopStrategy ? { agentLoopStrategy: this.harness.agentLoopStrategy } @@ -1083,28 +1126,75 @@ export class HarnessSession { // Agents // =========================== + /** The limits for children of this session, with the harness defaults. */ + private limits() { + return this.harness.subagents?.limits ?? DEFAULT_SUBAGENT_LIMITS + } + private runAgent( - name: string, + target: string | AnyAgent, input: unknown, options: AgentStartOptions, ): OperationImpl { + const name = typeof target === 'string' ? target : target.name const operation = new OperationImpl( 'agent', this.feed, - (target) => this.cancel(target.id), + (running) => this.cancel(running.id), name, ) this.operations.set(operation.id, operation) - void this.executeAgent(operation, name, input, options) + void this.executeAgent(operation, target, input, options) return operation } + /** Run a group of agents. Every child settles before the group returns. */ + private async agentGroup( + options: { onFailure?: 'cancel-siblings' | 'collect' }, + body: (group: AgentGroup) => Promise, + ): Promise { + const started: Array> = [] + const cancelOthers = (failed: OperationImpl) => { + if (options.onFailure === 'collect') return + for (const operation of started) { + if (operation !== failed && !operation.isSettled()) + void operation.cancel() + } + } + const group: AgentGroup = { + run: (target: string | AnyAgent, input?: unknown) => { + const operation = this.runAgent(target, input, { wake: false }) + started.push(operation) + return Promise.resolve(operation).catch((error: unknown) => { + cancelOthers(operation) + throw error + }) + }, + runSettled: (target: string | AnyAgent, input?: unknown) => { + const operation = this.runAgent(target, input, { wake: false }) + started.push(operation) + return Promise.resolve(operation).then( + (value) => ({ ok: true as const, value }), + (error: unknown) => ({ ok: false as const, error }), + ) + }, + } as AgentGroup + try { + return await body(group) + } finally { + await Promise.allSettled( + started.map((operation) => Promise.resolve(operation)), + ) + } + } + private async executeAgent( operation: OperationImpl, - name: string, + target: string | AnyAgent, input: unknown, options: AgentStartOptions, ): Promise { + const name = typeof target === 'string' ? target : target.name const inputId = createInputId() await this.accept(inputId, { op: 'agent', @@ -1112,7 +1202,8 @@ export class HarnessSession { input, ...(options.wake ? { detached: true } : {}), }) - const agent: AnyAgent | undefined = this.agentRegistry.get(name) + const agent: AnyAgent | undefined = + typeof target === 'string' ? this.agentRegistry.get(target) : target if (!agent) { this.reject(inputId, 'unknown_agent') operation.fail('failed', new Error(`Unknown agent: ${name}`)) @@ -1176,7 +1267,7 @@ export class HarnessSession { }, undefined, undefined, - this.binding(), + { ...this.binding(), budget: this.codeBudget() }, ) for await (const chunk of stream) { operation.publish(chunk) diff --git a/packages/ai-harness/tests/subagents.test.ts b/packages/ai-harness/tests/subagents.test.ts new file mode 100644 index 0000000000..f4cc99a04d --- /dev/null +++ b/packages/ai-harness/tests/subagents.test.ts @@ -0,0 +1,165 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { defineAgent } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + DEFAULT_SUBAGENT_LIMITS, + createHarnessHost, + defineCommand, + defineHarness, + definePlugin, + harnessAgent, +} from '../src' +import { mockAdapter, text, toolCall } from './helpers' + +describe('a harness as a child agent', () => { + it('lets the main model call another harness as a tool', async () => { + const inner = mockAdapter([() => text('Reviewed: looks good.')]) + const reviewer = defineHarness({ + name: 'acme/reviewer', + description: 'Reviews code', + adapter: inner.adapter, + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const child = harnessAgent(reviewer, { host }) + expect(child.name).toBe('acme_reviewer') + expect(child.description).toBe('Reviews code') + + const outer = mockAdapter([ + () => toolCall('acme_reviewer', {}), + () => text('The reviewer approved.'), + ]) + const lead = defineHarness({ + name: 'acme/lead', + adapter: outer.adapter, + subagents: { agents: [child] }, + }) + const session = await host.open(lead, { threadId: 't' }) + const turn = await session.prompt('review my change') + expect(turn.text).toBe('The reviewer approved.') + expect(JSON.stringify(outer.calls[1].messages)).toContain( + 'Reviewed: looks good.', + ) + await host.close() + }) +}) + +describe('ctx.agents in plugins', () => { + const echo = defineAgent({ + name: 'echo', + description: 'Echoes', + inputSchema: z.object({ text: z.string() }), + run: async (ctx) => ctx.input.text.toUpperCase(), + }) + const failing = defineAgent({ + name: 'failing', + description: 'Fails', + run: async () => { + throw new Error('boom') + }, + }) + + it('runs agents one by one and as a group from a command', async () => { + let slowStopped = false + const slow = defineAgent({ + name: 'slow', + description: 'Waits until stopped', + run: (ctx) => + new Promise((resolve) => { + ctx.abortSignal?.addEventListener('abort', () => { + slowStopped = true + resolve('stopped') + }) + }), + }) + const orchestrator = definePlugin({ + name: 'test/orchestrator', + setup: (ctx) => ({ + commands: { + one: defineCommand({ + description: 'Run echo', + run: async () => ctx.agents.run(echo, { text: 'hi' }), + }), + both: defineCommand({ + description: 'Run two echoes', + run: () => + ctx.agents.group({}, (group) => + Promise.all([ + group.run(echo, { text: 'a' }), + group.run('echo', { text: 'b' }), + ]), + ), + }), + failFast: defineCommand({ + description: 'One fails, the other is cancelled', + run: () => + ctx.agents.group({ onFailure: 'cancel-siblings' }, (group) => + Promise.all([group.run(slow), group.run(failing)]), + ), + }), + collect: defineCommand({ + description: 'Collect every outcome', + run: () => + ctx.agents.group({ onFailure: 'collect' }, (group) => + Promise.all([ + group.runSettled(echo, { text: 'ok' }), + group.runSettled(failing), + ]), + ), + }), + }, + }), + }) + const { adapter } = mockAdapter([]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/orchestrate', + adapter, + agents: [echo, failing, slow], + plugins: () => [orchestrator], + }), + { threadId: 't' }, + ) + + expect(await session.command('one')).toBe('HI') + expect(await session.command('both')).toEqual(['A', 'B']) + await expect(session.command('failFast')).rejects.toThrow('boom') + expect(slowStopped).toBe(true) + const collected = await session.command('collect') + expect(collected).toMatchObject([{ ok: true, value: 'OK' }, { ok: false }]) + await host.close() + }) +}) + +describe('default limits', () => { + it('applies depth, concurrency, and call limits to the main model', async () => { + expect(DEFAULT_SUBAGENT_LIMITS).toEqual({ + maxDepth: 2, + maxConcurrent: 3, + maxCalls: 12, + }) + const run = vi.fn(async () => 'done') + const worker = defineAgent({ name: 'worker', description: 'Works', run }) + const outer = mockAdapter([ + () => toolCall('worker', {}, 'c1'), + () => toolCall('worker', {}, 'c2'), + () => text('stopped'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/limits', + adapter: outer.adapter, + subagents: { agents: [worker], limits: { maxCalls: 1 } }, + }), + { threadId: 't' }, + ) + await session.prompt('work twice') + expect(run).toHaveBeenCalledTimes(1) + expect(JSON.stringify(outer.calls[2].messages)).toContain( + 'subagent limit reached (maxCalls 1)', + ) + await host.close() + }) +}) diff --git a/packages/ai/src/activities/chat/agents/bound.ts b/packages/ai/src/activities/chat/agents/bound.ts index bc73311b1a..8cbd182921 100644 --- a/packages/ai/src/activities/chat/agents/bound.ts +++ b/packages/ai/src/activities/chat/agents/bound.ts @@ -15,6 +15,7 @@ import type { GenerationMiddleware } from '../../middleware/types' import type { AnyChatMiddleware } from '../middleware/types' import type { RunAgentResumeItem } from '../../../types' import type { SubagentRunInput } from './define-agent' +import type { SubagentBudget } from './limits' /** * The fields a child `chat()` needs, in one spread: @@ -40,6 +41,11 @@ export interface SubagentBinding { chatMiddleware?: ReadonlyArray /** Added before the call's own middleware on every generation call. */ generationMiddleware?: ReadonlyArray + /** + * The subagent tree budget. A child's `ctx.chat({ subagents })` passes it + * down, so limits hold across the whole tree. + */ + budget?: SubagentBudget } /** @@ -109,6 +115,18 @@ export function createBoundActivities( ...(input.resume ? { resume: input.resume } : {}), abortController, ...options, + // Nested children share the tree budget. + ...(options.subagents && binding?.budget + ? { + subagents: { + ...options.subagents, + binding: { + ...options.subagents.binding, + budget: binding.budget, + }, + }, + } + : {}), middleware: [...chatMiddleware, ...(options.middleware ?? [])], } as never)) as typeof chat, summarize: ((options: LooseOptions) => 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 0b3733de42..2782be0e7e 100644 --- a/packages/ai/src/activities/chat/agents/spawn.ts +++ b/packages/ai/src/activities/chat/agents/spawn.ts @@ -26,7 +26,9 @@ import type { UIMessage, } from '../../../types' import { createBoundActivities } from './bound' +import { SubagentBudget } from './limits' import type { SubagentBinding } from './bound' +import type { SubagentLimits } from './limits' import type { DefinedAgent, SubagentRunContext, @@ -81,6 +83,12 @@ export interface SubagentsBag< */ order?: 'parallel' | 'sequence' sandbox?: 'own' | 'inherit' + /** + * 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. + */ + limits?: SubagentLimits /** * What a host adds to every activity call a child makes through `ctx` * (`ctx.chat`, `ctx.generateImage`, and the rest). A harness session sets @@ -240,6 +248,7 @@ function openAgentStream( ctx: SpawnContext, sink?: SubagentSink, parentToolCallId?: string, + binding: SubagentBinding | undefined = bag.binding, ) { const agent = agentByName(bag.agents, entry.name) const resume = entry.resume @@ -273,7 +282,7 @@ function openAgentStream( }, sink, parentToolCallId, - bag.binding, + binding, ) } @@ -818,6 +827,9 @@ export function createSyntheticSubagentTools( sink: SubagentSink }, ): Array { + // One budget per root run. A child run inherits its parent's budget. + const budget = bag.binding?.budget ?? SubagentBudget.root(bag.limits) + let active = 0 return bag.agents.map((agent) => ({ name: agent.name, description: agent.description, @@ -852,8 +864,35 @@ export function createSyntheticSubagentTools( }, }), } + // A resumed child already counted when it first started. + const refusal = suspended ? undefined : budget.reserve(active) + if (refusal !== undefined) { + return { + subagentRunId: '', + text: '', + error: refusal, + } satisfies SubagentToolOutcome + } + active += 1 const sink = createSubagentSink() const link = linkAbort(parent.abortSignal) + const timeout = budget.childTimeout() + const timer = + timeout === undefined + ? undefined + : setTimeout( + () => + 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 @@ -873,6 +912,7 @@ export function createSyntheticSubagentTools( }, sink, toolCallId, + childBinding, )) { if (chunk.type === SUBAGENT_STARTED && subagentRunId === '') { subagentRunId = chunk.subagentRunId @@ -900,6 +940,8 @@ export function createSyntheticSubagentTools( toolContext?.[EMIT_STREAM_CHUNK]?.(chunk) } } finally { + active -= 1 + if (timer !== undefined) clearTimeout(timer) link.dispose() } parent.sink.usage.push(...sink.usage) diff --git a/packages/ai/src/activities/index.ts b/packages/ai/src/activities/index.ts index c2d9013816..37fd9c9159 100644 --- a/packages/ai/src/activities/index.ts +++ b/packages/ai/src/activities/index.ts @@ -53,6 +53,7 @@ export type { SubagentBinding, SubagentForward, } from './chat/agents/bound' +export { SubagentBudget, type SubagentLimits } from './chat/agents/limits' export { subagentRoute, type SubagentRouteOptions } from './chat/agents/route' export type { SubagentOrder, diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 8e2889ab62..e442e22882 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -99,6 +99,10 @@ export type { SubagentBinding, SubagentForward, } from './activities/chat/agents/bound' +export { + SubagentBudget, + type SubagentLimits, +} from './activities/chat/agents/limits' export { subagentRoute, type SubagentRouteOptions, diff --git a/packages/ai/tests/subagent-limits.test.ts b/packages/ai/tests/subagent-limits.test.ts new file mode 100644 index 0000000000..031dcc4c86 --- /dev/null +++ b/packages/ai/tests/subagent-limits.test.ts @@ -0,0 +1,159 @@ +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).toBeDefined() + }) + + 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 } }, + }) as AsyncIterable, + ) + const errors = Object.values(results(chunks)) + .map((result) => result.error) + .filter(Boolean) + // Tool calls in one turn may run one after another; a refusal only happens + // when they overlap. Either way no more than two run at once. + expect( + errors.every( + (error) => error === 'subagent limit reached (maxConcurrent 2)', + ), + ).toBe(true) + }) +}) diff --git a/testing/e2e/fixtures/harness/basic.json b/testing/e2e/fixtures/harness/basic.json index edf1a61780..f717368b80 100644 --- a/testing/e2e/fixtures/harness/basic.json +++ b/testing/e2e/fixtures/harness/basic.json @@ -39,6 +39,43 @@ "response": { "content": "Hello over the harness protocol." } + }, + { + "match": { + "userMessage": "[harness-limits] work twice", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "worker", + "arguments": "{}" + } + ] + } + }, + { + "match": { + "userMessage": "[harness-limits] work twice", + "sequenceIndex": 1 + }, + "response": { + "toolCalls": [ + { + "name": "worker", + "arguments": "{}" + } + ] + } + }, + { + "match": { + "userMessage": "[harness-limits] work twice", + "sequenceIndex": 2 + }, + "response": { + "content": "The second call hit the limit." + } } ] } diff --git a/testing/e2e/src/routes/api.harness-test.ts b/testing/e2e/src/routes/api.harness-test.ts index 6b7be92b52..b66add840d 100644 --- a/testing/e2e/src/routes/api.harness-test.ts +++ b/testing/e2e/src/routes/api.harness-test.ts @@ -48,6 +48,30 @@ export const Route = createFileRoute('/api/harness-test')({ const host = createHarnessHost({ persistence }) try { const session = await host.open(harness, { threadId: 'e2e-thread' }) + if (body.scenario === 'limits') { + // The main model calls `worker` twice. The limit allows one. + let workerRuns = 0 + const worker = defineAgent({ + name: 'worker', + description: 'Does one unit of work', + run: async () => { + workerRuns += 1 + return 'worked' + }, + }) + const limited = defineHarness({ + name: 'e2e/harness-limits', + adapter: openai(), + subagents: { agents: [worker], limits: { maxCalls: 1 } }, + }) + const limitedSession = await host.open(limited, { + threadId: 'e2e-limits', + }) + const turn = await limitedSession.prompt( + '[harness-limits] work twice', + ) + return Response.json({ workerRuns, text: turn.text }) + } if (body.scenario === 'agent') { const result = await session.agents.pricer.run({ task: '[harness-agent] price vendor a', diff --git a/testing/e2e/tests/harness.spec.ts b/testing/e2e/tests/harness.spec.ts index 03e3c458a5..609a642592 100644 --- a/testing/e2e/tests/harness.spec.ts +++ b/testing/e2e/tests/harness.spec.ts @@ -15,6 +15,20 @@ test.describe('harness session', () => { expect(body.roles).toEqual(['user', 'assistant', 'user', 'assistant']) }) + test('subagent limits refuse the second child start', async ({ + request, + testId, + aimockPort, + }) => { + const response = await request.post('/api/harness-test', { + data: { scenario: 'limits', testId, aimockPort }, + }) + expect(response.ok()).toBe(true) + const body = await response.json() + expect(body.workerRuns).toBe(1) + expect(body.text).toBe('The second call hit the limit.') + }) + test('a typed agent runs from code and the next turn sees its result', async ({ request, testId, From bd1ac5bc77e49c5d54a7f7b7d3e8ef222cd142b6 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 13:22:50 +0200 Subject: [PATCH 011/286] feat(ai-harness): build artifacts, worker mode, and remote harnessText @tanstack/ai-harness/build: buildHarness bundles a harness with Bun into a worker artifact (harness.js) plus harness.manifest.json (name, agents, plugins, requirements, sha256 digest), and can compile a single executable. artifactText(dir) checks the digest, starts one worker process per thread, and uses it as a text adapter. @tanstack/ai-harness/worker: runHarnessWorker serves session-tier frames as NDJSON on stdin and stdout. harnessText({ url, token }) uses a harness served on another machine (for example with runCli --serve) as a text adapter. --- .changeset/harness-p5-build.md | 9 + docs/config.json | 5 + docs/harness/deploy.md | 97 +++++ packages/ai-harness/package.json | 8 + packages/ai-harness/src/build.ts | 428 +++++++++++++++++++++ packages/ai-harness/src/harness-text.ts | 136 +++++++ packages/ai-harness/src/index.ts | 2 +- packages/ai-harness/src/worker.ts | 84 ++++ packages/ai-harness/tests/build.test.ts | 195 ++++++++++ packages/ai-harness/vite.config.ts | 4 +- testing/e2e/src/routes/api.harness-test.ts | 25 +- testing/e2e/tests/harness.spec.ts | 14 + 12 files changed, 1003 insertions(+), 4 deletions(-) create mode 100644 .changeset/harness-p5-build.md create mode 100644 docs/harness/deploy.md create mode 100644 packages/ai-harness/src/build.ts create mode 100644 packages/ai-harness/src/worker.ts create mode 100644 packages/ai-harness/tests/build.test.ts diff --git a/.changeset/harness-p5-build.md b/.changeset/harness-p5-build.md new file mode 100644 index 0000000000..0a6e2a896d --- /dev/null +++ b/.changeset/harness-p5-build.md @@ -0,0 +1,9 @@ +--- +'@tanstack/ai-harness': minor +--- + +Run a harness outside your server. + +- **`@tanstack/ai-harness/build`**: `buildHarness` bundles a harness with Bun into a worker artifact (`harness.js`) and writes `harness.manifest.json` (name, agents, plugins, requirements, sha256 digest). With `compile`, it also builds a single executable. `artifactText(dir)` runs the artifact as worker processes and uses it as a text adapter, after it checks the digest. `readManifest` reads and checks a manifest. +- **`@tanstack/ai-harness/worker`**: `runHarnessWorker` serves a harness as session-tier frames over NDJSON on stdin and stdout. +- **`harnessText({ url, token })`** uses a harness served by `createHarnessHandler` (or `runCli --serve`) on another machine as a text adapter. diff --git a/docs/config.json b/docs/config.json index c1c32488d3..533ad81083 100644 --- a/docs/config.json +++ b/docs/config.json @@ -843,6 +843,11 @@ "label": "Run agents from a harness", "to": "harness/subagents", "addedAt": "2026-09-26" + }, + { + "label": "Deploy a harness", + "to": "harness/deploy", + "addedAt": "2026-09-26" } ] }, diff --git a/docs/harness/deploy.md b/docs/harness/deploy.md new file mode 100644 index 0000000000..d6d22b92db --- /dev/null +++ b/docs/harness/deploy.md @@ -0,0 +1,97 @@ +--- +title: Deploy a harness +id: harness-deploy +order: 9 +description: "Run a harness in your server, as a worker process, on another machine, or as a single executable." +keywords: + - tanstack ai + - harness + - deploy + - build + - worker +--- + +Your harness works on your laptop. Now it needs to run next to your API, in its own process, on a different machine, or as a CLI your team can download. You pick the place. The harness definition stays the same. + +| Place | How | +|---|---| +| Your server | `createHarnessHost` and `createHarnessHandler` in your routes. No build step. | +| A worker process | `buildHarness`, then `artifactText(dir)` as the model of a `chat()` call. | +| Another machine | `runCli(harness)` with `--serve` there, then `harnessText({ url, token })` here. | +| A single executable | `buildHarness({ compile })` with Bun. | + +## Build a worker artifact + +`buildHarness` bundles the harness with Bun and writes a manifest next to it. Install [Bun](https://bun.sh) on the build machine. + +```ts group=harness-deploy +import { artifactText, buildHarness } from '@tanstack/ai-harness/build' + +const { manifest } = await buildHarness({ + entry: './src/studio.ts', + export: 'studio', + outDir: './dist/studio', +}) +console.log(manifest.name, manifest.digest) +``` + +- `dist/studio/harness.js` starts the harness as a worker. It reads session frames on stdin and writes them on stdout. +- `dist/studio/harness.manifest.json` names the harness, its agents and plugins, what it needs (file system, processes, network), and the sha256 digest of the bundle. + +The build imports your entry in a child process to read the harness name, agents, and plugins. Importing a module runs its code. + +## Run the artifact as a model + +`artifactText` starts one worker process per outer thread and checks the digest first. A changed bundle is refused. + +```ts group=harness-deploy +import { chat } from '@tanstack/ai' + +const model = await artifactText('./dist/studio') +const stream = chat({ + adapter: model, + messages: [{ role: 'user', content: 'Plan the launch post.' }], + threadId: 'thread-1', +}) +// When you are done with the workers: +model.dispose() +``` + +The worker keeps the conversation of its thread between calls. Workers get the environment of the process that starts them, so they can read your provider keys. + +## Use a harness on another machine + +On the other machine, serve the harness with the CLI. It prints a token unless you set one: + +```bash +HARNESS_TOKEN=your-token npx tsx cli.ts --serve --host 0.0.0.0 +``` + +Here, use it as a model: + +```ts group=harness-deploy +import { harnessText } from '@tanstack/ai-harness' + +const remote = harnessText({ url: 'http://build-box:8787', token: 'your-token' }) +``` + +The remote session keeps the conversation, so each call sends only the new message. + +## Ship a single executable + +Point `compile.entry` at a file that calls `runCli`: + +```ts group=harness-deploy +await buildHarness({ + entry: './src/studio.ts', + export: 'studio', + outDir: './dist/studio', + compile: { entry: './src/cli.ts', outfile: './dist/studio-cli', target: 'bun-linux-x64' }, +}) +``` + +Bun can build for another platform with `target`, for example `bun-darwin-arm64` or `bun-windows-x64`. + +## What you have now + +- The same harness in your server, in a worker process, on another machine, or as a binary. diff --git a/packages/ai-harness/package.json b/packages/ai-harness/package.json index 01b208a614..fcecc105a1 100644 --- a/packages/ai-harness/package.json +++ b/packages/ai-harness/package.json @@ -37,6 +37,14 @@ "./plugins": { "types": "./dist/esm/first-party/index.d.ts", "import": "./dist/esm/first-party/index.js" + }, + "./build": { + "types": "./dist/esm/build.d.ts", + "import": "./dist/esm/build.js" + }, + "./worker": { + "types": "./dist/esm/worker.d.ts", + "import": "./dist/esm/worker.js" } }, "scripts": { diff --git a/packages/ai-harness/src/build.ts b/packages/ai-harness/src/build.ts new file mode 100644 index 0000000000..4d3cf074e8 --- /dev/null +++ b/packages/ai-harness/src/build.ts @@ -0,0 +1,428 @@ +import { spawn } from 'node:child_process' +import { createHash, randomUUID } from 'node:crypto' +import { existsSync } from 'node:fs' +import { mkdir, readFile, rm, writeFile } from 'node:fs/promises' +import { dirname, join, resolve } from 'node:path' +import { createInterface } from 'node:readline' +import { fileURLToPath } from 'node:url' +import { EventType } from '@tanstack/ai' +import { HARNESS_PROTOCOL_VERSION } from './protocol' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import type { HostFrame } from './protocol' + +/** What a built harness says about itself. Fields are claims; deploy policy decides. */ +export interface HarnessManifestV1 { + format: 'tanstack-ai-harness' + version: 1 + name: string + /** sha256 of the bundle. */ + digest: string + /** The bundle file, relative to the manifest. */ + entry: string + runtime: { kind: 'node'; major: number } + protocol: typeof HARNESS_PROTOCOL_VERSION + agents: Array<{ name: string; produces?: string }> + plugins: Array<{ name: string }> + requires: { + filesystem: boolean + processExecution: boolean + network: 'model-only' | 'declared' + } +} + +export interface BuildHarnessOptions { + /** The module that exports the harness. */ + entry: string + /** The export name. Default `'default'`. */ + export?: string + outDir: string + /** + * Also build a single executable of `compile.entry` (for example a file + * that calls `runCli`) with Bun. + */ + compile?: { entry: string; outfile: string; target?: string } + /** The Bun executable. Default `'bun'`. */ + bun?: string +} + +const MANIFEST = 'harness.manifest.json' +const BUNDLE = 'harness.js' + +function run( + command: string, + args: Array, + cwd: string, +): Promise { + // On Windows, npm installs Bun as a `.cmd` shim, which only a shell can + // start. The arguments are paths this module builds, quoted for cmd. + const viaShell = process.platform === 'win32' && !/\.exe$/i.test(command) + return new Promise((done, fail) => { + const child = viaShell + ? spawn( + command, + args.map((arg) => `"${arg}"`), + { cwd, stdio: ['ignore', 'pipe', 'pipe'], shell: true }, + ) + : spawn(command, args, { cwd, stdio: ['ignore', 'pipe', 'pipe'] }) + let stdout = '' + let stderr = '' + child.stdout.on('data', (data: Buffer) => (stdout += data.toString())) + child.stderr.on('data', (data: Buffer) => (stderr += data.toString())) + child.on('error', (error) => + fail( + new Error( + `Could not run ${command}. buildHarness uses Bun: install it from https://bun.sh. (${error.message})`, + ), + ), + ) + child.on('close', (code) => + code === 0 + ? done(stdout) + : fail( + new Error( + `${command} ${args[0] ?? ''} failed:\n${stderr || stdout}`, + ), + ), + ) + }) +} + +/** + * Bundle a harness into a worker artifact: `harness.js` (starts a worker on + * stdin and stdout) and `harness.manifest.json`. With `compile`, also build a + * single executable. + * + * The build imports `entry` in a child process to read the harness name, + * agents, and plugin names. Importing a module runs its code. + */ +export async function buildHarness( + options: BuildHarnessOptions, +): Promise<{ + manifest: HarnessManifestV1 + bundle: string + executable?: string +}> { + const bun = options.bun ?? 'bun' + const entry = resolve(options.entry) + const outDir = resolve(options.outDir) + const exportName = options.export ?? 'default' + // Bun resolves plain absolute paths, not file:// URLs. Forward slashes + // keep Windows paths valid inside the generated modules. + const slash = (path: string) => path.split('\\').join('/') + const entryPath = slash(entry) + const builtWorker = fileURLToPath(new URL('./worker.js', import.meta.url)) + // From source (tests), the worker module is still TypeScript. + const workerModule = slash( + existsSync(builtWorker) ? builtWorker : builtWorker.replace(/\.js$/, '.ts'), + ) + await mkdir(outDir, { recursive: true }) + + // Temporary files sit next to the entry, so its imports resolve the same way. + const id = randomUUID().slice(0, 8) + const bootstrap = join(dirname(entry), `.harness-worker-${id}.mjs`) + const describe = join(dirname(entry), `.harness-describe-${id}.mjs`) + await writeFile( + bootstrap, + [ + `import * as entry from ${JSON.stringify(entryPath)}`, + `import { runHarnessWorker } from ${JSON.stringify(workerModule)}`, + `await runHarnessWorker(entry[${JSON.stringify(exportName)}])`, + ].join('\n'), + ) + await writeFile( + describe, + [ + `const entry = await import(${JSON.stringify(entryPath)})`, + `const harness = entry[${JSON.stringify(exportName)}]`, + `if (!harness || harness.kind !== 'tanstack-ai-harness') throw new Error('Export ${exportName} of ${entry.replace(/\\/g, '/')} is not a harness.')`, + `const agents = [...(harness.agents ?? []), ...(harness.subagents?.agents ?? [])]`, + `const plugins = harness.plugins ? harness.plugins().map((plugin) => ({ name: plugin.name })) : []`, + `console.log(JSON.stringify({ name: harness.name, agents: agents.map((agent) => ({ name: agent.name, ...(agent.produces ? { produces: agent.produces } : {}) })), plugins }))`, + ].join('\n'), + ) + + try { + const bundle = join(outDir, BUNDLE) + await run( + bun, + ['build', bootstrap, '--target=node', `--outfile=${bundle}`], + dirname(entry), + ) + const described: { + name: string + agents: HarnessManifestV1['agents'] + plugins: HarnessManifestV1['plugins'] + } = JSON.parse( + (await run(bun, [describe], dirname(entry))).trim().split('\n').at(-1) ?? + '{}', + ) + const names = described.plugins.map((plugin) => plugin.name) + const usesWorkspace = names.includes('tanstack/workspace-tools') + const manifest: HarnessManifestV1 = { + format: 'tanstack-ai-harness', + version: 1, + name: described.name, + digest: createHash('sha256') + .update(await readFile(bundle)) + .digest('hex'), + entry: BUNDLE, + runtime: { + kind: 'node', + major: Number(process.versions.node.split('.')[0]), + }, + protocol: HARNESS_PROTOCOL_VERSION, + agents: described.agents, + plugins: described.plugins, + requires: { + filesystem: + usesWorkspace || names.includes('tanstack/project-instructions'), + processExecution: usesWorkspace, + network: names.some((name) => name.startsWith('connector/')) + ? 'declared' + : 'model-only', + }, + } + await writeFile( + join(outDir, MANIFEST), + `${JSON.stringify(manifest, null, 2)}\n`, + ) + + let executable: string | undefined + if (options.compile) { + executable = resolve(options.compile.outfile) + await run( + bun, + [ + 'build', + resolve(options.compile.entry), + '--compile', + `--outfile=${executable}`, + ...(options.compile.target + ? [`--target=${options.compile.target}`] + : []), + ], + dirname(resolve(options.compile.entry)), + ) + } + return { manifest, bundle, ...(executable ? { executable } : {}) } + } finally { + await rm(bootstrap, { force: true }) + await rm(describe, { force: true }) + } +} + +/** Read and check the manifest of a built harness. */ +export async function readManifest(dir: string): Promise { + const manifest: unknown = JSON.parse( + await readFile(join(resolve(dir), MANIFEST), 'utf8'), + ) + if ( + typeof manifest !== 'object' || + manifest === null || + !('format' in manifest) || + manifest.format !== 'tanstack-ai-harness' || + !('version' in manifest) || + manifest.version !== 1 + ) { + throw new Error(`${dir} has no tanstack-ai-harness v1 manifest.`) + } + // The format and version checks above identify the file. + return manifest as HarnessManifestV1 +} + +const FORWARDED = new Set([ + EventType.TEXT_MESSAGE_START, + EventType.TEXT_MESSAGE_CONTENT, + EventType.TEXT_MESSAGE_END, + EventType.REASONING_START, + EventType.REASONING_MESSAGE_START, + EventType.REASONING_MESSAGE_CONTENT, + EventType.REASONING_MESSAGE_END, + EventType.REASONING_END, +]) + +interface Worker { + send: (frame: object) => void + frames: (listener: (frame: HostFrame) => void) => () => void + kill: () => void +} + +function startWorker(node: string, bundle: string, threadId: string): Worker { + const child = spawn(node, [bundle], { stdio: ['pipe', 'pipe', 'inherit'] }) + const listeners = new Set<(frame: HostFrame) => void>() + createInterface({ input: child.stdout }).on('line', (line) => { + try { + const frame = JSON.parse(line) as HostFrame + for (const listener of listeners) listener(frame) + } catch { + // A line that is not a frame (for example a stray log) is skipped. + } + }) + const send = (frame: object) => + child.stdin.write(`${JSON.stringify(frame)}\n`) + send({ type: 'harness.subscribe', threadId }) + return { + send, + frames: (listener) => { + listeners.add(listener) + return () => listeners.delete(listener) + }, + kill: () => { + child.stdin.end() + child.kill() + }, + } +} + +function lastUserText(messages: unknown): string { + const list: ReadonlyArray = Array.isArray(messages) ? messages : [] + const message = list.findLast( + (entry) => + typeof entry === 'object' && + entry !== null && + 'role' in entry && + entry.role === 'user', + ) + if ( + typeof message !== 'object' || + message === null || + !('content' in message) + ) + return '' + return typeof message.content === 'string' ? message.content : '' +} + +/** + * Use a built harness as the model of a `chat()` call. Each outer thread gets + * its own worker process (`node harness.js`). Call `dispose()` to stop them. + */ +export async function artifactText( + dir: string, + options: { node?: string } = {}, +): Promise void }> { + const manifest = await readManifest(dir) + const bundle = join(resolve(dir), manifest.entry) + const digest = createHash('sha256') + .update(await readFile(bundle)) + .digest('hex') + if (digest !== manifest.digest) { + throw new Error(`The bundle in ${dir} does not match its manifest digest.`) + } + const workers = new Map() + let requests = 0 + + return { + kind: 'text', + name: 'harness-artifact', + model: manifest.name, + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + dispose: () => { + for (const worker of workers.values()) worker.kill() + workers.clear() + }, + structuredOutput: () => + Promise.reject( + new Error('artifactText does not support structured output.'), + ), + chatStream: (chatOptions) => + (async function* (): AsyncGenerator { + const threadId = chatOptions.threadId ?? 'default' + let worker = workers.get(threadId) + if (!worker) { + worker = startWorker( + options.node ?? process.execPath, + bundle, + threadId, + ) + workers.set(threadId, worker) + } + const runId = chatOptions.runId ?? `artifact-${Date.now().toString(36)}` + const requestId = `r${(requests += 1)}` + const queue: Array = [] + let wake: (() => void) | undefined + const stop = worker.frames((frame) => { + queue.push(frame) + wake?.() + }) + worker.send({ + type: 'harness.input', + requestId, + input: { op: 'prompt', message: lastUserText(chatOptions.messages) }, + }) + yield { + type: EventType.RUN_STARTED, + runId, + threadId, + timestamp: Date.now(), + } + let operationId: string | undefined + const early: Array> = [] + try { + while (true) { + if (queue.length === 0) { + await new Promise((resolveWait) => (wake = resolveWait)) + wake = undefined + } + const frame = queue.shift() + if (!frame) continue + if (frame.type === 'harness.error') throw new Error(frame.message) + if ( + frame.type === 'harness.receipt' && + frame.requestId === requestId + ) { + if (frame.status === 'rejected') + throw new Error( + frame.reason ?? 'The worker refused the prompt.', + ) + operationId = frame.operationId + queue.unshift( + ...early.filter((entry) => entry.operationId === operationId), + ) + continue + } + if (frame.type !== 'harness.event') continue + if (!operationId) { + early.push(frame) + continue + } + if (frame.operationId !== operationId) continue + const event = frame.event + if ( + FORWARDED.has(event.type) && + !('subagentRunId' in event && event.subagentRunId) + ) + yield event + if (event.type === EventType.RUN_ERROR) + throw new Error(event.message) + if ( + event.type === EventType.CUSTOM && + event.name === 'harness.operation.finished' + ) + break + } + } finally { + stop() + } + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + metadata: { tanstack: { finishReason: 'stop' } }, + } + })(), + } +} diff --git a/packages/ai-harness/src/harness-text.ts b/packages/ai-harness/src/harness-text.ts index 51b0aa7a21..2550d5c659 100644 --- a/packages/ai-harness/src/harness-text.ts +++ b/packages/ai-harness/src/harness-text.ts @@ -1,6 +1,7 @@ import { EventType } from '@tanstack/ai' import { createHarnessHost } from './host' import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import { isHarnessDefinition } from './define' import type { AnyHarness } from './define' import type { HarnessHost } from './host' @@ -54,10 +55,145 @@ const FORWARDED = new Set([ * const stream = chat({ adapter: harnessText(studio), messages, threadId }) * ``` */ +/** A harness served by `createHarnessHandler` (or `runCli --serve`) on another machine. */ +export interface RemoteHarness { + /** The handler base URL, for example `http://127.0.0.1:8787`. */ + url: string + /** The bearer token. */ + token?: string + fetch?: typeof fetch +} + +const TYPES = { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, +} + +/** Stream the SSE `data:` payloads of a response. */ +async function* sseData(response: Response): AsyncGenerator { + if (!response.body) return + const reader = response.body.getReader() + const decoder = new TextDecoder() + let buffer = '' + while (true) { + const { value, done } = await reader.read() + if (done) break + buffer += decoder.decode(value, { stream: true }) + const blocks = buffer.split('\n\n') + buffer = blocks.pop() ?? '' + for (const block of blocks) { + const data = block + .split('\n') + .find((line) => line.startsWith('data: ')) + ?.slice(6) + if (data) yield JSON.parse(data) + } + } +} + +function remoteHarnessText(remote: RemoteHarness): AnyTextAdapter { + const base = remote.url.replace(/\/$/, '') + const doFetch = remote.fetch ?? fetch + return { + kind: 'text', + name: 'harness-remote', + model: base, + '~types': TYPES, + structuredOutput: () => + Promise.reject( + new Error('harnessText does not support structured output.'), + ), + chatStream: (chatOptions) => + (async function* (): AsyncGenerator { + const threadId = chatOptions.threadId ?? 'default' + const runId = chatOptions.runId ?? `harness-${Date.now().toString(36)}` + const response = await doFetch(`${base}/run`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + ...(remote.token + ? { Authorization: `Bearer ${remote.token}` } + : {}), + }, + body: JSON.stringify({ + threadId, + runId, + // The remote session keeps its own history, so send the new message only. + messages: [ + { + id: `${runId}-user`, + role: 'user', + content: lastUserText(chatOptions.messages), + }, + ], + tools: [], + context: [], + state: {}, + forwardedProps: {}, + }), + }) + if (!response.ok) { + throw new Error( + `Remote harness failed (${response.status}): ${await response.text()}`, + ) + } + yield { + type: EventType.RUN_STARTED, + runId, + threadId, + timestamp: Date.now(), + } + for await (const data of sseData(response)) { + if (!isRecord(data) || typeof data.type !== 'string') continue + // The handler sends AG-UI chunks over SSE. + const chunk = data as StreamChunk + if ( + FORWARDED.has(chunk.type) && + !('subagentRunId' in chunk && chunk.subagentRunId) + ) { + yield chunk + } + if (chunk.type === EventType.RUN_ERROR) { + yield { + type: EventType.RUN_ERROR, + message: chunk.message, + timestamp: Date.now(), + } + return + } + } + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + metadata: { tanstack: { finishReason: 'stop' } }, + } + })(), + } +} + export function harnessText( harness: AnyHarness, + options?: HarnessTextOptions, +): AnyTextAdapter +export function harnessText(remote: RemoteHarness): AnyTextAdapter +export function harnessText( + target: AnyHarness | RemoteHarness, options: HarnessTextOptions = {}, ): AnyTextAdapter { + if (!isHarnessDefinition(target)) return remoteHarnessText(target) + const harness = target let host = options.host return { kind: 'text', diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts index 3fda086bd6..cc4fa6a623 100644 --- a/packages/ai-harness/src/index.ts +++ b/packages/ai-harness/src/index.ts @@ -97,7 +97,7 @@ export type { } from './http' export { harnessText } from './harness-text' -export type { HarnessTextOptions } from './harness-text' +export type { HarnessTextOptions, RemoteHarness } from './harness-text' export { buildAuthorizationUrl, diff --git a/packages/ai-harness/src/worker.ts b/packages/ai-harness/src/worker.ts new file mode 100644 index 0000000000..233b005c31 --- /dev/null +++ b/packages/ai-harness/src/worker.ts @@ -0,0 +1,84 @@ +import { createInterface } from 'node:readline' +import { createHarnessHost } from './host' +import { + HARNESS_PROTOCOL_VERSION, + applyInput, + parseControlFrame, +} from './protocol' +import type { AnyHarness } from './define' +import type { HarnessPersistence } from './host' +import type { HostFrame } from './protocol' + +export interface HarnessWorkerOptions { + input?: NodeJS.ReadableStream + output?: { write: (text: string) => unknown } + persistence?: HarnessPersistence +} + +/** + * Run a harness as a worker: session-tier frames as NDJSON on stdin and + * stdout. The first frame must be `harness.subscribe`. The worker exits when + * stdin closes. `artifactText` starts workers like this. + */ +export async function runHarnessWorker( + harness: AnyHarness, + options: HarnessWorkerOptions = {}, +): Promise { + const output = options.output ?? process.stdout + const send = (frame: HostFrame) => output.write(`${JSON.stringify(frame)}\n`) + const host = createHarnessHost( + options.persistence ? { persistence: options.persistence } : {}, + ) + const reader = new AbortController() + let session: Awaited> | undefined + const lines = createInterface({ + input: options.input ?? process.stdin, + crlfDelay: Infinity, + }) + + try { + for await (const line of lines) { + if (line.trim() === '') continue + try { + const frame = parseControlFrame(line) + if (frame.type === 'harness.subscribe') { + if (session) throw new Error('Already subscribed.') + session = await host.open(harness, { threadId: frame.threadId }) + send({ + type: 'harness.hello', + v: HARNESS_PROTOCOL_VERSION, + threadId: frame.threadId, + }) + const events = session.events({ + ...(frame.from ? { from: frame.from } : {}), + signal: reader.signal, + }) + void (async () => { + for await (const entry of events) + send({ type: 'harness.event', ...entry }) + })() + continue + } + if (!session) throw new Error('Send harness.subscribe first.') + if (frame.type === 'harness.snapshot') { + send({ type: 'harness.snapshot', snapshot: session.snapshot() }) + continue + } + const receipt = await applyInput(harness, session, frame.input) + send({ + type: 'harness.receipt', + requestId: frame.requestId, + ...receipt, + }) + } catch (error) { + send({ + type: 'harness.error', + message: error instanceof Error ? error.message : String(error), + }) + } + } + } finally { + reader.abort() + await host.close() + } +} diff --git a/packages/ai-harness/tests/build.test.ts b/packages/ai-harness/tests/build.test.ts new file mode 100644 index 0000000000..90cd18fbd8 --- /dev/null +++ b/packages/ai-harness/tests/build.test.ts @@ -0,0 +1,195 @@ +import { mkdir, readFile, rm, writeFile } from 'node:fs/promises' +import { PassThrough } from 'node:stream' +import { fileURLToPath } from 'node:url' +import { join } from 'node:path' +import { afterAll, describe, expect, it } from 'vitest' +import { EventType, chat } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + createHarnessHandler, + createHarnessHost, + defineHarness, + harnessText, +} from '../src' +import { artifactText, buildHarness, readManifest } from '../src/build' +import { runHarnessWorker } from '../src/worker' +import { mockAdapter, text } from './helpers' +import type { StreamChunk } from '@tanstack/ai' + +async function textOf(stream: AsyncIterable): Promise { + let out = '' + for await (const chunk of stream) { + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) out += chunk.delta + if (chunk.type === EventType.RUN_ERROR) throw new Error(chunk.message) + } + return out +} + +describe('worker mode', () => { + it('speaks session frames over NDJSON', async () => { + const { adapter } = mockAdapter([() => text('from the worker')]) + const input = new PassThrough() + const lines: Array = [] + const running = runHarnessWorker( + defineHarness({ name: 'test/worker', adapter }), + { + input, + output: { write: (line: string) => lines.push(JSON.parse(line)) }, + persistence: memoryPersistence(), + }, + ) + input.write('{"type":"harness.subscribe","threadId":"w1"}\n') + input.write( + '{"type":"harness.input","requestId":"r1","input":{"op":"prompt","message":"hi"}}\n', + ) + await new Promise((resolve) => { + const check = setInterval(() => { + if ( + lines.some( + (frame) => frame.event?.name === 'harness.operation.finished', + ) + ) { + clearInterval(check) + resolve() + } + }, 10) + }) + input.end() + await running + expect(lines[0]).toMatchObject({ type: 'harness.hello', threadId: 'w1' }) + expect( + lines.find((frame) => frame.type === 'harness.receipt'), + ).toMatchObject({ requestId: 'r1', status: 'accepted' }) + expect(JSON.stringify(lines)).toContain('from the worker') + }) +}) + +describe('harnessText({ url })', () => { + it('uses a harness served over HTTP as a model', async () => { + const { adapter } = mockAdapter([() => text('from far away')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const handler = createHarnessHandler({ + host, + harness: defineHarness({ name: 'test/remote', adapter }), + authorize: (request) => + request.headers.get('authorization') === 'Bearer t' + ? { id: 'u' } + : null, + }) + const remote = harnessText({ + url: 'http://remote.test/api', + token: 't', + fetch: (input, init) => handler(new Request(input, init)), + }) + const answer = await textOf( + chat({ + adapter: remote, + messages: [{ role: 'user', content: 'hi' }], + threadId: 'r', + }) as AsyncIterable, + ) + expect(answer).toBe('from far away') + await host.close() + }) +}) + +const tmp = fileURLToPath(new URL('./.tmp-build/', import.meta.url)) + +describe('buildHarness and artifactText', () => { + afterAll(async () => { + await rm(tmp, { recursive: true, force: true }) + }) + + it( + 'bundles a harness, writes a manifest, and runs it as a worker process', + { timeout: 60_000 }, + async () => { + await mkdir(tmp, { recursive: true }) + const entry = join(tmp, 'studio.ts') + await writeFile( + entry, + ` +import { EventType, defineAgent } from '@tanstack/ai' +import { defineHarness } from '../../src/index' + +let calls = 0 +const adapter = { + kind: 'text', + name: 'inline', + model: 'inline-model', + '~types': {}, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + chatStream: (options) => + (async function* () { + calls += 1 + const now = Date.now() + const messageId = 'm-' + calls + yield { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now } + yield { type: EventType.TEXT_MESSAGE_START, messageId, role: 'assistant', timestamp: now } + yield { type: EventType.TEXT_MESSAGE_CONTENT, messageId, delta: 'worker saw ' + options.messages.length + ' messages', timestamp: now } + yield { type: EventType.TEXT_MESSAGE_END, messageId, timestamp: now } + yield { type: EventType.RUN_FINISHED, runId: 'r', threadId: 't', timestamp: now, metadata: { tanstack: { finishReason: 'stop' } } } + })(), +} + +export const studio = defineHarness({ + name: 'acme/built', + adapter, + agents: [defineAgent({ name: 'painter', description: 'Paints', produces: 'image', run: async () => 'art' })], +}) +`, + ) + + const outDir = join(tmp, 'out') + const { manifest, bundle } = await buildHarness({ + entry, + export: 'studio', + outDir, + }) + expect(manifest).toMatchObject({ + format: 'tanstack-ai-harness', + version: 1, + name: 'acme/built', + entry: 'harness.js', + agents: [{ name: 'painter', produces: 'image' }], + requires: { + filesystem: false, + processExecution: false, + network: 'model-only', + }, + }) + expect(manifest.digest).toMatch(/^[0-9a-f]{64}$/) + expect(await readManifest(outDir)).toEqual(manifest) + expect((await readFile(bundle)).length).toBeGreaterThan(1000) + + const model = await artifactText(outDir) + try { + const ask = (content: string) => + textOf( + chat({ + adapter: model, + messages: [{ role: 'user', content }], + threadId: 'outer', + }) as AsyncIterable, + ) + expect(await ask('first')).toBe('worker saw 1 messages') + // The worker keeps the conversation between calls. + expect(await ask('second')).toBe('worker saw 3 messages') + } finally { + model.dispose() + } + }, + ) + + it( + 'refuses a bundle that does not match its manifest', + { timeout: 60_000 }, + async () => { + const outDir = join(tmp, 'out') + await writeFile(join(outDir, 'harness.js'), '// changed\n') + await expect(artifactText(outDir)).rejects.toThrow( + 'does not match its manifest digest', + ) + }, + ) +}) diff --git a/packages/ai-harness/vite.config.ts b/packages/ai-harness/vite.config.ts index c99a69e684..e0d86be566 100644 --- a/packages/ai-harness/vite.config.ts +++ b/packages/ai-harness/vite.config.ts @@ -28,7 +28,9 @@ const config = defineConfig({ export default mergeConfig( config, tanstackViteConfig({ - entry: ['./src/index.ts', './src/client.ts', './src/first-party/index.ts'], + entry: ['./src/index.ts', './src/client.ts', './src/first-party/index.ts', + './src/build.ts', + './src/worker.ts',], srcDir: './src', cjs: false, }), diff --git a/testing/e2e/src/routes/api.harness-test.ts b/testing/e2e/src/routes/api.harness-test.ts index b66add840d..ff75aba9b1 100644 --- a/testing/e2e/src/routes/api.harness-test.ts +++ b/testing/e2e/src/routes/api.harness-test.ts @@ -1,6 +1,10 @@ import { createFileRoute } from '@tanstack/react-router' -import { defineAgent } from '@tanstack/ai' -import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { chat, defineAgent } from '@tanstack/ai' +import { + createHarnessHost, + defineHarness, + harnessText, +} from '@tanstack/ai-harness' import { memoryPersistence } from '@tanstack/ai-persistence' import { z } from 'zod' import { createTextAdapter } from '@/lib/providers' @@ -48,6 +52,23 @@ export const Route = createFileRoute('/api/harness-test')({ const host = createHarnessHost({ persistence }) try { const session = await host.open(harness, { threadId: 'e2e-thread' }) + if (body.scenario === 'remote') { + // A harness on "another machine": the protocol route of this app, + // called over HTTP. + const remote = harnessText({ + url: new URL('/api/harness-protocol', request.url).href, + token: 'e2e-token', + }) + let answer = '' + for await (const chunk of chat({ + adapter: remote, + messages: [{ role: 'user', content: '[harness-protocol] hello' }], + threadId: `remote-${testId ?? 'default'}`, + })) { + if (chunk.type === 'TEXT_MESSAGE_CONTENT') answer += chunk.delta + } + return Response.json({ answer }) + } if (body.scenario === 'limits') { // The main model calls `worker` twice. The limit allows one. let workerRuns = 0 diff --git a/testing/e2e/tests/harness.spec.ts b/testing/e2e/tests/harness.spec.ts index 609a642592..667f53389d 100644 --- a/testing/e2e/tests/harness.spec.ts +++ b/testing/e2e/tests/harness.spec.ts @@ -15,6 +15,20 @@ test.describe('harness session', () => { expect(body.roles).toEqual(['user', 'assistant', 'user', 'assistant']) }) + test('a remote harness answers as the model of a chat call', async ({ + request, + testId, + aimockPort, + }) => { + const response = await request.post('/api/harness-test', { + data: { scenario: 'remote', testId, aimockPort }, + }) + expect(response.ok()).toBe(true) + expect((await response.json()).answer).toBe( + 'Hello over the harness protocol.', + ) + }) + test('subagent limits refuse the second child start', async ({ request, testId, From 5d27a1d638fad2b69323f716ce266a62b5a6427c Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 11:24:41 +0000 Subject: [PATCH 012/286] ci: apply automated fixes --- packages/ai-harness/src/build.ts | 4 +--- packages/ai-harness/vite.config.ts | 8 ++++++-- .../ai-persistence/src/testkit/conformance.ts | 17 ++++++----------- 3 files changed, 13 insertions(+), 16 deletions(-) diff --git a/packages/ai-harness/src/build.ts b/packages/ai-harness/src/build.ts index 4d3cf074e8..6ebfb5b5c2 100644 --- a/packages/ai-harness/src/build.ts +++ b/packages/ai-harness/src/build.ts @@ -95,9 +95,7 @@ function run( * The build imports `entry` in a child process to read the harness name, * agents, and plugin names. Importing a module runs its code. */ -export async function buildHarness( - options: BuildHarnessOptions, -): Promise<{ +export async function buildHarness(options: BuildHarnessOptions): Promise<{ manifest: HarnessManifestV1 bundle: string executable?: string diff --git a/packages/ai-harness/vite.config.ts b/packages/ai-harness/vite.config.ts index e0d86be566..b5a782675a 100644 --- a/packages/ai-harness/vite.config.ts +++ b/packages/ai-harness/vite.config.ts @@ -28,9 +28,13 @@ const config = defineConfig({ export default mergeConfig( config, tanstackViteConfig({ - entry: ['./src/index.ts', './src/client.ts', './src/first-party/index.ts', + entry: [ + './src/index.ts', + './src/client.ts', + './src/first-party/index.ts', './src/build.ts', - './src/worker.ts',], + './src/worker.ts', + ], srcDir: './src', cjs: false, }), diff --git a/packages/ai-persistence/src/testkit/conformance.ts b/packages/ai-persistence/src/testkit/conformance.ts index 27b6d0e874..3b8ab62ebf 100644 --- a/packages/ai-persistence/src/testkit/conformance.ts +++ b/packages/ai-persistence/src/testkit/conformance.ts @@ -1583,18 +1583,13 @@ export function runPersistenceConformance( const current = await store.getVersioned('cas', 'k') expect(current?.value).toEqual({ n: 1 }) const revision = current?.revision ?? null - const updated = await store.setIf( - 'cas', - 'k', - { n: 2 }, - revision, - ) + const updated = await store.setIf('cas', 'k', { n: 2 }, revision) expect(updated.ok).toBe(true) - expect( - await store.setIf('cas', 'k', { n: 3 }, revision), - ).toMatchObject({ - ok: false, - }) + expect(await store.setIf('cas', 'k', { n: 3 }, revision)).toMatchObject( + { + ok: false, + }, + ) expect(await store.get('cas', 'k')).toEqual({ n: 2 }) }) }) From 61ce0041fabed895ecdc7037cc51069ae76dfea2 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 13:36:41 +0200 Subject: [PATCH 013/286] feat(ai-dashboard): self-hosted dashboard, --dashboard, and a runnable example New package @tanstack/ai-dashboard. `npx @tanstack/ai-dashboard` (or startDashboard) runs a node:http server with no new dependencies. Hosts dial out with connectDashboard, pair with a one-time code, and get a revocable host token. The relay uses SSE plus POST and caches recent events per session; inputs for an offline host wait and are delivered on reconnect. The web app lists hosts and sessions, streams messages and tool calls, shows approval cards and plugin questions, and sends prompts, steers, and stops. It installs as a PWA on a phone. @tanstack/ai-harness-cli adds --dashboard . examples/harness-cli: a small coding agent (permissions, workspace tools, todos, model picker, typed agent) that runs with OpenAI, Anthropic, or a demo model. --- .changeset/harness-p6-dashboard.md | 8 + docs/config.json | 5 + docs/harness/dashboard.md | 69 ++++ examples/README.md | 17 + examples/harness-cli/README.md | 33 ++ examples/harness-cli/package.json | 26 ++ examples/harness-cli/playground/.gitkeep | 0 examples/harness-cli/playground/AGENTS.md | 1 + examples/harness-cli/src/cli.ts | 4 + examples/harness-cli/src/harness.ts | 136 +++++++ examples/harness-cli/tsconfig.json | 13 + packages/ai-dashboard/README.md | 56 +++ packages/ai-dashboard/bin/dashboard.mjs | 26 ++ packages/ai-dashboard/package.json | 85 +++++ packages/ai-dashboard/src/connect.ts | 242 ++++++++++++ packages/ai-dashboard/src/index.ts | 2 + packages/ai-dashboard/src/server.ts | 444 ++++++++++++++++++++++ packages/ai-dashboard/src/ui.ts | 194 ++++++++++ packages/ai-dashboard/tests/relay.test.ts | 219 +++++++++++ packages/ai-dashboard/tsconfig.json | 8 + packages/ai-dashboard/vite.config.ts | 35 ++ packages/ai-harness-cli/package.json | 5 + packages/ai-harness-cli/src/args.ts | 4 + packages/ai-harness-cli/src/index.ts | 32 ++ packages/ai-harness-cli/src/lines.ts | 34 +- pnpm-lock.yaml | 58 +++ testing/e2e/package.json | 1 + testing/e2e/tests/dashboard.spec.ts | 124 ++++++ 28 files changed, 1867 insertions(+), 14 deletions(-) create mode 100644 .changeset/harness-p6-dashboard.md create mode 100644 docs/harness/dashboard.md create mode 100644 examples/harness-cli/README.md create mode 100644 examples/harness-cli/package.json create mode 100644 examples/harness-cli/playground/.gitkeep create mode 100644 examples/harness-cli/playground/AGENTS.md create mode 100644 examples/harness-cli/src/cli.ts create mode 100644 examples/harness-cli/src/harness.ts create mode 100644 examples/harness-cli/tsconfig.json create mode 100644 packages/ai-dashboard/README.md create mode 100644 packages/ai-dashboard/bin/dashboard.mjs create mode 100644 packages/ai-dashboard/package.json create mode 100644 packages/ai-dashboard/src/connect.ts create mode 100644 packages/ai-dashboard/src/index.ts create mode 100644 packages/ai-dashboard/src/server.ts create mode 100644 packages/ai-dashboard/src/ui.ts create mode 100644 packages/ai-dashboard/tests/relay.test.ts create mode 100644 packages/ai-dashboard/tsconfig.json create mode 100644 packages/ai-dashboard/vite.config.ts create mode 100644 testing/e2e/tests/dashboard.spec.ts diff --git a/.changeset/harness-p6-dashboard.md b/.changeset/harness-p6-dashboard.md new file mode 100644 index 0000000000..dd71ffb871 --- /dev/null +++ b/.changeset/harness-p6-dashboard.md @@ -0,0 +1,8 @@ +--- +'@tanstack/ai-dashboard': minor +'@tanstack/ai-harness-cli': minor +--- + +New package `@tanstack/ai-dashboard`: a self-hosted dashboard for harness sessions. `npx @tanstack/ai-dashboard` starts it (`startDashboard` in code). Agents dial out with `connectDashboard` from `@tanstack/ai-dashboard/connect`, pair with a one-time code, and get a revocable host token. The web app lists hosts and sessions, streams messages and tool calls, shows approval cards and plugin questions, and sends prompts, steers, and stops. It installs as a web app on a phone. Inputs for an offline host wait at the dashboard and reach it on reconnect. + +`@tanstack/ai-harness-cli` adds `--dashboard `, which pairs on first use (or reads `HARNESS_DASHBOARD_TOKEN`) and keeps the session connected. diff --git a/docs/config.json b/docs/config.json index 533ad81083..c8e087348a 100644 --- a/docs/config.json +++ b/docs/config.json @@ -848,6 +848,11 @@ "label": "Deploy a harness", "to": "harness/deploy", "addedAt": "2026-09-26" + }, + { + "label": "Self-host the dashboard", + "to": "harness/dashboard", + "addedAt": "2026-09-26" } ] }, diff --git a/docs/harness/dashboard.md b/docs/harness/dashboard.md new file mode 100644 index 0000000000..8bc06bcd83 --- /dev/null +++ b/docs/harness/dashboard.md @@ -0,0 +1,69 @@ +--- +title: Self-host the dashboard +id: harness-dashboard +order: 10 +description: "Watch and steer your harness sessions from a browser or a phone. Agents dial out to your dashboard server, so they need no open port." +keywords: + - tanstack ai + - harness + - dashboard + - mobile + - relay +--- + +Your agent runs on a laptop, a server, or a CI runner, and it stops to ask for approval while you are away from that terminal. The dashboard shows every connected session in a browser or on your phone. You approve tools, answer questions, and send prompts from there. Agents connect out to the dashboard, so they need no open port. + +## 1. Start the dashboard + +```bash +npx @tanstack/ai-dashboard --port 8790 +``` + +It prints a sign-in link that holds the owner token. Set `DASHBOARD_OWNER_TOKEN` to keep the same token across restarts. Put the dashboard behind TLS when it is reachable from other machines. + +## 2. Connect an agent + +With the CLI, add `--dashboard`: + +```bash +npx tsx cli.ts --dashboard http://127.0.0.1:8790 +``` + +In your own host, call `connectDashboard`: + +```ts group=harness-dashboard +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { connectDashboard } from '@tanstack/ai-dashboard/connect' +import { openaiText } from '@tanstack/ai-openai' + +const studio = defineHarness({ name: 'acme/studio', adapter: openaiText('gpt-5.6') }) +const host = createHarnessHost() + +const connection = await connectDashboard({ + host, + harness: studio, + url: 'http://127.0.0.1:8790', + threads: ['main'], + onPairingCode: (code) => console.log(`Approve ${code} in the dashboard`), + onToken: (token) => console.log(`Save this host token: ${token}`), +}) +``` + +## 3. Pair it + +The first time, the agent prints a pairing code. Approve the code in the dashboard under "Pairing requests". The agent gets a host token. Pass it as `token` (or `HARNESS_DASHBOARD_TOKEN` for the CLI) to skip pairing next time. Revoke a host in the dashboard to cut it off. + +## What you can do in the dashboard + +- See hosts (online or offline) and sessions (running, waiting for you, idle). +- Read the messages and tool calls of a session as they stream. +- Approve or reject tool calls, and answer questions from plugins. +- Send a prompt, steer a running turn, or stop it. + +Inputs for an offline host wait at the dashboard for 10 minutes and reach the host when it reconnects. Set `allowRemoteStart: true` in `connectDashboard` to let the dashboard open new sessions on the host. + +On a phone, open the dashboard and add it to the home screen. It works as an installed web app. + +## What you have now + +- Your sessions on any screen, with approvals and prompts a tap away. diff --git a/examples/README.md b/examples/README.md index b050e1b370..9406c3054e 100644 --- a/examples/README.md +++ b/examples/README.md @@ -11,6 +11,7 @@ Choose an example based on your use case: - **Want a table that fills as JSON streams?** → [Streaming Structured Table](#streaming-structured-table) - **Want typed ticket routing?** → [Evaluate](#evaluate) - **Want a blog desk with three agents?** → [Subagents](#subagents) +- **Want your own coding agent in the terminal?** → [Harness CLI](#harness-cli) - **Want a full-stack TypeScript app?** → [TanStack Chat (ts-react-chat)](#tanstack-chat-ts-react-chat) - **Need a live world stream?** → [World generation (ts-react-media)](#world-generation-ts-react-media) - **Need a vanilla JS frontend?** → [Vanilla Chat](#vanilla-chat) @@ -19,6 +20,22 @@ Choose an example based on your use case: ## TypeScript Examples +### Harness CLI + +A small coding agent in the terminal, built with `@tanstack/ai-harness`. It edits files in `./playground`, asks before writes and commands, keeps a todo list, switches models, and connects to the self-hosted dashboard. Without an API key, a demo model answers. + +**Path:** `examples/harness-cli` + +**Getting Started:** + +```bash +OPENAI_API_KEY=... pnpm --filter harness-cli-example start +``` + +📖 [Full Documentation](harness-cli/README.md) + +--- + ### Basic Chat A slim TanStack Start chat. One provider (OpenRouter). You paste your own key. diff --git a/examples/harness-cli/README.md b/examples/harness-cli/README.md new file mode 100644 index 0000000000..cd5f93c5fe --- /dev/null +++ b/examples/harness-cli/README.md @@ -0,0 +1,33 @@ +# Harness CLI example + +A small coding agent in your terminal, built with `@tanstack/ai-harness`. It reads and edits files in `./playground`, asks before it writes a file or runs a command, keeps a todo list, and can switch models. + +## Run it + +From the repo root: + +1. `pnpm install` +2. `pnpm build:all` (the example uses the local packages) +3. Set a key: `OPENAI_API_KEY=...` or `ANTHROPIC_API_KEY=...`. Without one, a demo model answers. +4. `pnpm --filter harness-cli-example start` + +Try these: + +- `create hello.txt with a short poem`: the agent asks before `write_file`. Type `y`. +- `/mode plan`, then ask for another file: the write is refused. +- `/todos`, `/usage`, `/model fast`, `/agents` +- `/agent haiku {"topic":"rain"}`: runs a typed agent in the background. +- Press Esc to stop a long answer. Type while it works to steer it. + +## Other modes + +- One prompt for scripts and CI: `pnpm --filter harness-cli-example start -p "list the files"` +- AG-UI events as JSON lines: add `--output ndjson`. +- As an editor agent (ACP): `pnpm --filter harness-cli-example start --acp` +- As an HTTP server: `pnpm --filter harness-cli-example start --serve` (it prints a token). + +## Watch it from a browser or a phone + +1. In one terminal: `pnpm --filter harness-cli-example dashboard`. It prints a sign-in link. +2. In another: `pnpm --filter harness-cli-example start --dashboard http://127.0.0.1:8790`. It prints a pairing code. +3. Open the sign-in link, approve the code, and open the `main` session. Messages you send there reach the agent. diff --git a/examples/harness-cli/package.json b/examples/harness-cli/package.json new file mode 100644 index 0000000000..3038efe1ff --- /dev/null +++ b/examples/harness-cli/package.json @@ -0,0 +1,26 @@ +{ + "name": "harness-cli-example", + "private": true, + "type": "module", + "description": "A small coding agent in the terminal, built with @tanstack/ai-harness", + "scripts": { + "start": "tsx src/cli.ts", + "dashboard": "tanstack-ai-dashboard", + "test:types": "tsc --noEmit" + }, + "dependencies": { + "@tanstack/ai": "workspace:*", + "@tanstack/ai-acp": "workspace:*", + "@tanstack/ai-anthropic": "workspace:*", + "@tanstack/ai-dashboard": "workspace:*", + "@tanstack/ai-harness": "workspace:*", + "@tanstack/ai-harness-cli": "workspace:*", + "@tanstack/ai-openai": "workspace:*", + "@tanstack/ai-persistence": "workspace:*", + "zod": "^4.2.0" + }, + "devDependencies": { + "@types/node": "^24.10.1", + "tsx": "^4.23.12" + } +} diff --git a/examples/harness-cli/playground/.gitkeep b/examples/harness-cli/playground/.gitkeep new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/harness-cli/playground/AGENTS.md b/examples/harness-cli/playground/AGENTS.md new file mode 100644 index 0000000000..9fcc63b5ea --- /dev/null +++ b/examples/harness-cli/playground/AGENTS.md @@ -0,0 +1 @@ +This is a playground folder. Keep files short and add a comment at the top of every new file. diff --git a/examples/harness-cli/src/cli.ts b/examples/harness-cli/src/cli.ts new file mode 100644 index 0000000000..774a33ba2e --- /dev/null +++ b/examples/harness-cli/src/cli.ts @@ -0,0 +1,4 @@ +import { runCli } from '@tanstack/ai-harness-cli' +import { assistant } from './harness' + +process.exitCode = await runCli(assistant) diff --git a/examples/harness-cli/src/harness.ts b/examples/harness-cli/src/harness.ts new file mode 100644 index 0000000000..c87b1a6434 --- /dev/null +++ b/examples/harness-cli/src/harness.ts @@ -0,0 +1,136 @@ +import { fileURLToPath } from 'node:url' +import { EventType, defineAgent } from '@tanstack/ai' +import { defineHarness } from '@tanstack/ai-harness' +import { + compact, + modelPicker, + permissions, + projectInstructions, + todos, + usage, + workspaceTools, +} from '@tanstack/ai-harness/plugins' +import { anthropicText } from '@tanstack/ai-anthropic' +import { openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' +import type { AnyTextAdapter } from '@tanstack/ai' + +// The agent works in ./playground, so it cannot touch the rest of your disk. +const root = fileURLToPath(new URL('../playground', import.meta.url)) + +/** Without an API key, a stand-in model that explains how to add one. */ +function demoModel(): AnyTextAdapter { + let calls = 0 + return { + kind: 'text', + name: 'demo', + model: 'demo', + '~types': { + providerOptions: {}, + inputModalities: ['text'], + messageMetadataByModality: { + text: undefined, + image: undefined, + audio: undefined, + video: undefined, + document: undefined, + }, + toolCapabilities: [], + toolCallMetadata: undefined, + systemPromptMetadata: undefined as never, + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + chatStream: (options) => + (async function* () { + calls += 1 + const messageId = `demo-${calls}` + const last = options.messages.at(-1) + const said = typeof last?.content === 'string' ? last.content : '' + const now = Date.now() + yield { + type: EventType.RUN_STARTED, + runId: 'demo', + threadId: 'demo', + timestamp: now, + } + yield { + type: EventType.TEXT_MESSAGE_START, + messageId, + role: 'assistant', + timestamp: now, + } + yield { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId, + delta: `(demo model) You said: "${said}". Set OPENAI_API_KEY or ANTHROPIC_API_KEY to talk to a real model.`, + timestamp: now, + } + yield { type: EventType.TEXT_MESSAGE_END, messageId, timestamp: now } + yield { + type: EventType.RUN_FINISHED, + runId: 'demo', + threadId: 'demo', + timestamp: now, + metadata: { tanstack: { finishReason: 'stop' } }, + } + })(), + } +} + +function pickModels(): { + main: AnyTextAdapter + choices: Record +} { + if (process.env.OPENAI_API_KEY) { + const smart = openaiText('gpt-5.6') + return { main: smart, choices: { smart, fast: openaiText('gpt-5.6-luna') } } + } + if (process.env.ANTHROPIC_API_KEY) { + const smart = anthropicText('claude-opus-5-5') + return { + main: smart, + choices: { smart, fast: anthropicText('claude-sonnet-5') }, + } + } + const demo = demoModel() + return { main: demo, choices: { demo } } +} + +const { main, choices } = pickModels() + +/** A typed agent you can run with `/agent haiku {"topic":"rain"}`. */ +const haiku = defineAgent({ + name: 'haiku', + description: 'Writes a haiku about a topic', + inputSchema: z.object({ topic: z.string() }), + run: (ctx) => + ctx.chat({ + adapter: main, + messages: [ + { + role: 'user', + content: `Write one haiku about ${ctx.input.topic}. Only the haiku.`, + }, + ], + stream: false, + }), +}) + +export const assistant = defineHarness({ + name: 'example/coder', + description: 'A small coding agent that works in ./playground', + adapter: main, + systemPrompts: [ + 'You are a careful coding agent. Read files before you edit them. Keep answers short.', + ], + agents: [haiku], + plugins: () => [ + permissions(), + workspaceTools({ root }), + todos(), + modelPicker({ choices, default: Object.keys(choices)[0] }), + projectInstructions({ root }), + compact({ adapter: main }), + usage(), + ], +}) diff --git a/examples/harness-cli/tsconfig.json b/examples/harness-cli/tsconfig.json new file mode 100644 index 0000000000..619816780b --- /dev/null +++ b/examples/harness-cli/tsconfig.json @@ -0,0 +1,13 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "ESNext", + "moduleResolution": "Bundler", + "strict": true, + "noEmit": true, + "skipLibCheck": true, + "types": ["node"], + "esModuleInterop": true + }, + "include": ["src"] +} diff --git a/packages/ai-dashboard/README.md b/packages/ai-dashboard/README.md new file mode 100644 index 0000000000..be3f0a764b --- /dev/null +++ b/packages/ai-dashboard/README.md @@ -0,0 +1,56 @@ +
+ + + + TanStack AI + +
+ +
+ +# @tanstack/ai-dashboard + +A self-hosted dashboard for TanStack AI harnesses. Watch sessions, approve tool calls, answer questions, and send prompts from a browser or a phone. Agents dial out to it, so they need no open port. + +## Start the dashboard + +```bash +npx @tanstack/ai-dashboard --port 8790 +``` + +It prints a sign-in link with the owner token. Set `DASHBOARD_OWNER_TOKEN` to keep the same token across restarts. + +## Connect an agent + +```ts +import { connectDashboard } from '@tanstack/ai-dashboard/connect' + +await connectDashboard({ + host, + harness: studio, + url: 'http://127.0.0.1:8790', + threads: ['main'], + onPairingCode: (code) => console.log(`Approve ${code} in the dashboard`), + onToken: (token) => console.log(`Save this host token: ${token}`), +}) +``` + +With `@tanstack/ai-harness-cli`, run your CLI with `--dashboard `. + +## Documentation + +Read [Self-host the dashboard](https://tanstack.com/ai/latest/docs/harness/dashboard). + +## License + +MIT diff --git a/packages/ai-dashboard/bin/dashboard.mjs b/packages/ai-dashboard/bin/dashboard.mjs new file mode 100644 index 0000000000..bf9a07d52e --- /dev/null +++ b/packages/ai-dashboard/bin/dashboard.mjs @@ -0,0 +1,26 @@ +#!/usr/bin/env node +// Start a dashboard: npx @tanstack/ai-dashboard [--port 8790] [--host 127.0.0.1] +import { parseArgs } from 'node:util' +import { startDashboard } from '../dist/esm/index.js' + +const { values } = parseArgs({ + options: { port: { type: 'string' }, host: { type: 'string' } }, +}) +const dashboard = await startDashboard({ + port: Number(values.port ?? process.env.PORT ?? 8790), + hostname: values.host ?? process.env.HOST ?? '127.0.0.1', + ...(process.env.DASHBOARD_OWNER_TOKEN + ? { ownerToken: process.env.DASHBOARD_OWNER_TOKEN } + : {}), +}) +console.log(`Dashboard: ${dashboard.url}`) +if (!process.env.DASHBOARD_OWNER_TOKEN) { + console.log(`Owner token: ${dashboard.ownerToken}`) + console.log(`Sign in: ${dashboard.url}/#token=${dashboard.ownerToken}`) +} +const stop = async () => { + await dashboard.close() + process.exit(0) +} +process.once('SIGINT', stop) +process.once('SIGTERM', stop) diff --git a/packages/ai-dashboard/package.json b/packages/ai-dashboard/package.json new file mode 100644 index 0000000000..5e096733d6 --- /dev/null +++ b/packages/ai-dashboard/package.json @@ -0,0 +1,85 @@ +{ + "name": "@tanstack/ai-dashboard", + "version": "0.0.0", + "description": "A self-hosted dashboard for TanStack AI harnesses: watch sessions, approve tools, and send prompts from a browser or a phone", + "author": "", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/TanStack/ai.git", + "directory": "packages/ai-dashboard" + }, + "type": "module", + "module": "./dist/esm/index.js", + "types": "./dist/esm/index.d.ts", + "bin": { + "tanstack-ai-dashboard": "./bin/dashboard.mjs" + }, + "exports": { + ".": { + "types": "./dist/esm/index.d.ts", + "import": "./dist/esm/index.js" + }, + "./connect": { + "types": "./dist/esm/connect.d.ts", + "import": "./dist/esm/connect.js" + } + }, + "sideEffects": false, + "files": [ + "dist", + "src", + "bin" + ], + "scripts": { + "build": "vite build", + "clean": "premove ./build ./dist", + "lint:fix": "oxlint src --type-aware --fix", + "test:build": "publint --strict", + "test:oxlint": "oxlint src --type-aware", + "test:coverage": "vitest run --passWithNoTests --coverage --coverage.include='src/**' --coverage.reporter=text-summary --coverage.reporter=json-summary", + "test:lib": "vitest --passWithNoTests", + "test:lib:dev": "pnpm test:lib --watch", + "test:types": "tsc" + }, + "nx": { + "targets": { + "test:lib": { + "dependsOn": [ + "build", + "^build" + ] + }, + "test:coverage": { + "dependsOn": [ + "build", + "^build" + ] + }, + "test:types": { + "dependsOn": [ + "build", + "^build" + ] + } + } + }, + "keywords": [ + "ai", + "tanstack", + "harness", + "dashboard", + "relay", + "pwa" + ], + "peerDependencies": { + "@tanstack/ai": "workspace:^", + "@tanstack/ai-harness": "workspace:^" + }, + "devDependencies": { + "@tanstack/ai": "workspace:*", + "@tanstack/ai-harness": "workspace:*", + "@tanstack/ai-persistence": "workspace:*", + "@vitest/coverage-v8": "4.1.10" + } +} diff --git a/packages/ai-dashboard/src/connect.ts b/packages/ai-dashboard/src/connect.ts new file mode 100644 index 0000000000..07f4cd66ac --- /dev/null +++ b/packages/ai-dashboard/src/connect.ts @@ -0,0 +1,242 @@ +import { applyInput, parseHarnessInput } from '@tanstack/ai-harness' +import type { + AnyHarness, + HarnessHost, + HarnessSession, +} from '@tanstack/ai-harness' + +export interface ConnectDashboardOptions { + host: HarnessHost + harness: AnyHarness + /** The dashboard URL, for example `https://dash.example.com`. */ + url: string + /** The host token. Without one, the host pairs first (see `onPairingCode`). */ + token?: string + /** How this host shows up in the dashboard. */ + name?: string + /** Threads to show in the dashboard at once. */ + threads?: ReadonlyArray + /** Let the dashboard open new sessions on this host. Default false. */ + allowRemoteStart?: boolean + /** Called with the code to approve in the dashboard. */ + onPairingCode?: (code: string) => void + /** Called with the new host token after pairing. Save it to skip pairing next time. */ + onToken?: (token: string) => void + fetch?: typeof fetch + /** Wait between reconnects, doubled up to 30 s. Default 1000 ms. */ + reconnectDelayMs?: number +} + +const isRecord = (value: unknown): value is Record => + typeof value === 'object' && value !== null + +async function* sseData(response: Response): AsyncGenerator { + if (!response.body) return + const reader = response.body.getReader() + const decoder = new TextDecoder() + let buffer = '' + while (true) { + const { value, done } = await reader.read() + if (done) return + buffer += decoder.decode(value, { stream: true }) + const blocks = buffer.split('\n\n') + buffer = blocks.pop() ?? '' + for (const block of blocks) { + const data = block + .split('\n') + .find((line) => line.startsWith('data: ')) + ?.slice(6) + if (data) yield JSON.parse(data) + } + } +} + +/** + * Connect a harness host to a dashboard. The host dials out, so it needs no + * open port. The dashboard shows its sessions and sends inputs back. Every + * input goes through the same session API as a local client. + */ +export async function connectDashboard(options: ConnectDashboardOptions) { + const base = options.url.replace(/\/$/, '') + const doFetch = options.fetch ?? fetch + const stopped = new AbortController() + let token = options.token + + if (!token) { + const started = await ( + await doFetch(`${base}/api/pair/start`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ name: options.name ?? options.harness.name }), + }) + ).json() + if ( + !isRecord(started) || + typeof started.pairingId !== 'string' || + typeof started.code !== 'string' + ) { + throw new Error('The dashboard did not start a pairing.') + } + options.onPairingCode?.(started.code) + while (!token) { + if (stopped.signal.aborted) + throw new Error('Stopped before pairing finished.') + await new Promise((resolve) => setTimeout(resolve, 1000)) + const status: unknown = await ( + await doFetch( + `${base}/api/pair/status?pairingId=${encodeURIComponent(started.pairingId)}`, + ) + ).json() + if ( + isRecord(status) && + status.status === 'approved' && + typeof status.token === 'string' + ) { + token = status.token + } else if (!isRecord(status) || status.status !== 'pending') { + throw new Error('The pairing expired. Start again.') + } + } + options.onToken?.(token) + } + const hostToken = token + const headers = { + 'Content-Type': 'application/json', + Authorization: `Bearer ${hostToken}`, + } + + const post = (path: string, value: unknown) => + doFetch(`${base}${path}`, { + method: 'POST', + headers, + body: JSON.stringify(value), + }) + + await post('/api/host/hello', { + name: options.name ?? options.harness.name, + harnesses: [options.harness.name], + }) + + const attached = new Map() + const attach = async (threadId: string): Promise => { + const existing = attached.get(threadId) + if (existing) return existing + const session = await options.host.open(options.harness, { threadId }) + attached.set(threadId, session) + // Batch events so a streaming answer is a few requests, not one per token. + let pending: Array = [] + let timer: ReturnType | undefined + const flush = () => { + timer = undefined + const frames = pending + pending = [] + if (frames.length > 0) { + void post('/api/host/frames', { + threadId, + harness: options.harness.name, + frames, + }).catch(() => {}) + } + } + void (async () => { + for await (const entry of session.events({ + from: '0', + signal: stopped.signal, + })) { + pending.push({ type: 'harness.event', ...entry }) + timer ??= setTimeout(flush, 50) + } + })() + return session + } + for (const threadId of options.threads ?? []) await attach(threadId) + + const handle = async (envelope: unknown) => { + if ( + !isRecord(envelope) || + typeof envelope.threadId !== 'string' || + !isRecord(envelope.frame) + ) + return + const { threadId, frame } = envelope + if (frame.type === 'harness.subscribe') { + if (attached.has(threadId) || options.allowRemoteStart) + await attach(threadId) + return + } + const session = attached.get(threadId) + if (!session) { + if ( + frame.type === 'harness.input' && + typeof frame.requestId === 'string' + ) { + await post('/api/host/frames', { + threadId, + frames: [ + { + type: 'harness.receipt', + requestId: frame.requestId, + status: 'rejected', + reason: 'remote_start_disabled', + }, + ], + }) + } + return + } + if (frame.type === 'harness.input' && typeof frame.requestId === 'string') { + try { + const receipt = await applyInput( + options.harness, + session, + parseHarnessInput(frame.input), + ) + await post('/api/host/frames', { + threadId, + frames: [ + { type: 'harness.receipt', requestId: frame.requestId, ...receipt }, + ], + }) + } catch (error) { + await post('/api/host/frames', { + threadId, + frames: [ + { + type: 'harness.error', + message: error instanceof Error ? error.message : String(error), + }, + ], + }) + } + } + } + + // The control stream, with reconnects. + void (async () => { + let delay = options.reconnectDelayMs ?? 1000 + while (!stopped.signal.aborted) { + try { + const response = await doFetch(`${base}/api/host/stream`, { + headers, + signal: stopped.signal, + }) + if (response.status === 401) + throw new Error('The dashboard refused this host token.') + delay = options.reconnectDelayMs ?? 1000 + for await (const envelope of sseData(response)) void handle(envelope) + } catch (error) { + if (stopped.signal.aborted) return + if (error instanceof Error && error.message.includes('refused')) return + } + await new Promise((resolve) => setTimeout(resolve, delay)) + delay = Math.min(delay * 2, 30_000) + } + })() + + return { + token: hostToken, + /** Show another thread in the dashboard. */ + attach: (threadId: string) => attach(threadId).then(() => undefined), + close: () => stopped.abort(), + } +} diff --git a/packages/ai-dashboard/src/index.ts b/packages/ai-dashboard/src/index.ts new file mode 100644 index 0000000000..2d67dee7c8 --- /dev/null +++ b/packages/ai-dashboard/src/index.ts @@ -0,0 +1,2 @@ +export { startDashboard } from './server' +export type { ControlEnvelope, DashboardOptions } from './server' diff --git a/packages/ai-dashboard/src/server.ts b/packages/ai-dashboard/src/server.ts new file mode 100644 index 0000000000..9d561448bd --- /dev/null +++ b/packages/ai-dashboard/src/server.ts @@ -0,0 +1,444 @@ +import { createServer } from 'node:http' +import { randomBytes, randomInt, timingSafeEqual } from 'node:crypto' +import { DASHBOARD_HTML, MANIFEST_JSON, SERVICE_WORKER } from './ui' +import type { IncomingMessage, ServerResponse } from 'node:http' + +/** One frame the dashboard sends a host: a session-tier control frame for a thread. */ +export interface ControlEnvelope { + threadId: string + frame: Record +} + +export interface DashboardOptions { + port?: number + hostname?: string + /** The owner's login token. Default: a new random token, returned by `startDashboard`. */ + ownerToken?: string + /** Events kept per session for replay. Default 2000. */ + cacheSize?: number + /** How long an input waits for an offline host. Default 10 minutes. */ + queueTtlMs?: number +} + +interface HostRecord { + hostId: string + name: string + harnesses: Array + token: string + lastSeen: number + streams: Set + queue: Array<{ envelope: ControlEnvelope; expiresAt: number }> +} + +interface SessionRecord { + hostId: string + threadId: string + harness?: string + events: Array<{ + cursor: string + operationId: string + event: Record + }> + lastActivity: number + status: 'idle' | 'running' | 'waiting' + clients: Set +} + +const token = () => randomBytes(24).toString('base64url') + +function sameSecret( + given: string | null | undefined, + expected: string, +): boolean { + if (!given) return false + const a = Buffer.from(given) + const b = Buffer.from(expected) + return a.length === b.length && timingSafeEqual(a, b) +} + +function bearer(req: IncomingMessage, url: URL): string | null { + const header = req.headers.authorization + if (typeof header === 'string' && header.startsWith('Bearer ')) + return header.slice(7) + // EventSource cannot send headers, so SSE clients pass the token in the URL. + return url.searchParams.get('token') +} + +async function body(req: IncomingMessage): Promise> { + const chunks: Array = [] + let size = 0 + for await (const chunk of req) { + size += (chunk as Buffer).length + if (size > 5_000_000) throw new Error('Body too large.') + chunks.push(chunk as Buffer) + } + if (chunks.length === 0) return {} + const parsed: unknown = JSON.parse(Buffer.concat(chunks).toString('utf8')) + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) + throw new Error('Expected a JSON object.') + // Checked to be a plain object above. + return parsed as Record +} + +function json(res: ServerResponse, status: number, value: unknown): void { + res.writeHead(status, { + 'Content-Type': 'application/json', + 'Cache-Control': 'no-store', + }) + res.end(JSON.stringify(value)) +} + +function openSse(res: ServerResponse): void { + res.writeHead(200, { + 'Content-Type': 'text/event-stream', + 'Cache-Control': 'no-cache', + Connection: 'keep-alive', + }) + res.write(': connected\n\n') +} + +function sse(res: ServerResponse, data: unknown, id?: string): void { + res.write(`${id ? `id: ${id}\n` : ''}data: ${JSON.stringify(data)}\n\n`) +} + +/** A code a person can read out: 3 characters, a dash, 3 characters. */ +function pairingCode(): string { + const alphabet = 'ABCDEFGHJKLMNPQRSTUVWXYZ23456789' + const pick = () => alphabet.charAt(randomInt(alphabet.length)) + return `${pick()}${pick()}${pick()}-${pick()}${pick()}${pick()}` +} + +/** + * Start a self-hosted dashboard server. Harness hosts dial out to it + * (`connectDashboard`), and people open it in a browser or on a phone. + * + * - Hosts pair with a one-time code that the owner approves, and get a host token. + * - The owner logs in with the owner token. + * - The server relays session frames and keeps a cache of recent events. + * Hosts keep the transcripts. + */ +export async function startDashboard(options: DashboardOptions = {}) { + const ownerToken = options.ownerToken ?? token() + const cacheSize = options.cacheSize ?? 2000 + const queueTtlMs = options.queueTtlMs ?? 10 * 60_000 + const hosts = new Map() + const hostByToken = new Map() + const pairings = new Map< + string, + { code: string; name: string; expiresAt: number; hostId?: string } + >() + const sessions = new Map() + const sessionKey = (hostId: string, threadId: string) => + `${hostId}\u0000${threadId}` + + const isOwner = (req: IncomingMessage, url: URL) => + sameSecret(bearer(req, url), ownerToken) + const hostFor = (req: IncomingMessage, url: URL): HostRecord | undefined => { + const given = bearer(req, url) + const hostId = given ? hostByToken.get(given) : undefined + const host = hostId ? hosts.get(hostId) : undefined + return host && sameSecret(given, host.token) ? host : undefined + } + + const sessionFor = (hostId: string, threadId: string): SessionRecord => { + const key = sessionKey(hostId, threadId) + let session = sessions.get(key) + if (!session) { + session = { + hostId, + threadId, + events: [], + lastActivity: Date.now(), + status: 'idle', + clients: new Set(), + } + sessions.set(key, session) + } + return session + } + + const sendToHost = ( + host: HostRecord, + envelope: ControlEnvelope, + ): 'sent' | 'queued' => { + if (host.streams.size === 0) { + host.queue.push({ envelope, expiresAt: Date.now() + queueTtlMs }) + return 'queued' + } + for (const stream of host.streams) sse(stream, envelope) + return 'sent' + } + + const record = (session: SessionRecord, frame: Record) => { + session.lastActivity = Date.now() + if (frame.type === 'harness.event') { + const event = frame.event as Record + const name = event.type === 'CUSTOM' ? event.name : undefined + if (name === 'harness.operation.started') session.status = 'running' + if (name === 'harness.question') session.status = 'waiting' + if (name === 'harness.operation.finished') { + const value = event.value as { status?: unknown } | undefined + session.status = value?.status === 'interrupted' ? 'waiting' : 'idle' + } + session.events.push({ + cursor: String(frame.cursor), + operationId: String(frame.operationId), + event, + }) + if (session.events.length > cacheSize) + session.events.splice(0, session.events.length - cacheSize) + } + for (const client of session.clients) + sse( + client, + frame, + typeof frame.cursor === 'string' ? frame.cursor : undefined, + ) + } + + const route = async (req: IncomingMessage, res: ServerResponse) => { + const url = new URL(req.url ?? '/', 'http://dashboard') + const path = url.pathname + const method = req.method ?? 'GET' + + // The app shell. It holds no data; every API call needs a token. + if (method === 'GET' && (path === '/' || path === '/index.html')) { + res.writeHead(200, { 'Content-Type': 'text/html; charset=utf-8' }) + res.end(DASHBOARD_HTML) + return + } + if (method === 'GET' && path === '/manifest.webmanifest') { + res.writeHead(200, { 'Content-Type': 'application/manifest+json' }) + res.end(MANIFEST_JSON) + return + } + if (method === 'GET' && path === '/sw.js') { + res.writeHead(200, { 'Content-Type': 'text/javascript' }) + res.end(SERVICE_WORKER) + return + } + + // Pairing: a host asks, the owner approves, the host gets a token. + if (method === 'POST' && path === '/api/pair/start') { + const input = await body(req) + const pairingId = token() + const code = pairingCode() + pairings.set(pairingId, { + code, + name: + typeof input.name === 'string' ? input.name.slice(0, 100) : 'host', + expiresAt: Date.now() + 10 * 60_000, + }) + json(res, 200, { pairingId, code }) + return + } + if (method === 'GET' && path === '/api/pair/status') { + const pairing = pairings.get(url.searchParams.get('pairingId') ?? '') + if (!pairing || pairing.expiresAt < Date.now()) + return json(res, 404, { error: 'unknown or expired pairing' }) + if (!pairing.hostId) return json(res, 200, { status: 'pending' }) + const host = hosts.get(pairing.hostId) + pairings.delete(url.searchParams.get('pairingId') ?? '') + return json(res, 200, { + status: 'approved', + hostId: pairing.hostId, + token: host?.token, + }) + } + + if (path.startsWith('/api/host/')) { + const host = hostFor(req, url) + if (!host) return json(res, 401, { error: 'unauthorized' }) + host.lastSeen = Date.now() + if (method === 'GET' && path === '/api/host/stream') { + openSse(res) + host.streams.add(res) + const now = Date.now() + for (const queued of host.queue.splice(0)) { + if (queued.expiresAt > now) sse(res, queued.envelope) + } + const heartbeat = setInterval(() => res.write(': ping\n\n'), 15_000) + res.on('close', () => { + clearInterval(heartbeat) + host.streams.delete(res) + }) + return + } + if (method === 'POST' && path === '/api/host/hello') { + const input = await body(req) + if (typeof input.name === 'string') host.name = input.name.slice(0, 100) + if (Array.isArray(input.harnesses)) + host.harnesses = input.harnesses.map(String).slice(0, 50) + return json(res, 200, { hostId: host.hostId }) + } + if (method === 'POST' && path === '/api/host/frames') { + const input = await body(req) + if ( + typeof input.threadId !== 'string' || + !Array.isArray(input.frames) + ) { + return json(res, 400, { error: 'expected { threadId, frames }' }) + } + const session = sessionFor(host.hostId, input.threadId) + if (typeof input.harness === 'string') session.harness = input.harness + for (const frame of input.frames) { + if (typeof frame === 'object' && frame !== null) + record(session, frame as Record) + } + return json(res, 200, { ok: true }) + } + return json(res, 404, { error: 'not found' }) + } + + // Everything below is for the owner. + if (!path.startsWith('/api/')) return json(res, 404, { error: 'not found' }) + if (!isOwner(req, url)) return json(res, 401, { error: 'unauthorized' }) + + if (method === 'GET' && path === '/api/pairings') { + const now = Date.now() + return json( + res, + 200, + [...pairings.values()] + .filter((pairing) => !pairing.hostId && pairing.expiresAt > now) + .map((pairing) => ({ code: pairing.code, name: pairing.name })), + ) + } + if (method === 'POST' && path === '/api/pair/approve') { + const input = await body(req) + const entry = [...pairings.entries()].find( + ([, pairing]) => + pairing.code === input.code && + !pairing.hostId && + pairing.expiresAt > Date.now(), + ) + if (!entry) return json(res, 404, { error: 'unknown or expired code' }) + const hostId = `host-${randomBytes(6).toString('hex')}` + const hostToken = token() + hosts.set(hostId, { + hostId, + name: entry[1].name, + harnesses: [], + token: hostToken, + lastSeen: Date.now(), + streams: new Set(), + queue: [], + }) + hostByToken.set(hostToken, hostId) + entry[1].hostId = hostId + return json(res, 200, { hostId }) + } + if (method === 'POST' && path === '/api/hosts/revoke') { + const input = await body(req) + const host = hosts.get(String(input.hostId)) + if (!host) return json(res, 404, { error: 'unknown host' }) + hostByToken.delete(host.token) + for (const stream of host.streams) stream.end() + hosts.delete(host.hostId) + return json(res, 200, { ok: true }) + } + if (method === 'GET' && path === '/api/hosts') { + return json( + res, + 200, + [...hosts.values()].map((host) => ({ + hostId: host.hostId, + name: host.name, + harnesses: host.harnesses, + online: host.streams.size > 0, + lastSeen: host.lastSeen, + })), + ) + } + if (method === 'GET' && path === '/api/sessions') { + return json( + res, + 200, + [...sessions.values()] + .sort((a, b) => b.lastActivity - a.lastActivity) + .map((session) => ({ + hostId: session.hostId, + threadId: session.threadId, + harness: session.harness, + status: session.status, + lastActivity: session.lastActivity, + online: (hosts.get(session.hostId)?.streams.size ?? 0) > 0, + })), + ) + } + + const match = + /^\/api\/sessions\/([^/]+)\/([^/]+)\/(events|input|open)$/.exec(path) + if (match) { + const hostId = decodeURIComponent(match[1] ?? '') + const threadId = decodeURIComponent(match[2] ?? '') + const host = hosts.get(hostId) + if (!host) return json(res, 404, { error: 'unknown host' }) + if (method === 'GET' && match[3] === 'events') { + const session = sessionFor(hostId, threadId) + openSse(res) + const from = Number( + url.searchParams.get('from') ?? req.headers['last-event-id'] ?? 0, + ) + for (const entry of session.events) { + if (Number(entry.cursor) > from) + sse(res, { type: 'harness.event', ...entry }, entry.cursor) + } + session.clients.add(res) + const heartbeat = setInterval(() => res.write(': ping\n\n'), 15_000) + res.on('close', () => { + clearInterval(heartbeat) + session.clients.delete(res) + }) + return + } + if (method === 'POST' && match[3] === 'input') { + const input = await body(req) + const requestId = `dash-${randomBytes(6).toString('hex')}` + const delivery = sendToHost(host, { + threadId, + frame: { type: 'harness.input', requestId, input: input.input }, + }) + return json(res, 202, { requestId, status: delivery }) + } + if (method === 'POST' && match[3] === 'open') { + const delivery = sendToHost(host, { + threadId, + frame: { type: 'harness.subscribe', threadId }, + }) + return json(res, 202, { status: delivery }) + } + } + return json(res, 404, { error: 'not found' }) + } + + const server = createServer((req, res) => { + route(req, res).catch((error: unknown) => { + if (!res.headersSent) + json(res, 400, { + error: error instanceof Error ? error.message : String(error), + }) + else res.end() + }) + }) + await new Promise((resolve, reject) => { + server.once('error', reject) + server.listen(options.port ?? 8790, options.hostname ?? '127.0.0.1', () => + resolve(), + ) + }) + const address = server.address() + const port = + typeof address === 'object' && address + ? address.port + : (options.port ?? 8790) + return { + url: `http://${options.hostname ?? '127.0.0.1'}:${port}`, + ownerToken, + close: () => + new Promise((resolve) => { + server.closeAllConnections() + server.close(() => resolve()) + }), + } +} diff --git a/packages/ai-dashboard/src/ui.ts b/packages/ai-dashboard/src/ui.ts new file mode 100644 index 0000000000..4f0b2fd6ae --- /dev/null +++ b/packages/ai-dashboard/src/ui.ts @@ -0,0 +1,194 @@ +// The dashboard web app: one static page, no build step. It renders every +// value with textContent, never as HTML, so session content cannot inject +// markup. + +export const MANIFEST_JSON = JSON.stringify({ + name: 'TanStack AI Dashboard', + short_name: 'Harness', + start_url: '/', + display: 'standalone', + background_color: '#0b0d12', + theme_color: '#0b0d12', + icons: [], +}) + +// Caches the app shell only. API calls always go to the network. +export const SERVICE_WORKER = ` +const SHELL = ['/', '/manifest.webmanifest'] +self.addEventListener('install', (event) => { + event.waitUntil(caches.open('harness-shell-v1').then((cache) => cache.addAll(SHELL))) +}) +self.addEventListener('fetch', (event) => { + const url = new URL(event.request.url) + if (url.pathname.startsWith('/api/')) return + event.respondWith(fetch(event.request).catch(() => caches.match(event.request))) +}) +` + +const SCRIPT = ` +const $ = (tag, props = {}, ...children) => { + const node = document.createElement(tag) + for (const [key, value] of Object.entries(props)) { + if (key === 'onclick' || key === 'onsubmit') node[key] = value + else if (key === 'class') node.className = value + else node.setAttribute(key, value) + } + for (const child of children) node.append(child instanceof Node ? child : document.createTextNode(String(child))) + return node +} +const fromHash = new URLSearchParams(location.hash.slice(1)).get('token') +if (fromHash) { localStorage.setItem('harness-token', fromHash); history.replaceState(null, '', '/') } +let token = localStorage.getItem('harness-token') +let current = null +let source = null + +async function api(path, options = {}) { + const response = await fetch(path, { + ...options, + headers: { 'Content-Type': 'application/json', Authorization: 'Bearer ' + token, ...(options.headers || {}) }, + }) + if (response.status === 401) { localStorage.removeItem('harness-token'); token = null; render(); throw new Error('unauthorized') } + return response.json() +} + +function login() { + const input = $('input', { type: 'password', placeholder: 'Owner token', autocomplete: 'off' }) + return $('form', { class: 'card', onsubmit: (event) => { + event.preventDefault(); token = input.value.trim(); localStorage.setItem('harness-token', token); render() + } }, $('h2', {}, 'Sign in'), input, $('button', {}, 'Sign in')) +} + +async function refresh() { + if (!token) return + const [pairings, hosts, sessions] = await Promise.all([api('/api/pairings'), api('/api/hosts'), api('/api/sessions')]) + const side = document.getElementById('side') + side.replaceChildren( + $('h3', {}, 'Pairing requests'), + ...(pairings.length ? pairings.map((pairing) => $('div', { class: 'row' }, + $('span', {}, pairing.name + ' ', $('code', {}, pairing.code)), + $('button', { onclick: async () => { await api('/api/pair/approve', { method: 'POST', body: JSON.stringify({ code: pairing.code }) }); refresh() } }, 'Approve'))) : [$('p', { class: 'muted' }, 'None')]), + $('h3', {}, 'Hosts'), + ...(hosts.length ? hosts.map((host) => $('div', { class: 'row' }, + $('span', { class: host.online ? 'dot on' : 'dot' }), $('span', {}, host.name + ' (' + host.harnesses.join(', ') + ')'))) : [$('p', { class: 'muted' }, 'No hosts yet')]), + $('h3', {}, 'Sessions'), + ...(sessions.length ? sessions.map((session) => $('button', { class: 'session ' + session.status, onclick: () => open(session) }, + session.threadId + ' | ' + (session.harness || '') + ' | ' + session.status)) : [$('p', { class: 'muted' }, 'No sessions yet')]), + ) +} + +function open(session) { + current = { ...session, messages: [], interrupts: [], questions: new Map(), running: false } + if (source) source.close() + const path = '/api/sessions/' + encodeURIComponent(session.hostId) + '/' + encodeURIComponent(session.threadId) + current.path = path + source = new EventSource(path + '/events?token=' + encodeURIComponent(token)) + source.onmessage = (message) => { apply(JSON.parse(message.data)); draw() } + draw() +} + +function apply(frame) { + if (frame.type !== 'harness.event' || !current) return + const event = frame.event + if (event.subagentRunId) return + if (event.type === 'TEXT_MESSAGE_CONTENT') { + const last = current.messages[current.messages.length - 1] + if (last && last.kind === 'assistant' && last.op === frame.operationId) last.text += event.delta + else current.messages.push({ kind: 'assistant', op: frame.operationId, text: event.delta }) + } else if (event.type === 'TOOL_CALL_START') { + current.messages.push({ kind: 'tool', text: 'tool ' + event.toolCallName }) + } else if (event.type === 'RUN_ERROR') { + current.messages.push({ kind: 'error', text: event.message }) + } else if (event.type === 'RUN_FINISHED' && event.outcome && event.outcome.type === 'interrupt') { + current.interrupts = event.outcome.interrupts + } else if (event.type === 'CUSTOM') { + if (event.name === 'harness.operation.started') current.running = true + if (event.name === 'harness.operation.finished') current.running = false + if (event.name === 'harness.question') current.questions.set(event.value.questionId, event.value) + if (event.name === 'harness.question.answered') current.questions.delete(event.value.questionId) + if (event.name === 'harness.auth_required') current.messages.push({ kind: 'notice', text: 'Sign in to ' + event.value.connector + (event.value.url ? ': ' + event.value.url : '') }) + } +} + +async function send(input) { + if (input.op === 'prompt' || input.op === 'steer') current.messages.push({ kind: 'user', text: input.message }) + const receipt = await api(current.path + '/input', { method: 'POST', body: JSON.stringify({ input }) }) + if (receipt.status === 'queued') current.messages.push({ kind: 'notice', text: 'Queued: the host is offline.' }) + draw() +} + +function draw() { + const main = document.getElementById('main') + if (!current) { main.replaceChildren($('p', { class: 'muted' }, 'Pick a session.')); return } + const log = $('div', { class: 'log' }, ...current.messages.map((message) => $('div', { class: 'msg ' + message.kind }, message.text))) + const actions = [] + if (current.interrupts.length) { + const decide = (approved) => { const resume = current.interrupts.map((interrupt) => ({ interruptId: interrupt.id, status: 'resolved', payload: approved })); current.interrupts = []; send({ op: 'resolve', resume }) } + actions.push($('div', { class: 'card' }, $('strong', {}, 'Waiting for approval: '), current.interrupts.map((interrupt) => interrupt.message || interrupt.toolCallId || interrupt.id).join(', '), + $('div', { class: 'row' }, $('button', { onclick: () => decide(true) }, 'Approve'), $('button', { class: 'ghost', onclick: () => decide(false) }, 'Reject')))) + } + for (const question of current.questions.values()) { + const answer = $('input', { placeholder: 'Your answer' }) + const isBoolean = question.schema && question.schema.type === 'boolean' + actions.push($('div', { class: 'card' }, $('strong', {}, question.message), + isBoolean + ? $('div', { class: 'row' }, $('button', { onclick: () => send({ op: 'answer', questionId: question.questionId, value: true }) }, 'Yes'), $('button', { class: 'ghost', onclick: () => send({ op: 'answer', questionId: question.questionId, value: false }) }, 'No')) + : $('form', { class: 'row', onsubmit: (event) => { event.preventDefault(); send({ op: 'answer', questionId: question.questionId, value: answer.value }) } }, answer, $('button', {}, 'Answer')))) + } + const box = $('input', { placeholder: current.running ? 'Steer the running turn' : 'Send a message' }) + const form = $('form', { class: 'row prompt', onsubmit: (event) => { + event.preventDefault(); const text = box.value.trim(); if (!text) return; box.value = '' + send({ op: current.running ? 'steer' : 'prompt', message: text }) + } }, box, $('button', {}, 'Send'), $('button', { type: 'button', class: 'ghost', onclick: () => send({ op: 'cancel' }) }, 'Stop')) + main.replaceChildren($('h2', {}, current.threadId + (current.running ? ' (working)' : '')), log, ...actions, form) + log.scrollTop = log.scrollHeight +} + +function render() { + const app = document.getElementById('app') + if (!token) { app.replaceChildren(login()); return } + app.replaceChildren($('aside', { id: 'side' }), $('main', { id: 'main' })) + draw(); refresh().catch(() => {}) +} +setInterval(() => refresh().catch(() => {}), 5000) +if ('serviceWorker' in navigator) navigator.serviceWorker.register('/sw.js').catch(() => {}) +render() +` + +export const DASHBOARD_HTML = ` + + + + + + +TanStack AI Dashboard + + + +
+ + +` diff --git a/packages/ai-dashboard/tests/relay.test.ts b/packages/ai-dashboard/tests/relay.test.ts new file mode 100644 index 0000000000..daee6fa05f --- /dev/null +++ b/packages/ai-dashboard/tests/relay.test.ts @@ -0,0 +1,219 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { EventType } from '@tanstack/ai' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { startDashboard } from '../src' +import { connectDashboard } from '../src/connect' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' + +let calls = 0 +function adapterSaying(answer: string): AnyTextAdapter { + return { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + chatStream: () => + (async function* (): AsyncGenerator { + calls += 1 + const messageId = `m-${calls}` + const now = Date.now() + yield { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: now, + } + yield { + type: EventType.TEXT_MESSAGE_START, + messageId, + role: 'assistant', + timestamp: now, + } + yield { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId, + delta: answer, + timestamp: now, + } + yield { type: EventType.TEXT_MESSAGE_END, messageId, timestamp: now } + yield { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now, + metadata: { tanstack: { finishReason: 'stop' } }, + } + })(), + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } +} + +const cleanups: Array<() => unknown> = [] +afterEach(async () => { + for (const cleanup of cleanups.splice(0).reverse()) await cleanup() +}) + +describe('dashboard relay', () => { + it( + 'pairs a host, shows its session, and relays a prompt and its answer', + { timeout: 20_000 }, + async () => { + const dashboard = await startDashboard({ port: 0 }) + cleanups.push(() => dashboard.close()) + const owner = { + Authorization: `Bearer ${dashboard.ownerToken}`, + 'Content-Type': 'application/json', + } + + // The app shell loads without a token; the API does not. + expect((await fetch(`${dashboard.url}/`)).status).toBe(200) + expect((await fetch(`${dashboard.url}/api/hosts`)).status).toBe(401) + + const host = createHarnessHost({ persistence: memoryPersistence() }) + cleanups.push(() => host.close()) + const harness = defineHarness({ + name: 'acme/remote', + adapter: adapterSaying('Hello from the host.'), + }) + + let savedToken = '' + const connecting = connectDashboard({ + host, + harness, + url: dashboard.url, + threads: ['main'], + onPairingCode: (code) => { + // The owner approves the code in the dashboard. + void fetch(`${dashboard.url}/api/pair/approve`, { + method: 'POST', + headers: owner, + body: JSON.stringify({ code }), + }) + }, + onToken: (token) => { + savedToken = token + }, + }) + const connection = await connecting + cleanups.push(() => connection.close()) + expect(savedToken).toBe(connection.token) + + await vi.waitFor(async () => { + const hosts = await ( + await fetch(`${dashboard.url}/api/hosts`, { headers: owner }) + ).json() + expect(hosts).toMatchObject([ + { name: 'acme/remote', online: true, harnesses: ['acme/remote'] }, + ]) + }) + const [{ hostId }] = await ( + await fetch(`${dashboard.url}/api/hosts`, { headers: owner }) + ).json() + const receipt = await ( + await fetch(`${dashboard.url}/api/sessions/${hostId}/main/input`, { + method: 'POST', + headers: owner, + body: JSON.stringify({ + input: { op: 'prompt', message: 'hi from my phone' }, + }), + }) + ).json() + expect(receipt.status).toBe('sent') + + // The answer comes back through the relay cache. + await vi.waitFor(async () => { + const sessions = await ( + await fetch(`${dashboard.url}/api/sessions`, { headers: owner }) + ).json() + expect(sessions).toMatchObject([ + { hostId, threadId: 'main', harness: 'acme/remote', status: 'idle' }, + ]) + }) + const controller = new AbortController() + const events = await fetch( + `${dashboard.url}/api/sessions/${hostId}/main/events?token=${dashboard.ownerToken}`, + { + signal: controller.signal, + }, + ) + const reader = events.body!.getReader() + let text = '' + while (!text.includes('harness.operation.finished')) { + const { value, done } = await reader.read() + if (done) break + text += new TextDecoder().decode(value) + } + controller.abort() + expect(text).toContain('Hello from the host.') + }, + ) + + it( + 'queues inputs for an offline host and refuses revoked hosts', + { timeout: 20_000 }, + async () => { + const dashboard = await startDashboard({ port: 0 }) + cleanups.push(() => dashboard.close()) + const owner = { + Authorization: `Bearer ${dashboard.ownerToken}`, + 'Content-Type': 'application/json', + } + + const started = await ( + await fetch(`${dashboard.url}/api/pair/start`, { + method: 'POST', + body: JSON.stringify({ name: 'laptop' }), + }) + ).json() + await fetch(`${dashboard.url}/api/pair/approve`, { + method: 'POST', + headers: owner, + body: JSON.stringify({ code: started.code }), + }) + const status = await ( + await fetch( + `${dashboard.url}/api/pair/status?pairingId=${started.pairingId}`, + ) + ).json() + expect(status.status).toBe('approved') + + const queued = await ( + await fetch(`${dashboard.url}/api/sessions/${status.hostId}/t/input`, { + method: 'POST', + headers: owner, + body: JSON.stringify({ input: { op: 'prompt', message: 'later' } }), + }) + ).json() + expect(queued.status).toBe('queued') + + await fetch(`${dashboard.url}/api/hosts/revoke`, { + method: 'POST', + headers: owner, + body: JSON.stringify({ hostId: status.hostId }), + }) + const refused = await fetch(`${dashboard.url}/api/host/hello`, { + method: 'POST', + headers: { + Authorization: `Bearer ${status.token}`, + 'Content-Type': 'application/json', + }, + body: '{}', + }) + expect(refused.status).toBe(401) + }, + ) +}) diff --git a/packages/ai-dashboard/tsconfig.json b/packages/ai-dashboard/tsconfig.json new file mode 100644 index 0000000000..29112eff9f --- /dev/null +++ b/packages/ai-dashboard/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "dist" + }, + "include": ["vite.config.ts", "./src", "./tests"], + "exclude": ["node_modules", "dist"] +} diff --git a/packages/ai-dashboard/vite.config.ts b/packages/ai-dashboard/vite.config.ts new file mode 100644 index 0000000000..dadc8e1df5 --- /dev/null +++ b/packages/ai-dashboard/vite.config.ts @@ -0,0 +1,35 @@ +import { defineConfig, mergeConfig } from 'vitest/config' +import { tanstackViteConfig } from '@tanstack/vite-config' +import packageJson from './package.json' + +const config = defineConfig({ + test: { + name: packageJson.name, + dir: './', + watch: false, + globals: true, + environment: 'node', + include: ['tests/**/*.test.ts'], + coverage: { + provider: 'v8', + reporter: ['text', 'json', 'html', 'lcov'], + exclude: [ + 'node_modules/', + 'dist/', + 'tests/', + '**/*.test.ts', + '**/*.config.ts', + ], + include: ['src/**/*.ts'], + }, + }, +}) + +export default mergeConfig( + config, + tanstackViteConfig({ + entry: ['./src/index.ts', './src/connect.ts'], + srcDir: './src', + cjs: false, + }), +) diff --git a/packages/ai-harness-cli/package.json b/packages/ai-harness-cli/package.json index e9224933f2..4f91fbfe82 100644 --- a/packages/ai-harness-cli/package.json +++ b/packages/ai-harness-cli/package.json @@ -72,17 +72,22 @@ "peerDependencies": { "@tanstack/ai": "workspace:^", "@tanstack/ai-acp": "workspace:^", + "@tanstack/ai-dashboard": "workspace:^", "@tanstack/ai-harness": "workspace:^", "@tanstack/ai-persistence": "workspace:^" }, "peerDependenciesMeta": { "@tanstack/ai-acp": { "optional": true + }, + "@tanstack/ai-dashboard": { + "optional": true } }, "devDependencies": { "@tanstack/ai": "workspace:*", "@tanstack/ai-acp": "workspace:*", + "@tanstack/ai-dashboard": "workspace:*", "@tanstack/ai-harness": "workspace:*", "@tanstack/ai-persistence": "workspace:*", "@tanstack/ai-sandbox": "workspace:*", diff --git a/packages/ai-harness-cli/src/args.ts b/packages/ai-harness-cli/src/args.ts index c83474e715..9d49cf37c7 100644 --- a/packages/ai-harness-cli/src/args.ts +++ b/packages/ai-harness-cli/src/args.ts @@ -10,6 +10,7 @@ export interface CliArgs { hostname: string token?: string thread: string + dashboard?: string } export const USAGE = `Usage: [options] @@ -25,6 +26,7 @@ Options: --host Host name for --serve (default 127.0.0.1) --token Bearer token for --serve (default: HARNESS_TOKEN, or a new random token) --thread Conversation id (default: main) + --dashboard Connect to a dashboard (pairs on first use; set HARNESS_DASHBOARD_TOKEN to skip) -h, --help Show this help Exit codes for --print: 0 done, 1 failed, 2 waiting for approval, 130 cancelled.` @@ -43,6 +45,7 @@ export function parseCliArgs(argv: ReadonlyArray): CliArgs { host: { type: 'string' }, token: { type: 'string' }, thread: { type: 'string' }, + dashboard: { type: 'string' }, }, strict: true, }) @@ -64,5 +67,6 @@ export function parseCliArgs(argv: ReadonlyArray): CliArgs { hostname: values.host ?? '127.0.0.1', ...(values.token !== undefined ? { token: values.token } : {}), thread: values.thread ?? 'main', + ...(values.dashboard !== undefined ? { dashboard: values.dashboard } : {}), } } diff --git a/packages/ai-harness-cli/src/index.ts b/packages/ai-harness-cli/src/index.ts index d91dd0d9d7..18dbac479d 100644 --- a/packages/ai-harness-cli/src/index.ts +++ b/packages/ai-harness-cli/src/index.ts @@ -63,6 +63,38 @@ export async function runCli( return EXIT.ok } + if (args.dashboard) { + const { connectDashboard } = + await import('@tanstack/ai-dashboard/connect').catch(() => { + throw new Error( + '--dashboard needs @tanstack/ai-dashboard. Install it next to @tanstack/ai-harness-cli.', + ) + }) + const savedToken = env.HARNESS_DASHBOARD_TOKEN + const connection = await connectDashboard({ + host, + harness, + url: args.dashboard, + threads: [args.thread], + ...(savedToken ? { token: savedToken } : {}), + onPairingCode: (code) => + stderr.write( + `Pair this host in the dashboard with the code ${code}. Waiting for approval...\n`, + ), + onToken: (token) => + stderr.write( + `Paired. Set HARNESS_DASHBOARD_TOKEN=${token} to skip pairing next time.\n`, + ), + }) + stderr.write(`Connected to ${args.dashboard}. Press Ctrl+C to stop.\n`) + await new Promise((resolve) => { + process.once('SIGINT', resolve) + process.once('SIGTERM', resolve) + }) + connection.close() + return EXIT.ok + } + if (args.serve) { const token = args.token ?? env.HARNESS_TOKEN ?? createToken() const server = await serve({ diff --git a/packages/ai-harness-cli/src/lines.ts b/packages/ai-harness-cli/src/lines.ts index 26cc914f8e..741005b97d 100644 --- a/packages/ai-harness-cli/src/lines.ts +++ b/packages/ai-harness-cli/src/lines.ts @@ -25,7 +25,13 @@ export async function runLines( ): Promise { const reader = new AbortController() let entries: Array = [] - let printed = 0 + // True while streamed text has no line break at its end yet. + let midLine = false + const line = (text: string) => { + if (midLine) stdout.write('\n') + midLine = false + stdout.write(`${text}\n`) + } const printing = (async () => { for await (const entry of session.events({ from: session.snapshot().cursor, @@ -42,40 +48,40 @@ export async function runLines( previous.operationId === last.operationId ) { stdout.write(last.text.slice(previous.text.length)) + midLine = true + } else if (last.kind === 'assistant') { + if (midLine) stdout.write('\n') + stdout.write(last.text) + midLine = true } else { - if (printed > 0) stdout.write('\n') - stdout.write(last.kind === 'assistant' ? last.text : `[${last.text}]`) + line(`[${last.text}]`) } - printed += 1 entries = next } })() const lines = createInterface({ input, crlfDelay: Infinity }) const pendingLater: Array> = [] - for await (const line of lines) { + for await (const typed of lines) { const snapshot = session.snapshot() if (snapshot.status === 'requires_action') { await resolveAll( session, snapshot.pendingInterrupts, - /^y(es)?$/i.test(line.trim()), + /^y(es)?$/i.test(typed.trim()), ) } else { - const result = await handleLine(session, line) + const result = await handleLine(session, typed) if (result.type === 'exit') break - if (result.type === 'notice' && result.text) - stdout.write(`${result.text}\n`) + if (result.type === 'notice' && result.text) line(result.text) if (result.type === 'notice' && result.later) { - pendingLater.push( - result.later.then((text) => void stdout.write(`\n${text}\n`)), - ) + pendingLater.push(result.later.then((text) => line(text))) } } await waitIdle(session) const after = session.snapshot() if (after.status === 'requires_action') { - stdout.write(`\n${approvalQuestion(after.pendingInterrupts)}\n`) + line(approvalQuestion(after.pendingInterrupts)) } } lines.close() @@ -83,5 +89,5 @@ export async function runLines( await waitIdle(session) reader.abort() await printing - if (printed > 0) stdout.write('\n') + if (midLine) stdout.write('\n') } diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 4a914a15c6..d8f2711c60 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -174,6 +174,43 @@ importers: specifier: ^8.2.1 version: 8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0) + examples/harness-cli: + dependencies: + '@tanstack/ai': + specifier: workspace:* + version: link:../../packages/ai + '@tanstack/ai-acp': + specifier: workspace:* + version: link:../../packages/ai-acp + '@tanstack/ai-anthropic': + specifier: workspace:* + version: link:../../packages/ai-anthropic + '@tanstack/ai-dashboard': + specifier: workspace:* + version: link:../../packages/ai-dashboard + '@tanstack/ai-harness': + specifier: workspace:* + version: link:../../packages/ai-harness + '@tanstack/ai-harness-cli': + specifier: workspace:* + version: link:../../packages/ai-harness-cli + '@tanstack/ai-openai': + specifier: workspace:* + version: link:../../packages/ai-openai + '@tanstack/ai-persistence': + specifier: workspace:* + version: link:../../packages/ai-persistence + zod: + specifier: ^4.2.0 + version: 4.3.6 + devDependencies: + '@types/node': + specifier: ^24.10.1 + version: 24.13.6 + tsx: + specifier: ^4.23.12 + version: 4.23.12 + examples/react/basic-chat: dependencies: '@tailwindcss/vite': @@ -2335,6 +2372,21 @@ importers: specifier: 4.1.10 version: 4.1.10(vitest@4.1.11) + packages/ai-dashboard: + devDependencies: + '@tanstack/ai': + specifier: workspace:* + version: link:../ai + '@tanstack/ai-harness': + specifier: workspace:* + version: link:../ai-harness + '@tanstack/ai-persistence': + specifier: workspace:* + version: link:../ai-persistence + '@vitest/coverage-v8': + specifier: 4.1.10 + version: 4.1.10(vitest@4.1.11) + packages/ai-devtools: dependencies: '@tanstack/ai': @@ -2556,6 +2608,9 @@ importers: '@tanstack/ai-acp': specifier: workspace:* version: link:../ai-acp + '@tanstack/ai-dashboard': + specifier: workspace:* + version: link:../ai-dashboard '@tanstack/ai-harness': specifier: workspace:* version: link:../ai-harness @@ -3611,6 +3666,9 @@ importers: '@tanstack/ai-compaction': specifier: workspace:* version: link:../../packages/ai-compaction + '@tanstack/ai-dashboard': + specifier: workspace:* + version: link:../../packages/ai-dashboard '@tanstack/ai-elevenlabs': specifier: workspace:* version: link:../../packages/ai-elevenlabs diff --git a/testing/e2e/package.json b/testing/e2e/package.json index 41b655bed1..03584db0fa 100644 --- a/testing/e2e/package.json +++ b/testing/e2e/package.json @@ -26,6 +26,7 @@ "@tanstack/ai-client": "workspace:*", "@tanstack/ai-cloudflare": "workspace:*", "@tanstack/ai-compaction": "workspace:*", + "@tanstack/ai-dashboard": "workspace:*", "@tanstack/ai-elevenlabs": "workspace:*", "@tanstack/ai-fal": "workspace:*", "@tanstack/ai-gemini": "workspace:*", diff --git a/testing/e2e/tests/dashboard.spec.ts b/testing/e2e/tests/dashboard.spec.ts new file mode 100644 index 0000000000..786a8f0363 --- /dev/null +++ b/testing/e2e/tests/dashboard.spec.ts @@ -0,0 +1,124 @@ +import { EventType } from '@tanstack/ai' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { startDashboard } from '@tanstack/ai-dashboard' +import { connectDashboard } from '@tanstack/ai-dashboard/connect' +import { test, expect } from './fixtures' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' + +/** Answers with what the user said, so the test can follow the round trip. */ +function echoAdapter(): AnyTextAdapter { + let calls = 0 + return { + kind: 'text', + name: 'echo', + model: 'echo', + '~types': { + providerOptions: {}, + inputModalities: ['text'], + messageMetadataByModality: { + text: undefined, + image: undefined, + audio: undefined, + video: undefined, + document: undefined, + }, + toolCapabilities: [], + toolCallMetadata: undefined, + systemPromptMetadata: undefined as never, + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + chatStream: (options) => + (async function* (): AsyncGenerator { + calls += 1 + const messageId = `echo-${calls}` + const now = Date.now() + const last = options.messages.at(-1) + const said = typeof last?.content === 'string' ? last.content : '' + yield { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: now, + } + yield { + type: EventType.TEXT_MESSAGE_START, + messageId, + role: 'assistant', + timestamp: now, + } + yield { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId, + delta: `You said: ${said}`, + timestamp: now, + } + yield { type: EventType.TEXT_MESSAGE_END, messageId, timestamp: now } + yield { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now, + metadata: { tanstack: { finishReason: 'stop' } }, + } + })(), + } +} + +test.describe('dashboard', () => { + test('a paired host shows up, and a prompt from a phone reaches it', async ({ + page, + }) => { + const dashboard = await startDashboard({ port: 0 }) + const owner = { + Authorization: `Bearer ${dashboard.ownerToken}`, + 'Content-Type': 'application/json', + } + const host = createHarnessHost({ persistence: memoryPersistence() }) + const connection = await connectDashboard({ + host, + harness: defineHarness({ + name: 'acme/phone-demo', + adapter: echoAdapter(), + }), + url: dashboard.url, + threads: ['main'], + onPairingCode: (code) => + void fetch(`${dashboard.url}/api/pair/approve`, { + method: 'POST', + headers: owner, + body: JSON.stringify({ code }), + }), + }) + try { + const [{ hostId }] = await ( + await fetch(`${dashboard.url}/api/hosts`, { headers: owner }) + ).json() + await fetch(`${dashboard.url}/api/sessions/${hostId}/main/input`, { + method: 'POST', + headers: owner, + body: JSON.stringify({ input: { op: 'prompt', message: 'first' } }), + }) + + await page.setViewportSize({ width: 390, height: 844 }) + await page.goto(`${dashboard.url}/#token=${dashboard.ownerToken}`) + await expect( + page.getByText('acme/phone-demo (acme/phone-demo)'), + ).toBeVisible() + await page.getByRole('button', { name: /main/ }).click() + await expect(page.getByText('You said: first')).toBeVisible() + + await page + .getByPlaceholder(/Send a message|Steer/) + .fill('hello from the phone') + await page.getByRole('button', { name: 'Send' }).click() + await expect( + page.getByText('You said: hello from the phone'), + ).toBeVisible() + } finally { + connection.close() + await host.close() + await dashboard.close() + } + }) +}) From 551c8670ef921be19be1b2fa7aeae8ae244f0497 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 14:23:35 +0200 Subject: [PATCH 014/286] test(ai): cover every bound activity wrapper --- packages/ai/tests/bound-activities.test.ts | 148 +++++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 packages/ai/tests/bound-activities.test.ts diff --git a/packages/ai/tests/bound-activities.test.ts b/packages/ai/tests/bound-activities.test.ts new file mode 100644 index 0000000000..1eda9780c6 --- /dev/null +++ b/packages/ai/tests/bound-activities.test.ts @@ -0,0 +1,148 @@ +import { describe, expect, it, vi } from 'vitest' + +// Replace every activity with a spy that returns the options it got, so the +// test sees exactly what each bound wrapper passes on. +const record = (name: string) => + vi.fn((options: unknown) => ({ name, options })) +vi.mock('../src/activities/chat/index', async (importOriginal) => ({ + ...(await importOriginal()), + chat: record('chat'), +})) +vi.mock('../src/activities/summarize/index', () => ({ + summarize: record('summarize'), +})) +vi.mock('../src/activities/generateImage/index', () => ({ + generateImage: record('image'), +})) +vi.mock('../src/activities/generateVideo/index', () => ({ + generateVideo: record('video'), +})) +vi.mock('../src/activities/generateAudio/index', () => ({ + generateAudio: record('audio'), +})) +vi.mock('../src/activities/generateSpeech/index', () => ({ + generateSpeech: record('speech'), +})) +vi.mock('../src/activities/generateVoice/index', () => ({ + generateVoice: record('voice'), +})) +vi.mock('../src/activities/generateTranscription/index', () => ({ + generateTranscription: record('transcription'), +})) +vi.mock('../src/activities/generateWorld/index', () => ({ + generateWorld: record('world'), +})) +vi.mock('../src/activities/generateLiveVideo/index', () => ({ + generateLiveVideo: record('liveVideo'), +})) +vi.mock('../src/activities/embed/index', () => ({ embed: record('embed') })) +vi.mock('../src/activities/rerank/index', () => ({ rerank: record('rerank') })) +vi.mock('../src/activities/evaluate/index', () => ({ + decide: record('decide'), +})) + +const { createBoundActivities } = + await import('../src/activities/chat/agents/bound') + +const input = { + input: undefined, + messages: [{ role: 'user' as const, content: 'hi' }], + threadId: 'thread-1', + runId: 'run-1', + parentRunId: 'parent-1', + subagentRunId: 'child-1', +} + +type Recorded = { name: string; options: Record } + +describe('createBoundActivities', () => { + const generationMiddleware = [{ name: 'host-generation' }] + const chatMiddleware = [{ name: 'host-chat' }] + + function bound() { + const controller = new AbortController() + const activities = createBoundActivities(input, controller, { + chatMiddleware: chatMiddleware as never, + generationMiddleware: generationMiddleware as never, + }) + return { activities, controller } + } + + it('gives every id-taking activity its own run id, the thread, the signal, and host middleware first', () => { + const { activities, controller } = bound() + const own = { name: 'own' } + const calls: Array<[keyof typeof activities, string]> = [ + ['summarize', 'summarize'], + ['generateImage', 'image'], + ['generateVideo', 'video'], + ['generateAudio', 'audio'], + ['generateSpeech', 'speech'], + ['generateVoice', 'voice'], + ['generateTranscription', 'transcription'], + ['generateWorld', 'world'], + ['generateLiveVideo', 'liveVideo'], + ] + calls.forEach(([method, activity], index) => { + const call = activities[method] as unknown as ( + options: object, + ) => Recorded + const result = call({ prompt: 'x', middleware: [own] }) + expect(result.name).toBe(activity) + expect(result.options).toMatchObject({ + threadId: 'thread-1', + runId: `run-1:${activity}-${index + 1}`, + abortSignal: controller.signal, + prompt: 'x', + }) + expect(result.options.middleware).toEqual([...generationMiddleware, own]) + }) + }) + + it('passes only middleware to embed, and the signal plus middleware to rerank and decide', () => { + const { activities, controller } = bound() + const embed = ( + activities.embed as unknown as (options: object) => Recorded + )({ input: 'x' }) + expect(embed.options).toEqual({ + input: 'x', + middleware: generationMiddleware, + }) + + for (const method of ['rerank', 'decide'] as const) { + const call = activities[method] as unknown as ( + options: object, + ) => Recorded + const result = call({ query: 'q' }) + expect(result.options).toEqual({ + query: 'q', + abortSignal: controller.signal, + middleware: generationMiddleware, + }) + } + }) + + it('fills chat with the child ids and messages, and lets options win', () => { + const { activities, controller } = bound() + const chat = activities.chat as unknown as (options: object) => Recorded + const defaults = chat({ adapter: 'a' }) + expect(defaults.options).toMatchObject({ + messages: input.messages, + threadId: 'thread-1', + runId: 'run-1', + parentRunId: 'parent-1', + subagentRunId: 'child-1', + abortController: controller, + middleware: chatMiddleware, + }) + const custom = chat({ adapter: 'a', messages: [], threadId: 'mine' }) + expect(custom.options).toMatchObject({ messages: [], threadId: 'mine' }) + }) + + it('works without a binding', () => { + const activities = createBoundActivities(input, new AbortController()) + const result = ( + activities.generateImage as unknown as (options: object) => Recorded + )({ prompt: 'x' }) + expect(result.options.middleware).toEqual([]) + }) +}) From 24c16d5a870535de80f166a512a27b1e479162da Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 14:27:18 +0200 Subject: [PATCH 015/286] test(ai-sandbox): type the snapshot fixture by the snapshot persistence shape --- packages/ai-sandbox/tests/snapshot-lifecycle.test.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/ai-sandbox/tests/snapshot-lifecycle.test.ts b/packages/ai-sandbox/tests/snapshot-lifecycle.test.ts index f46eb5a61b..23371010c8 100644 --- a/packages/ai-sandbox/tests/snapshot-lifecycle.test.ts +++ b/packages/ai-sandbox/tests/snapshot-lifecycle.test.ts @@ -10,6 +10,7 @@ import { import { defineSandbox } from '../src/sandbox' import { withSandbox } from '../src/middleware' import { memorySandboxSnapshots } from '../src/memory-snapshots' +import type { MemorySnapshotPersistence } from '../src/memory-snapshot-types' import { SandboxCapability } from '../src/capabilities' import { InMemorySandboxInstanceStore } from '../src/instance-store' import { InMemorySandboxCheckpointStore } from '../src/checkpoint-store' @@ -72,7 +73,7 @@ type WorkspaceSeed = | { path: string; type: 'dir' } | { path: string; type: 'file'; data: Uint8Array } type FixtureOptions = { - persistence?: ReturnType + persistence?: MemorySnapshotPersistence checkpoints?: SandboxCheckpointStore workspace?: Array onWorkspaceList?: (path: string) => void | Promise @@ -88,7 +89,7 @@ type Fixture = { events: Array instances: InMemorySandboxInstanceStore checkpoints: SandboxCheckpointStore - persistence: ReturnType + persistence: MemorySnapshotPersistence definition: ReturnType resumed?: SandboxHandle } From 924aa4f7008ea0f2f6f155552323983796882512 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 14:52:51 +0200 Subject: [PATCH 016/286] test(ai-harness): cover buildHarness without Bun, run the real build only where Bun exists --- packages/ai-harness/tests/build-stub.test.ts | 300 +++++++++++++++++++ packages/ai-harness/tests/build.test.ts | 7 +- 2 files changed, 306 insertions(+), 1 deletion(-) create mode 100644 packages/ai-harness/tests/build-stub.test.ts diff --git a/packages/ai-harness/tests/build-stub.test.ts b/packages/ai-harness/tests/build-stub.test.ts new file mode 100644 index 0000000000..838085ef57 --- /dev/null +++ b/packages/ai-harness/tests/build-stub.test.ts @@ -0,0 +1,300 @@ +import { createHash } from 'node:crypto' +import { + chmod, + mkdir, + readFile, + readdir, + rm, + writeFile, +} from 'node:fs/promises' +import { join } from 'node:path' +import { fileURLToPath } from 'node:url' +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { EventType, chat } from '@tanstack/ai' +import { artifactText, buildHarness, readManifest } from '../src/build' +import type { StreamChunk } from '@tanstack/ai' + +// These tests stand in a small script for Bun, so they run where Bun is not +// installed. The script "bundles" by copying a hand-written worker that +// speaks the harness frames. build.test.ts covers a real Bun build. +const tmp = fileURLToPath(new URL('./.tmp-build-stub/', import.meta.url)) +const app = join(tmp, 'app') +const broken = join(tmp, 'broken') +const entry = join(app, 'studio.ts') + +const stubWorker = ` +import { createInterface } from 'node:readline' +const send = (frame) => process.stdout.write(JSON.stringify(frame) + '\\n') +let turns = 0 +createInterface({ input: process.stdin }).on('line', (line) => { + const frame = JSON.parse(line) + if (frame.type !== 'harness.input') return + const message = frame.input.message + if (message === 'reject') { + send({ type: 'harness.receipt', requestId: frame.requestId, status: 'rejected', reason: 'busy' }) + return + } + if (message === 'reject quietly') { + send({ type: 'harness.receipt', requestId: frame.requestId, status: 'rejected' }) + return + } + if (message === 'crash') { + send({ type: 'harness.error', message: 'worker crashed' }) + return + } + turns += 1 + const operationId = 'op-' + turns + const event = (value) => + send({ type: 'harness.event', cursor: String(turns), operationId, event: { timestamp: 1, ...value } }) + // One event arrives before the receipt, one belongs to another operation. + event({ type: 'TEXT_MESSAGE_START', messageId: 'm' + turns, role: 'assistant' }) + send({ type: 'harness.event', cursor: 'x', operationId: 'other', event: { type: 'TEXT_MESSAGE_CONTENT', messageId: 'x', delta: 'ignored', timestamp: 1 } }) + send({ type: 'harness.receipt', requestId: frame.requestId, status: 'accepted', operationId }) + process.stdout.write('a stray log line\\n') + if (message === 'fail') { + event({ type: 'RUN_ERROR', message: 'turn failed' }) + return + } + send({ type: 'harness.event', cursor: 'y', operationId: 'other', event: { type: 'TEXT_MESSAGE_CONTENT', messageId: 'x', delta: 'ignored', timestamp: 1 } }) + event({ type: 'TOOL_CALL_START', toolCallId: 't', toolCallName: 'hidden' }) + event({ type: 'TEXT_MESSAGE_CONTENT', messageId: 'm' + turns, delta: 'turn ' + turns + ': ' + message }) + event({ type: 'TEXT_MESSAGE_CONTENT', messageId: 'c', delta: ' (child)', subagentRunId: 'child-1' }) + event({ type: 'TEXT_MESSAGE_END', messageId: 'm' + turns }) + event({ type: 'CUSTOM', name: 'harness.operation.finished', value: {} }) +}) +` + +const fakeBun = ` +import { copyFileSync, writeFileSync } from 'node:fs' +const args = process.argv.slice(2) +const outfile = args.find((arg) => arg.startsWith('--outfile='))?.slice('--outfile='.length) +if (process.cwd().endsWith('broken')) { + process.stderr.write('bundle exploded') + process.exit(1) +} +if (args[0] === 'build' && args.includes('--compile')) { + writeFileSync(outfile, 'executable for ' + (args.find((arg) => arg.startsWith('--target=')) ?? 'this machine')) +} else if (args[0] === 'build') { + copyFileSync(new URL('./stub-worker.mjs', import.meta.url), outfile) +} else { + console.log('noise from the entry module') + console.log(JSON.stringify({ + name: 'acme/stub', + agents: [{ name: 'painter', produces: 'image' }], + plugins: [{ name: 'tanstack/workspace-tools' }, { name: 'connector/notion' }], + })) +} +` + +let bun: string + +beforeAll(async () => { + await mkdir(app, { recursive: true }) + await mkdir(broken, { recursive: true }) + await writeFile(entry, 'export default {}\n') + await writeFile(join(broken, 'studio.ts'), 'export default {}\n') + await writeFile(join(tmp, 'stub-worker.mjs'), stubWorker) + await writeFile(join(tmp, 'fake-bun.mjs'), fakeBun) + if (process.platform === 'win32') { + bun = join(tmp, 'fake-bun.cmd') + await writeFile(bun, `@"${process.execPath}" "%~dp0fake-bun.mjs" %*\r\n`) + } else { + bun = join(tmp, 'fake-bun') + await writeFile( + bun, + `#!/bin/sh\nexec "${process.execPath}" "$(dirname "$0")/fake-bun.mjs" "$@"\n`, + ) + await chmod(bun, 0o755) + } +}) + +afterAll(async () => { + await rm(tmp, { recursive: true, force: true }) +}) + +async function textOf(stream: AsyncIterable): Promise { + let out = '' + for await (const chunk of stream) { + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) out += chunk.delta + } + return out +} + +describe('buildHarness with a stand-in bundler', () => { + it('writes a manifest from what the entry describes, then cleans up', async () => { + const outDir = join(tmp, 'out') + const { manifest, bundle, executable } = await buildHarness({ + entry, + outDir, + bun, + }) + + const bytes = await readFile(bundle) + expect(manifest).toMatchObject({ + format: 'tanstack-ai-harness', + version: 1, + name: 'acme/stub', + entry: 'harness.js', + digest: createHash('sha256').update(bytes).digest('hex'), + runtime: { kind: 'node' }, + agents: [{ name: 'painter', produces: 'image' }], + plugins: [ + { name: 'tanstack/workspace-tools' }, + { name: 'connector/notion' }, + ], + requires: { + filesystem: true, + processExecution: true, + network: 'declared', + }, + }) + expect(executable).toBeUndefined() + expect(await readManifest(outDir)).toEqual(manifest) + // The temporary bootstrap and describe modules are gone. + expect( + (await readdir(app)).filter((name) => name.startsWith('.harness-')), + ).toEqual([]) + }) + + it('also compiles an executable when asked', async () => { + const outfile = join(tmp, 'studio-bin') + const { executable } = await buildHarness({ + entry, + export: 'studio', + outDir: join(tmp, 'out-compiled'), + bun, + compile: { entry, outfile, target: 'bun-linux-x64' }, + }) + expect(executable).toBe(outfile) + expect(await readFile(outfile, 'utf8')).toBe( + 'executable for --target=bun-linux-x64', + ) + + await buildHarness({ + entry, + outDir: join(tmp, 'out-compiled-default'), + bun, + compile: { entry, outfile }, + }) + expect(await readFile(outfile, 'utf8')).toBe('executable for this machine') + }) + + it('reports a failed bundle with its output and still cleans up', async () => { + await expect( + buildHarness({ + entry: join(broken, 'studio.ts'), + outDir: join(tmp, 'out-broken'), + bun, + }), + ).rejects.toThrow('bundle exploded') + expect( + (await readdir(broken)).filter((name) => name.startsWith('.harness-')), + ).toEqual([]) + }) + + it('says to install Bun when the command cannot start', async () => { + // On Windows the command runs through a shell, which reports the missing + // command as a failed run instead of a spawn error. + await expect( + buildHarness({ + entry, + outDir: join(tmp, 'out-missing'), + bun: join(tmp, 'no-such-bun.exe'), + }), + ).rejects.toThrow(/Could not run|failed/) + }) +}) + +describe('readManifest', () => { + it('refuses a directory whose manifest is not a harness manifest', async () => { + const dir = join(tmp, 'not-a-harness') + await mkdir(dir, { recursive: true }) + for (const manifest of [ + null, + 'text', + { format: 'other' }, + { format: 'tanstack-ai-harness', version: 2 }, + ]) { + await writeFile( + join(dir, 'harness.manifest.json'), + JSON.stringify(manifest), + ) + await expect(readManifest(dir)).rejects.toThrow( + 'has no tanstack-ai-harness v1 manifest', + ) + } + }) +}) + +describe('artifactText with a built stub', () => { + it('streams the worker text for each turn and keeps one worker per thread', async () => { + const outDir = join(tmp, 'out-chat') + await buildHarness({ entry, outDir, bun }) + const model = await artifactText(outDir) + try { + expect(model.model).toBe('acme/stub') + const ask = (content: string, threadId: string) => + textOf( + chat({ + adapter: model, + messages: [ + { role: 'assistant', content: 'earlier' }, + { role: 'user', content }, + ], + threadId, + }) as AsyncIterable, + ) + expect(await ask('hello', 'a')).toBe('turn 1: hello') + expect(await ask('again', 'a')).toBe('turn 2: again') + // Another thread starts its own worker, so it counts from one. + expect(await ask('fresh', 'b')).toBe('turn 1: fresh') + await expect(model.structuredOutput({} as never)).rejects.toThrow( + 'does not support structured output', + ) + } finally { + model.dispose() + } + }) + + it('turns worker refusals and errors into errors', async () => { + const outDir = join(tmp, 'out-errors') + await buildHarness({ entry, outDir, bun }) + const model = await artifactText(outDir, { node: process.execPath }) + const stream = (content: unknown) => + model.chatStream({ + model: 'acme/stub', + messages: [{ role: 'user', content }], + threadId: 'errors', + } as never) + const drain = async (content: unknown) => { + for await (const _chunk of stream(content)) { + // Only the error matters. + } + } + try { + await expect(drain('reject')).rejects.toThrow('busy') + await expect(drain('reject quietly')).rejects.toThrow( + 'The worker refused the prompt.', + ) + await expect(drain('crash')).rejects.toThrow('worker crashed') + await expect(drain('fail')).rejects.toThrow('turn failed') + // A message without text content is sent as an empty prompt. + let text = '' + for await (const chunk of stream([{ type: 'text', content: 'x' }])) { + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) text += chunk.delta + } + expect(text).toBe('turn 2: ') + } finally { + model.dispose() + } + }) + + it('refuses a bundle that does not match its manifest', async () => { + const outDir = join(tmp, 'out-tampered') + await buildHarness({ entry, outDir, bun }) + await writeFile(join(outDir, 'harness.js'), '// changed\n') + await expect(artifactText(outDir)).rejects.toThrow( + 'does not match its manifest digest', + ) + }) +}) diff --git a/packages/ai-harness/tests/build.test.ts b/packages/ai-harness/tests/build.test.ts index 90cd18fbd8..748e14b619 100644 --- a/packages/ai-harness/tests/build.test.ts +++ b/packages/ai-harness/tests/build.test.ts @@ -1,3 +1,4 @@ +import { spawnSync } from 'node:child_process' import { mkdir, readFile, rm, writeFile } from 'node:fs/promises' import { PassThrough } from 'node:stream' import { fileURLToPath } from 'node:url' @@ -94,8 +95,12 @@ describe('harnessText({ url })', () => { }) const tmp = fileURLToPath(new URL('./.tmp-build/', import.meta.url)) +// A real build needs Bun. build-stub.test.ts covers the rest without it. +const hasBun = + spawnSync('bun', ['--version'], { shell: process.platform === 'win32' }) + .status === 0 -describe('buildHarness and artifactText', () => { +describe.skipIf(!hasBun)('buildHarness and artifactText', () => { afterAll(async () => { await rm(tmp, { recursive: true, force: true }) }) From 0a8be1078978461d2d1facca8de4344f3a6c4f03 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 15:03:46 +0200 Subject: [PATCH 017/286] fix(ai-harness): read AG-UI text parts in /run prompts; cover protocol, client, resume, and ACP edges --- packages/ai-acp/tests/agent-handlers.test.ts | 392 ++++++++++++ packages/ai-harness/src/http.ts | 7 +- .../ai-harness/tests/resume-edges.test.ts | 279 ++++++++ .../ai-harness/tests/transport-edges.test.ts | 596 ++++++++++++++++++ 4 files changed, 1271 insertions(+), 3 deletions(-) create mode 100644 packages/ai-acp/tests/agent-handlers.test.ts create mode 100644 packages/ai-harness/tests/resume-edges.test.ts create mode 100644 packages/ai-harness/tests/transport-edges.test.ts diff --git a/packages/ai-acp/tests/agent-handlers.test.ts b/packages/ai-acp/tests/agent-handlers.test.ts new file mode 100644 index 0000000000..817c5bc7b1 --- /dev/null +++ b/packages/ai-acp/tests/agent-handlers.test.ts @@ -0,0 +1,392 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { EventType, toolDefinition } from '@tanstack/ai' +import { client, ndJsonStream } from '@agentclientprotocol/sdk/experimental/v2' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { createAcpAgent, serveAcp, toSessionUpdate } from '../src/agent' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import type { + ClientApp, + SessionUpdate, +} from '@agentclientprotocol/sdk/experimental/v2' + +const now = () => Date.now() + +/** + * A model that answers from the last user message: `wait` waits until the + * turn is cancelled, `remove` calls the remove tool, anything else echoes. + */ +function model(): AnyTextAdapter { + return { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + chatStream: (options: any) => + (async function* (): AsyncGenerator { + const last = options.messages.at(-1) + const said = typeof last?.content === 'string' ? last.content : '' + yield { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: now(), + } + if (said === 'wait') { + const signal: AbortSignal | undefined = + options.abortController?.signal ?? options.request?.signal + await new Promise((resolve) => { + if (!signal || signal.aborted) return resolve() + signal.addEventListener('abort', () => resolve(), { once: true }) + }) + return + } + if (said.startsWith('remove')) { + yield { + type: EventType.TOOL_CALL_START, + toolCallId: 'call_1', + toolCallName: 'remove', + timestamp: now(), + } + yield { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call_1', + delta: '{"path":"a.txt"}', + timestamp: now(), + } + yield { + type: EventType.TOOL_CALL_END, + toolCallId: 'call_1', + timestamp: now(), + } + yield { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'tool_calls' } }, + } + return + } + const reply = last?.role === 'tool' ? 'Kept a.txt.' : `echo: ${said}` + yield { + type: EventType.TEXT_MESSAGE_START, + messageId: 'm', + role: 'assistant', + timestamp: now(), + } + yield { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm', + delta: reply, + timestamp: now(), + } + yield { + type: EventType.TEXT_MESSAGE_END, + messageId: 'm', + timestamp: now(), + } + yield { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'stop' } }, + } + })(), + } +} + +function setup() { + const remove = vi.fn(async () => ({ removed: true })) + const harness = defineHarness({ + name: 'test/acp-handlers', + adapter: model(), + tools: [ + toolDefinition({ + name: 'remove', + description: 'Remove a file', + needsApproval: true, + inputSchema: { + type: 'object', + properties: { path: { type: 'string' } }, + required: ['path'], + }, + }).server(remove), + ], + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const updates: Array<{ sessionId: string; update: SessionUpdate }> = [] + const permissions: Array<{ title: string; toolCallId: string }> = [] + const editor: ClientApp = client() + .onNotification('session/update', ({ params }) => { + updates.push({ sessionId: params.sessionId, update: params.update }) + }) + .onRequest('session/request_permission', ({ params }) => { + permissions.push({ + title: params.title, + // The subject names the tool call: `{ type, toolCall: { toolCallId } }`. + toolCallId: JSON.stringify(params.subject ?? null).includes('call_1') + ? 'call_1' + : '', + }) + return { outcome: { outcome: 'selected', optionId: 'reject' } } + }) + const idle = (sessionId: string) => + updates.find( + (entry) => + entry.sessionId === sessionId && + entry.update.sessionUpdate === 'state_update' && + entry.update.state === 'idle', + )?.update + const textOf = (sessionId: string) => + updates + .filter((entry) => entry.sessionId === sessionId) + .map((entry) => { + if (entry.update.sessionUpdate !== 'agent_message_chunk') return '' + const content: unknown = entry.update.content + return typeof content === 'object' && + content !== null && + 'text' in content && + typeof content.text === 'string' + ? content.text + : '' + }) + .join('') + const running = (sessionId: string) => + updates.some( + (entry) => + entry.sessionId === sessionId && + entry.update.sessionUpdate === 'state_update' && + entry.update.state === 'running', + ) + return { + harness, + host, + editor, + updates, + permissions, + remove, + idle, + running, + textOf, + } +} + +describe('toSessionUpdate', () => { + it('maps reasoning and tool results, and skips child and unknown events', () => { + expect( + toSessionUpdate({ + type: EventType.REASONING_MESSAGE_CONTENT, + messageId: 'r1', + delta: 'thinking', + timestamp: 1, + }), + ).toEqual({ + sessionUpdate: 'agent_thought_chunk', + messageId: 'r1', + content: { type: 'text', text: 'thinking' }, + }) + expect( + toSessionUpdate({ + type: EventType.TOOL_CALL_RESULT, + messageId: 'm', + toolCallId: 'call_9', + content: 'done', + timestamp: 1, + }), + ).toEqual({ + sessionUpdate: 'tool_call_update', + toolCallId: 'call_9', + status: 'completed', + }) + expect( + toSessionUpdate({ + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm', + delta: 'from a child', + subagentRunId: 'child-1', + timestamp: 1, + } as StreamChunk), + ).toBeUndefined() + expect( + toSessionUpdate({ + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: 1, + }), + ).toBeUndefined() + }) +}) + +describe('ACP agent handlers', () => { + let close: (() => Promise) | undefined + afterEach(async () => { + await close?.() + close = undefined + }) + + it('keeps the tool from running when the editor rejects it', async () => { + const { harness, host, editor, permissions, remove, idle, textOf } = setup() + close = () => host.close() + await editor.connectWith(createAcpAgent({ host, harness }), async (ctx) => { + const init = await ctx.request('initialize', { + protocolVersion: 2, + info: { name: 'editor', version: '1.0.0' }, + }) + expect(init.info.version).toBe('0.0.0') + const { sessionId } = await ctx.request('session/new', { + cwd: '/tmp', + mcpServers: [], + }) + await ctx.request('session/prompt', { + sessionId, + prompt: [ + { type: 'image', data: 'AA==', mimeType: 'image/png' }, + { type: 'text', text: 'remove a.txt' }, + ], + }) + await vi.waitFor(() => expect(idle(sessionId)).toBeDefined()) + expect(idle(sessionId)).toMatchObject({ stopReason: 'end_turn' }) + expect(textOf(sessionId)).toBe('Kept a.txt.') + }) + expect(remove).not.toHaveBeenCalled() + expect(permissions).toEqual([ + { title: expect.any(String), toolCallId: 'call_1' }, + ]) + }) + + it('resumes a session by id, cancels a running turn, and closes it', async () => { + const { harness, host, editor, idle, running, textOf } = setup() + close = () => host.close() + await editor.connectWith(createAcpAgent({ host, harness }), async (ctx) => { + await ctx.request('initialize', { + protocolVersion: 2, + info: { name: 'editor', version: '1.0.0' }, + }) + await ctx.request('session/resume', { + sessionId: 'kept-session', + cwd: '/tmp', + }) + await ctx.request('session/prompt', { + sessionId: 'kept-session', + prompt: [{ type: 'text', text: 'hello' }], + }) + await vi.waitFor(() => expect(idle('kept-session')).toBeDefined()) + expect(textOf('kept-session')).toBe('echo: hello') + + await ctx.request('session/prompt', { + sessionId: 'slow-session', + prompt: [{ type: 'text', text: 'wait' }], + }) + await vi.waitFor(() => expect(running('slow-session')).toBe(true)) + await ctx.notify('session/cancel', { sessionId: 'slow-session' }) + await vi.waitFor(() => expect(idle('slow-session')).toBeDefined()) + expect(idle('slow-session')).toMatchObject({ stopReason: 'cancelled' }) + + await ctx.request('session/close', { sessionId: 'kept-session' }) + // Closing an unknown session and cancelling one are both quiet. + await ctx.request('session/close', { sessionId: 'never-opened' }) + await ctx.notify('session/cancel', { sessionId: 'never-opened' }) + }) + }) +}) + +describe('serveAcp', () => { + it('serves the agent over a stream', async () => { + const { harness, host, editor, idle, textOf } = setup() + const toClient = new TransformStream() + const toAgent = new TransformStream() + const connection = serveAcp({ + host, + harness, + version: '2.0.0', + stream: ndJsonStream(toClient.writable, toAgent.readable), + }) + await editor.connectWith( + ndJsonStream(toAgent.writable, toClient.readable), + async (ctx) => { + const init = await ctx.request('initialize', { + protocolVersion: 2, + info: { name: 'editor', version: '1.0.0' }, + }) + expect(init.info).toEqual({ + name: 'test/acp-handlers', + version: '2.0.0', + }) + const { sessionId } = await ctx.request('session/new', { + cwd: '/tmp', + mcpServers: [], + }) + await ctx.request('session/prompt', { + sessionId, + prompt: [{ type: 'text', text: 'over a stream' }], + }) + await vi.waitFor(() => expect(idle(sessionId)).toBeDefined()) + expect(textOf(sessionId)).toBe('echo: over a stream') + }, + ) + connection.close() + await host.close() + }) + + it('uses stdin and stdout when no stream is given', async () => { + const { harness, host } = setup() + const listeners = new Map void>() + const on = vi + .spyOn(process.stdin, 'on') + .mockImplementation((event: string | symbol, listener: any) => { + listeners.set(String(event), listener) + return process.stdin + }) + const written: Array = [] + const write = vi + .spyOn(process.stdout, 'write') + .mockImplementation((chunk: any, callback?: any) => { + written.push(Buffer.from(chunk).toString()) + if (typeof callback === 'function') callback() + return true + }) + try { + const connection = serveAcp({ host, harness }) + await vi.waitFor(() => expect(listeners.has('data')).toBe(true)) + listeners.get('data')?.( + new TextEncoder().encode( + `${JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'initialize', + params: { + protocolVersion: 2, + info: { name: 'editor', version: '1.0.0' }, + }, + })}\n`, + ), + ) + await vi.waitFor(() => + expect(written.join('')).toContain('test/acp-handlers'), + ) + listeners.get('end')?.() + connection.close() + } finally { + on.mockRestore() + write.mockRestore() + await host.close() + } + }) +}) diff --git a/packages/ai-harness/src/http.ts b/packages/ai-harness/src/http.ts index e9be9a5e24..3f222d5425 100644 --- a/packages/ai-harness/src/http.ts +++ b/packages/ai-harness/src/http.ts @@ -54,11 +54,12 @@ function lastUserText(messages: ReadonlyArray): string | undefined { ) if (!isRecord(message)) return undefined if (typeof message.content === 'string') return message.content - const parts = Array.isArray(message.parts) ? message.parts : [] + // AG-UI content parts: `{ type: 'text', text }` and binary parts. + const parts = Array.isArray(message.content) ? message.content : [] return parts .map((part: unknown) => - isRecord(part) && part.type === 'text' && typeof part.content === 'string' - ? part.content + isRecord(part) && part.type === 'text' && typeof part.text === 'string' + ? part.text : '', ) .join('') diff --git a/packages/ai-harness/tests/resume-edges.test.ts b/packages/ai-harness/tests/resume-edges.test.ts new file mode 100644 index 0000000000..7c9367ec33 --- /dev/null +++ b/packages/ai-harness/tests/resume-edges.test.ts @@ -0,0 +1,279 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { EventType, chat, defineAgent } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, harnessText } from '../src' +import { + INTERRUPTED_TOOL_RESULT, + findCrashedRuns, + repairTranscript, +} from '../src/resume' +import { messageTexts, mockAdapter, text, untilAborted } from './helpers' +import type { RunRecord, StreamChunk } from '@tanstack/ai' +import type { HarnessPersistence } from '../src/host' + +async function crash( + persistence: ReturnType, + runId: string, + startedAt: number, + pendingTools: NonNullable['pendingTools'] = [], +) { + await persistence.stores.runs.createOrResume({ + runId, + threadId: 't1', + startedAt, + }) + await persistence.stores.runs.update(runId, { + leaseOwner: 'host-gone', + leaseExpiresAt: Date.now() - 1_000, + checkpoint: { at: startedAt, pendingTools }, + }) +} + +describe('crash recovery edges', () => { + it('finds nothing when the run store cannot list a thread', async () => { + const noRuns = { stores: {} } as unknown as HarnessPersistence + expect(await findCrashedRuns(noRuns, 't1')).toEqual([]) + const noList = { + stores: { runs: { get: async () => null } }, + } as unknown as HarnessPersistence + expect(await findCrashedRuns(noList, 't1')).toEqual([]) + }) + + it('continues only the newest crashed turn and fails the older ones', async () => { + const persistence = memoryPersistence() + await persistence.stores.messages.saveThread('t1', [ + { id: 'u1', role: 'user', content: 'hello' }, + ]) + await crash(persistence, 'old-run', Date.now() - 120_000) + await crash(persistence, 'new-run', Date.now() - 60_000) + const { adapter, calls } = mockAdapter([() => text('picked up')]) + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ name: 'test/two-crashes', adapter }), + { threadId: 't1' }, + ) + await vi.waitFor(() => expect(calls).toHaveLength(1)) + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + + expect(await persistence.stores.runs.get('old-run')).toMatchObject({ + status: 'failed', + error: { message: 'The host stopped during this turn.' }, + }) + expect(await persistence.stores.runs.get('new-run')).toMatchObject({ + status: 'failed', + error: { + message: + 'The host stopped. The session continued this turn in a new run.', + }, + }) + await host.close() + }) + + it('leaves the transcript alone when no tool needs a note', async () => { + const persistence = memoryPersistence() + const history = [ + { id: 'u1', role: 'user' as const, content: 'go' }, + { + id: 'a1', + role: 'assistant' as const, + content: '', + toolCalls: [ + { + id: 'call-a', + type: 'function' as const, + function: { name: 'charge', arguments: '{}' }, + }, + ], + }, + { + id: 't1', + role: 'tool' as const, + toolCallId: 'call-a', + content: 'done', + }, + ] + await persistence.stores.messages.saveThread('t1', history) + const save = vi.spyOn(persistence.stores.messages, 'saveThread') + + // Only safe tools pending: nothing to note. + await repairTranscript(persistence, { + runId: 'r', + threadId: 't1', + status: 'running', + startedAt: 1, + checkpoint: { + at: 1, + pendingTools: [ + { toolCallId: 'call-b', name: 'lookup', replay: 'safe' }, + ], + }, + }) + // A never-replay tool that already has a result: nothing to note. + await repairTranscript(persistence, { + runId: 'r', + threadId: 't1', + status: 'running', + startedAt: 1, + checkpoint: { + at: 1, + pendingTools: [ + { toolCallId: 'call-a', name: 'charge', replay: 'never' }, + ], + }, + }) + // No checkpoint at all. + await repairTranscript(persistence, { + runId: 'r', + threadId: 't1', + status: 'running', + startedAt: 1, + }) + expect(save).not.toHaveBeenCalled() + expect( + JSON.stringify(await persistence.stores.messages.loadThread('t1')), + ).not.toContain(JSON.stringify(INTERRUPTED_TOOL_RESULT)) + }) + + it('releases the lease timer when a turn fails', async () => { + const persistence = memoryPersistence() + const { adapter } = mockAdapter([ + () => + (async function* (): AsyncGenerator { + yield { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: Date.now(), + } + throw new Error('model down') + })(), + ]) + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ name: 'test/fails', adapter }), + { threadId: 't-fail' }, + ) + const turn = session.prompt('hi') + await expect(turn).rejects.toThrow('model down') + expect(turn.status()).toBe('failed') + await host.close() + }) +}) + +describe('harnessText edges', () => { + it('uses a memory host by default and reads text content parts', async () => { + const inner = mockAdapter([() => text('parts answer'), () => text('empty')]) + const studio = defineHarness({ name: 'test/parts', adapter: inner.adapter }) + const model = harnessText(studio) + const collect = async (messages: Array) => { + let out = '' + for await (const chunk of chat({ adapter: model, messages })) { + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) out += chunk.delta + } + return out + } + expect( + await collect([ + { + role: 'user', + content: [ + { type: 'text', content: 'one ' }, + { + type: 'image', + source: { type: 'url', value: 'https://x/y.png' }, + }, + { type: 'text', content: 'two' }, + ], + }, + ]), + ).toBe('parts answer') + expect(messageTexts(inner.calls[0])).toEqual(['one two']) + await expect(model.structuredOutput({} as never)).rejects.toThrow( + 'does not support structured output', + ) + }) + + it('reports an inner failure as a run error', async () => { + const inner = mockAdapter([ + () => + (async function* (): AsyncGenerator { + yield { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: Date.now(), + } + throw new Error('inner broke') + })(), + ]) + const studio = defineHarness({ + name: 'test/inner-fail', + adapter: inner.adapter, + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const chunks: Array = [] + for await (const chunk of harnessText(studio, { host }).chatStream({ + model: 'test/inner-fail', + messages: [{ role: 'assistant', content: 'no user message here' }], + } as never)) { + chunks.push(chunk) + } + expect(chunks.at(-1)).toMatchObject({ + type: EventType.RUN_ERROR, + message: 'inner broke', + }) + await host.close() + }) + + it('cancels the inner turn when the outer request aborts', async () => { + const inner = mockAdapter([untilAborted()]) + const studio = defineHarness({ + name: 'test/inner-abort', + adapter: inner.adapter, + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const controller = new AbortController() + const stream = harnessText(studio, { host }).chatStream({ + model: 'test/inner-abort', + threadId: 'outer', + messages: [{ role: 'user', content: 'wait' }], + request: { signal: controller.signal }, + } as never) + const first = await stream[Symbol.asyncIterator]().next() + expect(first.value).toMatchObject({ type: EventType.RUN_STARTED }) + await vi.waitFor(() => expect(inner.calls).toHaveLength(1)) + controller.abort() + for await (const _chunk of stream) { + // Drains until the cancelled inner turn ends. + } + await host.close() + }) +}) + +describe('session agent and operation lookups', () => { + it('finds agents and operations by name and id', async () => { + const echo = defineAgent({ + name: 'echo', + description: 'Echoes', + inputSchema: z.object({ text: z.string() }), + run: async (ctx) => ctx.input.text.toUpperCase(), + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/lookups', + adapter: mockAdapter([]).adapter, + agents: [echo], + }), + { threadId: 't' }, + ) + expect(session.agent('missing')).toBeUndefined() + const handle = session.agent('echo') + const run = handle?.run({ text: 'hi' }) + expect(await run).toBe('HI') + expect(session.operation(run?.id ?? '')).toBe(run) + expect(session.operation('nope')).toBeUndefined() + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/transport-edges.test.ts b/packages/ai-harness/tests/transport-edges.test.ts new file mode 100644 index 0000000000..95c842698b --- /dev/null +++ b/packages/ai-harness/tests/transport-edges.test.ts @@ -0,0 +1,596 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { defineAgent, toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + applyInput, + capabilitiesOf, + createHarnessHandler, + createHarnessHost, + defineHarness, + handleHarnessSocket, + parseControlFrame, + parseHarnessInput, +} from '../src' +import { createHarnessClient } from '../src/client' +import { gate, mockAdapter, text, toolCall, untilAborted } from './helpers' +import type { WebSocketLike } from '@tanstack/ai' +import type { HostFrame } from '../src' + +const pricer = defineAgent({ + name: 'pricer', + description: 'Prices a vendor', + produces: 'data', + inputSchema: z.object({ vendor: z.string() }), + run: async (ctx) => ({ cents: ctx.input.vendor.length }), +}) +const lookup = toolDefinition({ + name: 'lookup', + description: 'Look up a fact', + inputSchema: z.object({ q: z.string() }), +}).server(async () => 'found') +const remove = toolDefinition({ + name: 'remove', + description: 'Remove a file', + needsApproval: true, + inputSchema: z.object({ path: z.string() }), +}).server(async () => ({ removed: true })) + +const auth = { authorization: 'Bearer good' } + +function setup(replies: Parameters[0] = []) { + const { adapter, calls } = mockAdapter(replies) + const harness = defineHarness({ + name: 'test/transport', + adapter, + tools: [lookup, remove], + agents: [pricer], + subagents: { agents: [pricer] }, + expose: { agents: ['pricer'] }, + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const handler = createHarnessHandler({ + host, + harness, + authorize: (request) => + request.headers.get('authorization') === 'Bearer good' + ? { id: 'user-1' } + : null, + canAccess: (principal, threadId) => threadId.startsWith(principal.id), + }) + const call = (path: string, init: RequestInit = {}) => + handler( + new Request(`http://x/api/harness/${path}`, { + ...init, + headers: { + ...auth, + 'content-type': 'application/json', + ...init.headers, + }, + }), + ) + return { handler, host, harness, calls, call } +} + +const runBody = (threadId: string, extra: Record = {}) => + JSON.stringify({ + threadId, + runId: 'client-run', + messages: [], + tools: [], + context: [], + state: {}, + forwardedProps: {}, + ...extra, + }) + +describe('HTTP handler edges', () => { + it('answers bad requests with 400, 403, and 404', async () => { + const { host, call } = setup() + const statusOf = async (path: string, init?: RequestInit) => + (await call(path, init)).status + + expect(await statusOf('events')).toBe(400) + expect(await statusOf('events?threadId=someone-else')).toBe(403) + expect(await statusOf('snapshot')).toBe(400) + expect(await statusOf('nowhere')).toBe(404) + expect(await statusOf('run', { method: 'GET' })).toBe(404) + + const control = (body: string) => + statusOf('control', { method: 'POST', body }) + expect(await control('null')).toBe(400) + expect(await control('{"input":{}}')).toBe(400) + expect(await control('{"threadId":7}')).toBe(400) + expect(await control('{"threadId":"user-1-a"}')).toBe(400) + expect(await control('not json')).toBe(400) + expect( + await control( + '{"threadId":"someone-else","input":{"op":"prompt","message":"hi"}}', + ), + ).toBe(403) + + expect( + await statusOf('run', { method: 'POST', body: runBody('someone-else') }), + ).toBe(403) + const noMessage = await call('run', { + method: 'POST', + body: runBody('user-1-r', { + messages: [{ id: 'a', role: 'assistant', content: 'hello' }], + }), + }) + expect(noMessage.status).toBe(400) + expect(await noMessage.json()).toEqual({ error: 'no user message' }) + await host.close() + }) + + it('reads the prompt from AG-UI content parts', async () => { + const { host, call, calls } = setup([() => text('from parts')]) + const response = await call('run', { + method: 'POST', + body: runBody('user-1-parts', { + messages: [ + { + id: 'u1', + role: 'user', + content: [ + { type: 'text', text: 'part one, ' }, + { type: 'binary', mimeType: 'image/png', data: 'AA==' }, + { type: 'text', text: 'part two' }, + ], + }, + ], + }), + }) + expect(await response.text()).toContain('from parts') + expect(JSON.stringify(calls[0].messages)).toContain('part one, part two') + await host.close() + }) + + it('continues an interrupted turn when a run request carries resume', async () => { + const { host, call } = setup([ + () => toolCall('remove', { path: 'a.txt' }), + () => text('removed it'), + ]) + const control = await call('control', { + method: 'POST', + body: JSON.stringify({ + threadId: 'user-1-resume', + input: { op: 'prompt', message: 'remove a.txt' }, + }), + }) + expect((await control.json()).status).toBe('accepted') + const snapshotOf = async () => + (await call('snapshot?threadId=user-1-resume')).json() + await vi.waitFor(async () => + expect((await snapshotOf()).status).toBe('requires_action'), + ) + const [interrupt] = (await snapshotOf()).pendingInterrupts + + // A thread with no stopped turn has nothing to resume. + const nothing = await call('run', { + method: 'POST', + body: runBody('user-1-idle', { + resume: [{ interruptId: 'nope', status: 'resolved', payload: true }], + }), + }) + expect(nothing.status).toBe(409) + + const resumed = await call('run', { + method: 'POST', + body: runBody('user-1-resume', { + resume: [ + { interruptId: interrupt.id, status: 'resolved', payload: true }, + ], + }), + }) + expect(resumed.headers.get('content-type')).toContain('text/event-stream') + const body = await resumed.text() + expect(body).toContain('removed it') + expect(body).toContain('harness.operation.finished') + await host.close() + }) + + it('streams session events from a `from` query cursor', async () => { + const { host, call } = setup([() => text('one'), () => text('two')]) + const prompt = (message: string) => + call('control', { + method: 'POST', + body: JSON.stringify({ + threadId: 'user-1-from', + input: { op: 'prompt', message }, + }), + }) + await prompt('first') + await vi.waitFor(async () => + expect( + (await (await call('snapshot?threadId=user-1-from')).json()).status, + ).toBe('idle'), + ) + const snapshot = await (await call('snapshot?threadId=user-1-from')).json() + await prompt('second') + const response = await call( + `events?threadId=user-1-from&from=${snapshot.cursor}`, + ) + const reader = response.body!.getReader() + let seen = '' + while (!seen.includes('harness.operation.finished')) { + const { value, done } = await reader.read() + if (done) break + seen += new TextDecoder().decode(value) + } + await reader.cancel() + expect(seen).toContain('two') + expect(seen).not.toContain('"delta":"one"') + await host.close() + }) +}) + +describe('WebSocket edges', () => { + function fakeSocket(options: { failSend?: boolean } = {}) { + const handlers: Record void>> = {} + const sent: Array = [] + const closed: Array = [] + const socket: WebSocketLike = { + send: (data) => { + if (options.failSend) throw new Error('socket gone') + sent.push(JSON.parse(data)) + }, + close: (code?: number) => { + closed.push(code) + handlers.close?.forEach((handler) => handler()) + }, + addEventListener: (type: string, handler: (event?: any) => void) => { + ;(handlers[type] ??= []).push(handler) + }, + } + const receive = (frame: unknown) => + handlers.message?.forEach((handler) => + handler({ + data: typeof frame === 'string' ? frame : JSON.stringify(frame), + }), + ) + const fail = () => handlers.error?.forEach((handler) => handler()) + return { socket, sent, closed, receive, fail } + } + + it('refuses a thread the principal cannot access', async () => { + const { host, harness } = setup() + const { socket, sent, closed, receive } = fakeSocket() + handleHarnessSocket({ + host, + harness, + socket, + principal: { id: 'user-1' }, + canAccess: () => false, + }) + receive({ type: 'harness.subscribe', threadId: 'user-1-x' }) + await vi.waitFor(() => expect(closed).toEqual([4403])) + expect(sent).toEqual([{ type: 'harness.error', message: 'forbidden' }]) + await host.close() + }) + + it('reports a second subscribe and bad frames, and resumes from a cursor', async () => { + const { host, harness } = setup([() => text('hello')]) + const first = fakeSocket() + handleHarnessSocket({ + host, + harness, + socket: first.socket, + principal: { id: 'user-1' }, + }) + first.receive({ type: 'harness.subscribe', threadId: 'user-1-ws2' }) + await vi.waitFor(() => + expect(first.sent.some((frame) => frame.type === 'harness.hello')).toBe( + true, + ), + ) + first.receive({ type: 'harness.subscribe', threadId: 'user-1-ws2' }) + await vi.waitFor(() => + expect(first.sent.at(-1)).toEqual({ + type: 'harness.error', + message: 'Already subscribed.', + }), + ) + first.receive('[1, 2]') + await vi.waitFor(() => + expect(first.sent.at(-1)).toMatchObject({ type: 'harness.error' }), + ) + first.receive({ + type: 'harness.input', + requestId: 'p1', + input: { op: 'prompt', message: 'hi' }, + }) + await vi.waitFor(() => + expect( + first.sent.some( + (frame) => + frame.type === 'harness.event' && + frame.event.type === 'CUSTOM' && + frame.event.name === 'harness.operation.finished', + ), + ).toBe(true), + ) + const cursor = first.sent + .filter((frame) => frame.type === 'harness.event') + .at(-1) + first.fail() + + // A new socket that resumes after the last cursor sees nothing old. + const second = fakeSocket() + handleHarnessSocket({ + host, + harness, + socket: second.socket, + principal: { id: 'user-1' }, + }) + second.receive({ + type: 'harness.subscribe', + threadId: 'user-1-ws2', + from: cursor && 'cursor' in cursor ? cursor.cursor : undefined, + }) + await vi.waitFor(() => + expect(second.sent.some((frame) => frame.type === 'harness.hello')).toBe( + true, + ), + ) + await new Promise((resolve) => setTimeout(resolve, 20)) + expect( + second.sent.filter((frame) => frame.type === 'harness.event'), + ).toEqual([]) + second.socket.close() + await host.close() + }) + + it('stops reading when the socket cannot send', async () => { + const { host, harness } = setup() + const { socket, receive } = fakeSocket({ failSend: true }) + handleHarnessSocket({ host, harness, socket, principal: { id: 'user-1' } }) + receive({ type: 'harness.subscribe', threadId: 'user-1-dead' }) + // Nothing to assert on the socket: the send failure must not throw. + await new Promise((resolve) => setTimeout(resolve, 20)) + await host.close() + }) +}) + +describe('protocol parsing and inputs', () => { + it('checks each input op for its fields', () => { + expect(() => parseHarnessInput(null)).toThrow('known op') + expect(() => parseHarnessInput({ op: 'launch' })).toThrow('known op') + expect(() => parseHarnessInput({ op: 'steer' })).toThrow( + 'steer needs a message', + ) + expect(() => parseHarnessInput({ op: 'resolve' })).toThrow( + 'resolve needs a resume array', + ) + expect(() => parseHarnessInput({ op: 'agent' })).toThrow( + 'agent needs an agent name', + ) + expect( + parseHarnessInput({ + op: 'followUp', + message: [{ type: 'text', content: 'x' }], + }), + ).toMatchObject({ op: 'followUp' }) + expect(parseHarnessInput({ op: 'cancel' })).toEqual({ op: 'cancel' }) + }) + + it('parses subscribe frames with a cursor and refuses non-objects', () => { + expect( + parseControlFrame( + '{"type":"harness.subscribe","threadId":"t","from":"4"}', + ), + ).toEqual({ type: 'harness.subscribe', threadId: 't', from: '4' }) + expect( + parseControlFrame('{"type":"harness.subscribe","threadId":"t","from":4}'), + ).toEqual({ type: 'harness.subscribe', threadId: 't' }) + expect(() => parseControlFrame('7')).toThrow('Invalid frame.') + expect(() => parseControlFrame('null')).toThrow('Invalid frame.') + }) + + it('applies steer, follow-up, resolve, cancel, busy prompts, and agents', async () => { + const release = gate() + const { host, harness } = setup([ + () => + (async function* () { + await release.opened + yield* text('slow answer') + })(), + () => text('queued answer'), + () => text('follow-up answer'), + ]) + const session = await host.open(harness, { threadId: 'user-1-apply' }) + + const first = await applyInput(harness, session, { + op: 'prompt', + message: 'slow', + }) + expect(first.status).toBe('accepted') + const queued = await applyInput(harness, session, { + op: 'prompt', + message: 'next', + busy: 'queue', + }) + expect(queued.status).toBe('queued') + expect( + await applyInput(harness, session, { op: 'steer', message: 'faster' }), + ).toHaveProperty('status') + expect( + await applyInput(harness, session, { op: 'followUp', message: 'after' }), + ).toHaveProperty('status') + release.open() + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + + expect( + await applyInput(harness, session, { op: 'resolve', resume: [] }), + ).toMatchObject({ status: 'rejected' }) + expect(await applyInput(harness, session, { op: 'cancel' })).toHaveProperty( + 'status', + ) + // defineHarness refuses this, but a hand-built harness object can name an + // exposed agent that the session does not have. + const loose = { ...harness, expose: { agents: ['pricer', 'ghost'] } } + expect( + await applyInput(loose, session, { op: 'agent', agent: 'ghost' }), + ).toEqual({ inputId: '', status: 'rejected', reason: 'unknown_agent' }) + const detached = await applyInput(harness, session, { + op: 'agent', + agent: 'pricer', + input: { vendor: 'acme' }, + detached: true, + }) + expect(detached).toMatchObject({ status: 'accepted' }) + await session.operation(detached.operationId ?? '') + await host.close() + }) + + it('describes tools, subagents, and what exposed agents produce', () => { + const { harness } = setup() + const capabilities = capabilitiesOf(harness) + expect(capabilities.tools.items.map((tool) => tool.name)).toEqual([ + 'lookup', + 'remove', + ]) + expect(capabilities.multiAgent).toMatchObject({ + supported: true, + subagents: [{ name: 'pricer', description: 'Prices a vendor' }], + }) + expect(capabilities.custom.tanstack.agents[0]).toMatchObject({ + name: 'pricer', + produces: 'data', + }) + const bare = capabilitiesOf( + defineHarness({ name: 'test/bare', adapter: mockAdapter([]).adapter }), + ) + expect(bare.multiAgent.supported).toBe(false) + expect(bare.tools.items).toEqual([]) + }) +}) + +describe('createHarnessClient edges', () => { + it('sends every input op with headers from a function', async () => { + const { host, harness, handler } = setup([ + () => text('a'), + () => text('b'), + () => text('c'), + ]) + const seenAuth: Array = [] + const client = createHarnessClient({ + url: 'http://x/api/harness', + threadId: 'user-1-client', + headers: () => auth, + fetch: (input, init) => { + const request = new Request(input, init) + seenAuth.push(request.headers.get('authorization')) + return handler(request) + }, + }) + expect((await client.prompt('hi', { busy: 'queue' })).status).toBe( + 'accepted', + ) + await client.steer('go on') + await client.followUp('then') + expect((await client.resolve([])).status).toBe('rejected') + await client.cancel() + await client.cancel('op-missing') + expect( + (await client.agents.pricer.start({ vendor: 'x' }, { detached: true })) + .status, + ).toBe('accepted') + expect(seenAuth.every((value) => value === 'Bearer good')).toBe(true) + // Symbol keys on the agents proxy are not agents. + expect( + (client.agents as unknown as Record)[Symbol.iterator], + ).toBeUndefined() + await host.close() + }) + + it('reports errors from the handler and a failed snapshot', async () => { + const client = createHarnessClient({ + url: 'http://x/api/harness/', + threadId: 't', + fetch: async (input) => + String(input).includes('snapshot') + ? new Response('nope', { status: 500 }) + : new Response('"plain"', { status: 502, statusText: 'Bad Gateway' }), + }) + await expect(client.prompt('hi')).rejects.toThrow( + 'Harness request failed (502): Bad Gateway', + ) + await expect(client.snapshot()).rejects.toThrow( + 'Harness snapshot failed (500)', + ) + }) + + it('reconnects after a failed events request and resumes from the cursor', async () => { + const urls: Array = [] + const frame = (cursor: string) => + `id: ${cursor}\ndata: ${JSON.stringify({ + type: 'harness.event', + cursor, + operationId: 'op', + event: { type: 'CUSTOM', name: 'tick', value: cursor, timestamp: 1 }, + })}\n\n` + let calls = 0 + const client = createHarnessClient({ + url: 'http://x/api/harness', + threadId: 't', + reconnectDelayMs: 1, + fetch: async (input) => { + urls.push(String(input)) + calls += 1 + if (calls === 1) return new Response('down', { status: 503 }) + if (calls === 2) throw new Error('network') + const body = calls === 3 ? `: comment\n\n${frame('1')}` : frame('2') + return new Response(body, { status: 200 }) + }, + }) + const controller = new AbortController() + const seen: Array = [] + for await (const entry of client.events({ + from: '0', + signal: controller.signal, + })) { + seen.push(entry.cursor) + if (seen.length === 2) controller.abort() + } + expect(seen).toEqual(['1', '2']) + expect(urls[0]).toContain('from=0') + expect(urls.at(-1)).toContain('from=1') + }) + + it('stops quietly when aborted during a request', async () => { + const controller = new AbortController() + const client = createHarnessClient({ + url: 'http://x/api/harness', + threadId: 't', + fetch: async () => { + controller.abort() + throw new Error('aborted') + }, + }) + const seen: Array = [] + for await (const entry of client.events({ signal: controller.signal })) { + seen.push(entry) + } + expect(seen).toEqual([]) + }) +}) + +describe('an operation that waits for cancel', () => { + it('cancels through a control input', async () => { + const { host, harness } = setup([untilAborted()]) + const session = await host.open(harness, { threadId: 'user-1-cancel' }) + const operation = session.prompt('wait') + await vi.waitFor(() => expect(session.snapshot().status).toBe('running')) + await applyInput(harness, session, { + op: 'cancel', + operationId: operation.id, + }) + await operation.then( + () => undefined, + () => undefined, + ) + expect(operation.status()).toBe('cancelled') + await host.close() + }) +}) From 37c04d30dc0bb020e6d3abddbaa10fde88299c8d Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 15:08:34 +0200 Subject: [PATCH 018/286] feat(ai-mcp, ai-harness): MCP connectors with browser sign-in, run-time tool discovery, and media in the example --- .changeset/harness-p7-connectors.md | 11 + docs/config.json | 5 + docs/harness/dashboard.md | 2 +- docs/harness/deploy.md | 2 +- docs/harness/mcp.md | 121 ++++++ docs/harness/subagents.md | 2 +- examples/harness-cli/.env.example | 7 + examples/harness-cli/.gitignore | 2 + examples/harness-cli/README.md | 13 +- examples/harness-cli/package.json | 4 +- examples/harness-cli/src/cli.ts | 9 +- examples/harness-cli/src/credentials.ts | 53 +++ examples/harness-cli/src/harness.ts | 37 +- examples/harness-cli/src/media.ts | 90 +++++ packages/ai-harness/src/define.ts | 1 + packages/ai-harness/src/index.ts | 1 + packages/ai-harness/src/oauth.ts | 62 +++ packages/ai-harness/src/plugins.ts | 18 + packages/ai-harness/src/session.ts | 54 ++- .../ai-harness/tests/discover-tools.test.ts | 165 ++++++++ packages/ai-mcp/package.json | 14 + packages/ai-mcp/src/connector.ts | 291 ++++++++++++++ packages/ai-mcp/tests/connector.test.ts | 361 ++++++++++++++++++ packages/ai-mcp/vite.config.ts | 7 +- packages/ai-persistence/src/types.ts | 5 + pnpm-lock.yaml | 12 + 26 files changed, 1332 insertions(+), 17 deletions(-) create mode 100644 .changeset/harness-p7-connectors.md create mode 100644 docs/harness/mcp.md create mode 100644 examples/harness-cli/.env.example create mode 100644 examples/harness-cli/.gitignore create mode 100644 examples/harness-cli/src/credentials.ts create mode 100644 examples/harness-cli/src/media.ts create mode 100644 packages/ai-harness/tests/discover-tools.test.ts create mode 100644 packages/ai-mcp/src/connector.ts create mode 100644 packages/ai-mcp/tests/connector.test.ts diff --git a/.changeset/harness-p7-connectors.md b/.changeset/harness-p7-connectors.md new file mode 100644 index 0000000000..17ea459419 --- /dev/null +++ b/.changeset/harness-p7-connectors.md @@ -0,0 +1,11 @@ +--- +'@tanstack/ai-mcp': minor +'@tanstack/ai-harness': minor +'@tanstack/ai-persistence': minor +--- + +`@tanstack/ai-mcp/connector` adds `mcpConnector`: a harness plugin that signs the user in to a remote MCP server (Notion, Linear, and others) with OAuth, then gives the model the server's tools. `/connect ` registers a client, signs in with PKCE on a `127.0.0.1` loopback, and keeps the token in the credential store. Tools the server does not mark read-only ask for approval. + +`@tanstack/ai-harness` adds `discoverTools` to plugins, for tools found at run time, and `startLoopbackReceiver` for OAuth redirects. A harness turn now runs up to 50 model calls by default (`agentLoopStrategy` still overrides it). + +`@tanstack/ai-persistence`: an `oauth` credential can keep the OAuth `client` it was issued to. diff --git a/docs/config.json b/docs/config.json index 9ac515446a..8b9dea312f 100644 --- a/docs/config.json +++ b/docs/config.json @@ -839,6 +839,11 @@ "to": "harness/auth", "addedAt": "2026-09-26" }, + { + "label": "Use MCP servers", + "to": "harness/mcp", + "addedAt": "2026-09-26" + }, { "label": "Run agents from a harness", "to": "harness/subagents", diff --git a/docs/harness/dashboard.md b/docs/harness/dashboard.md index 8bc06bcd83..3fab5e3e0a 100644 --- a/docs/harness/dashboard.md +++ b/docs/harness/dashboard.md @@ -1,7 +1,7 @@ --- title: Self-host the dashboard id: harness-dashboard -order: 10 +order: 11 description: "Watch and steer your harness sessions from a browser or a phone. Agents dial out to your dashboard server, so they need no open port." keywords: - tanstack ai diff --git a/docs/harness/deploy.md b/docs/harness/deploy.md index d6d22b92db..e75187f2a4 100644 --- a/docs/harness/deploy.md +++ b/docs/harness/deploy.md @@ -1,7 +1,7 @@ --- title: Deploy a harness id: harness-deploy -order: 9 +order: 10 description: "Run a harness in your server, as a worker process, on another machine, or as a single executable." keywords: - tanstack ai diff --git a/docs/harness/mcp.md b/docs/harness/mcp.md new file mode 100644 index 0000000000..52dff3f012 --- /dev/null +++ b/docs/harness/mcp.md @@ -0,0 +1,121 @@ +--- +title: Use MCP servers +id: harness-mcp +order: 8 +description: "Connect a harness to Notion, Linear, or another MCP server that signs in with OAuth. The user signs in once in the browser, and the tools appear after sign-in." +keywords: + - tanstack ai + - harness + - mcp + - connectors + - oauth + - notion + - linear +--- + +Your agent needs to read the user's Linear issues and Notion pages. Both services run MCP servers that sign in with OAuth, so you do not register an app or copy an API key. `mcpConnector` adds `/connect linear`, keeps the token in your credential store, and gives the model the server's tools after sign-in. + +## 1. Add the connectors + + + +react: @tanstack/ai-harness @tanstack/ai-mcp +vue: @tanstack/ai-harness @tanstack/ai-mcp +solid: @tanstack/ai-harness @tanstack/ai-mcp +svelte: @tanstack/ai-harness @tanstack/ai-mcp +preact: @tanstack/ai-harness @tanstack/ai-mcp +angular: @tanstack/ai-harness @tanstack/ai-mcp +octane: @tanstack/ai-harness @tanstack/ai-mcp +vanilla: @tanstack/ai-harness @tanstack/ai-mcp + + + +```ts group=harness-mcp +import { defineHarness } from '@tanstack/ai-harness' +import { mcpConnector } from '@tanstack/ai-mcp/connector' +import { openaiText } from '@tanstack/ai-openai' + +const notion = mcpConnector({ + id: 'notion', + label: 'Notion', + url: 'https://mcp.notion.com/mcp', +}) +const linear = mcpConnector({ + id: 'linear', + label: 'Linear', + url: 'https://mcp.linear.app/mcp', +}) + +export const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-5.6'), + plugins: () => [notion, linear], +}) +``` + +Any MCP server that speaks Streamable HTTP and OAuth works the same way. Give it an `id`, a `label`, and its `url`. + +## 2. Sign in + +1. Run `/connect linear` in the CLI. +2. The browser opens the Linear consent page. Approve it. +3. The CLI prints `Connected to Linear.` +4. Ask for your issues. The next turn has the Linear tools. + +Before sign-in, the model knows that Linear is not connected. If the user asks for Linear, the model tells them to run `/connect linear`. + +The sign-in follows the MCP authorization spec: + +- The connector reads the server's OAuth metadata and registers a client for this sign-in. +- It signs in with PKCE and a random `state`, and listens on `127.0.0.1` for one callback. +- The token and the registered client go into `stores.credentials`. The next start of the harness uses them, so the user signs in once. +- `/disconnect linear` deletes them. + +The model never sees the token. To keep sign-ins after a restart, give the host a credential store. [Auth and connectors](./auth) shows one. + +## 3. Choose which tools ask first + +The connector names each tool `_`, for example `linear_list_issues`. A tool that the server does not mark read-only asks for approval before it runs. To change that rule, pass `needsApproval`: + +```ts group=harness-mcp +export const linearWithoutWrites = mcpConnector({ + id: 'linear', + label: 'Linear', + url: 'https://mcp.linear.app/mcp', + // Ask before every tool whose name starts with "save". + needsApproval: (tool) => tool.name.startsWith('save'), +}) +``` + +Other options: + +- `prefix`: the tool name prefix. Default: the `id`. +- `scopes`: the OAuth scopes to ask for. Default: what the server offers. +- `clientName`: the app name on the consent page. Default: `TanStack AI Harness`. + +## Tools that appear later + +A plugin can give tools that it finds at run time, like the tools of a server that the user signed in to after the session started. Return them from `discoverTools`. The session asks each plugin before every turn: + +```ts group=harness-mcp +import { definePlugin } from '@tanstack/ai-harness' +import type { AnyTool } from '@tanstack/ai' + +const found: Array = [] + +export const lateTools = definePlugin({ + name: 'acme/late-tools', + setup: () => ({ + discoverTools: () => found, + }), +}) +``` + +- A tool whose name is already taken is skipped. The first tool keeps the name. +- If `discoverTools` throws, the turn runs without that plugin's tools, and clients get a `harness.plugin.warning` event. + +## What you have now + +- `/connect` and `/disconnect` for Notion, Linear, and other MCP servers. +- One browser sign-in per service, kept in your credential store. +- Server tools in the model's tool list, with approval for tools that change data. diff --git a/docs/harness/subagents.md b/docs/harness/subagents.md index 71b89b77ea..169770cb0d 100644 --- a/docs/harness/subagents.md +++ b/docs/harness/subagents.md @@ -1,7 +1,7 @@ --- title: Run agents from a harness id: harness-subagents -order: 8 +order: 9 description: "Start typed agents from commands and plugins, run them in groups, call a whole harness as a child, and keep the tree within limits." keywords: - tanstack ai diff --git a/examples/harness-cli/.env.example b/examples/harness-cli/.env.example new file mode 100644 index 0000000000..bb89717ef3 --- /dev/null +++ b/examples/harness-cli/.env.example @@ -0,0 +1,7 @@ +# Copy to .env. Without keys, a demo model answers. +# The main model: OpenAI first, then Anthropic. +# OPENAI_API_KEY also turns on the image and video tools. +OPENAI_API_KEY= +ANTHROPIC_API_KEY= +# Optional: make videos with Grok Imagine instead of OpenAI Sora. +XAI_API_KEY= diff --git a/examples/harness-cli/.gitignore b/examples/harness-cli/.gitignore new file mode 100644 index 0000000000..ec9309579e --- /dev/null +++ b/examples/harness-cli/.gitignore @@ -0,0 +1,2 @@ +.env +playground/media/ diff --git a/examples/harness-cli/README.md b/examples/harness-cli/README.md index cd5f93c5fe..bd0afebc98 100644 --- a/examples/harness-cli/README.md +++ b/examples/harness-cli/README.md @@ -1,6 +1,6 @@ # Harness CLI example -A small coding agent in your terminal, built with `@tanstack/ai-harness`. It reads and edits files in `./playground`, asks before it writes a file or runs a command, keeps a todo list, and can switch models. +A small coding agent in your terminal, built with `@tanstack/ai-harness`. It reads and edits files in `./playground`, asks before it writes a file or runs a command, keeps a todo list, and can switch models. It can also read Notion and Linear, and make images and videos. ## Run it @@ -8,7 +8,7 @@ From the repo root: 1. `pnpm install` 2. `pnpm build:all` (the example uses the local packages) -3. Set a key: `OPENAI_API_KEY=...` or `ANTHROPIC_API_KEY=...`. Without one, a demo model answers. +3. Copy `.env.example` to `.env` and set `OPENAI_API_KEY` or `ANTHROPIC_API_KEY`. Without a key, a demo model answers. 4. `pnpm --filter harness-cli-example start` Try these: @@ -19,6 +19,15 @@ Try these: - `/agent haiku {"topic":"rain"}`: runs a typed agent in the background. - Press Esc to stop a long answer. Type while it works to steer it. +## Use Notion, Linear, images, and video + +1. Run `/connect notion`. Approve the consent page that opens in the browser. Do the same with `/connect linear`. +2. Ask: `find my latest Linear issue, look for a related Notion page, then make an image and a short video about it`. +3. The files land in `playground/media`. + +- Sign-ins are kept in `~/.tanstack-harness-example/credentials.json`, so you sign in once. `/disconnect notion` deletes one. +- Images use `OPENAI_API_KEY`. Videos use Grok Imagine when `XAI_API_KEY` is set, and OpenAI Sora when it is not. + ## Other modes - One prompt for scripts and CI: `pnpm --filter harness-cli-example start -p "list the files"` diff --git a/examples/harness-cli/package.json b/examples/harness-cli/package.json index 3038efe1ff..96c40cb7d6 100644 --- a/examples/harness-cli/package.json +++ b/examples/harness-cli/package.json @@ -4,7 +4,7 @@ "type": "module", "description": "A small coding agent in the terminal, built with @tanstack/ai-harness", "scripts": { - "start": "tsx src/cli.ts", + "start": "tsx --env-file-if-exists=.env src/cli.ts", "dashboard": "tanstack-ai-dashboard", "test:types": "tsc --noEmit" }, @@ -13,8 +13,10 @@ "@tanstack/ai-acp": "workspace:*", "@tanstack/ai-anthropic": "workspace:*", "@tanstack/ai-dashboard": "workspace:*", + "@tanstack/ai-grok": "workspace:*", "@tanstack/ai-harness": "workspace:*", "@tanstack/ai-harness-cli": "workspace:*", + "@tanstack/ai-mcp": "workspace:*", "@tanstack/ai-openai": "workspace:*", "@tanstack/ai-persistence": "workspace:*", "zod": "^4.2.0" diff --git a/examples/harness-cli/src/cli.ts b/examples/harness-cli/src/cli.ts index 774a33ba2e..a4b42ec82e 100644 --- a/examples/harness-cli/src/cli.ts +++ b/examples/harness-cli/src/cli.ts @@ -1,4 +1,11 @@ import { runCli } from '@tanstack/ai-harness-cli' +import { composePersistence, memoryPersistence } from '@tanstack/ai-persistence' +import { fileCredentials } from './credentials' import { assistant } from './harness' -process.exitCode = await runCli(assistant) +// Everything in memory, except sign-ins, which are kept in a file. +const persistence = composePersistence(memoryPersistence(), { + overrides: { credentials: fileCredentials() }, +}) + +process.exitCode = await runCli(assistant, { persistence }) diff --git a/examples/harness-cli/src/credentials.ts b/examples/harness-cli/src/credentials.ts new file mode 100644 index 0000000000..06c358356b --- /dev/null +++ b/examples/harness-cli/src/credentials.ts @@ -0,0 +1,53 @@ +import { mkdir, readFile, writeFile } from 'node:fs/promises' +import { homedir } from 'node:os' +import { dirname, join } from 'node:path' +import { defineCredentialStore } from '@tanstack/ai-persistence' +import type { Credential } from '@tanstack/ai-persistence' + +/** + * Sign-ins (Notion, Linear) saved in a file in your home folder, so they + * survive a restart. The file is readable by your user only. A production + * store would encrypt the values. + */ +export function fileCredentials( + file = join(homedir(), '.tanstack-harness-example', 'credentials.json'), +) { + type Saved = Record> + const load = async (): Promise => { + try { + const parsed: unknown = JSON.parse(await readFile(file, 'utf8')) + return typeof parsed === 'object' && parsed !== null + ? (parsed as Saved) + : {} + } catch { + return {} + } + } + const save = async (all: Saved) => { + await mkdir(dirname(file), { recursive: true }) + await writeFile(file, JSON.stringify(all, null, 2), { mode: 0o600 }) + } + const owner = (scope: { userId?: string; tenantId?: string }) => + `${scope.tenantId ?? '-'}/${scope.userId ?? '-'}` + + return defineCredentialStore({ + get: async (scope, id) => (await load())[owner(scope)]?.[id] ?? null, + set: async (scope, id, credential) => { + const all = await load() + all[owner(scope)] = { ...all[owner(scope)], [id]: credential } + await save(all) + }, + delete: async (scope, id) => { + const all = await load() + delete all[owner(scope)]?.[id] + await save(all) + }, + list: async (scope) => + Object.entries((await load())[owner(scope)] ?? {}).map( + ([id, credential]) => ({ + id, + type: credential.type, + }), + ), + }) +} diff --git a/examples/harness-cli/src/harness.ts b/examples/harness-cli/src/harness.ts index c87b1a6434..86699642c0 100644 --- a/examples/harness-cli/src/harness.ts +++ b/examples/harness-cli/src/harness.ts @@ -11,12 +11,16 @@ import { workspaceTools, } from '@tanstack/ai-harness/plugins' import { anthropicText } from '@tanstack/ai-anthropic' -import { openaiText } from '@tanstack/ai-openai' +import { mcpConnector } from '@tanstack/ai-mcp/connector' +import { grokVideo } from '@tanstack/ai-grok' +import { openaiText, openaiVideo } from '@tanstack/ai-openai' import { z } from 'zod' +import { imageAgent, videoAgent } from './media' import type { AnyTextAdapter } from '@tanstack/ai' // The agent works in ./playground, so it cannot touch the rest of your disk. const root = fileURLToPath(new URL('../playground', import.meta.url)) +const mediaDir = fileURLToPath(new URL('../playground/media', import.meta.url)) /** Without an API key, a stand-in model that explains how to add one. */ function demoModel(): AnyTextAdapter { @@ -116,15 +120,46 @@ const haiku = defineAgent({ }), }) +// The model calls the media agents as tools. Images use OpenAI. Videos use +// Grok Imagine when XAI_API_KEY is set, else OpenAI Sora. +const media = process.env.OPENAI_API_KEY + ? [ + imageAgent(mediaDir), + videoAgent( + mediaDir, + process.env.XAI_API_KEY + ? grokVideo('grok-imagine-video') + : openaiVideo('sora-2'), + ), + ] + : [] + +// Notion and Linear through their MCP servers. `/connect notion` signs in +// through the browser. No app setup or API key is needed. +const notion = mcpConnector({ + id: 'notion', + label: 'Notion', + url: 'https://mcp.notion.com/mcp', +}) +const linear = mcpConnector({ + id: 'linear', + label: 'Linear', + url: 'https://mcp.linear.app/mcp', +}) + export const assistant = defineHarness({ name: 'example/coder', description: 'A small coding agent that works in ./playground', adapter: main, systemPrompts: [ 'You are a careful coding agent. Read files before you edit them. Keep answers short.', + 'You can read Notion and Linear when they are connected, and make images and videos with the image and video tools. Media files are saved under ./playground/media.', ], agents: [haiku], + subagents: { agents: media }, plugins: () => [ + notion, + linear, permissions(), workspaceTools({ root }), todos(), diff --git a/examples/harness-cli/src/media.ts b/examples/harness-cli/src/media.ts new file mode 100644 index 0000000000..40b9329ae0 --- /dev/null +++ b/examples/harness-cli/src/media.ts @@ -0,0 +1,90 @@ +import { mkdir, writeFile } from 'node:fs/promises' +import { join } from 'node:path' +import { defineAgent, getVideoJobStatus } from '@tanstack/ai' +import { openaiImage } from '@tanstack/ai-openai' +import { z } from 'zod' +import type { AnyVideoAdapter } from '@tanstack/ai' + +/** A short file name from a prompt: `a-fox-in-the-snow-1714000000000`. */ +function fileName(prompt: string): string { + const slug = prompt + .toLowerCase() + .replace(/[^a-z0-9]+/g, '-') + .replace(/^-|-$/g, '') + .slice(0, 40) + return `${slug || 'media'}-${Date.now()}` +} + +async function download(url: string): Promise { + const response = await fetch(url) + if (!response.ok) throw new Error(`Download failed (${response.status}).`) + return Buffer.from(await response.arrayBuffer()) +} + +/** Makes an image with OpenAI and saves it as a PNG in `mediaDir`. */ +export function imageAgent(mediaDir: string) { + return defineAgent({ + name: 'image', + description: + 'Generates one image from a detailed visual prompt and saves it as a PNG file. Returns the file path.', + produces: 'image', + inputSchema: z.object({ + prompt: z.string().describe('A detailed description of the image'), + }), + run: async (ctx) => { + const result = await ctx.generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: ctx.input.prompt, + size: '1024x1024', + }) + const image = result.images[0] + if (!image) throw new Error('The image model returned no image.') + let bytes: Buffer + if (image.b64Json) bytes = Buffer.from(image.b64Json, 'base64') + else if (image.url) bytes = await download(image.url) + else throw new Error('The image model returned no image data.') + await mkdir(mediaDir, { recursive: true }) + const path = join(mediaDir, `${fileName(ctx.input.prompt)}.png`) + await writeFile(path, bytes) + return { saved: path, model: result.model } + }, + }) +} + +/** Makes a short video with `adapter` and saves it as an MP4 in `mediaDir`. */ +export function videoAgent(mediaDir: string, adapter: AnyVideoAdapter) { + return defineAgent({ + name: 'video', + description: + 'Generates one short video clip from a detailed visual prompt and saves it as an MP4 file. Takes a minute or two. Returns the file path.', + produces: 'video', + inputSchema: z.object({ + prompt: z + .string() + .describe('A detailed description of the scene and the motion'), + }), + run: async (ctx) => { + const { jobId } = await ctx.generateVideo({ + adapter, + prompt: ctx.input.prompt, + }) + // Video models work in the background, so poll until the job is done. + const deadline = Date.now() + 10 * 60_000 + let url: string | undefined + while (!url) { + if (ctx.abortSignal?.aborted) throw new Error('Stopped.') + if (Date.now() > deadline) + throw new Error('The video took longer than 10 minutes.') + await new Promise((resolve) => setTimeout(resolve, 4000)) + const status = await getVideoJobStatus({ adapter, jobId }) + if (status.status === 'failed') + throw new Error(status.error ?? 'The video job failed.') + if (status.status === 'completed') url = status.url + } + await mkdir(mediaDir, { recursive: true }) + const path = join(mediaDir, `${fileName(ctx.input.prompt)}.mp4`) + await writeFile(path, await download(url)) + return { saved: path, model: adapter.model } + }, + }) +} diff --git a/packages/ai-harness/src/define.ts b/packages/ai-harness/src/define.ts index aef7787bd7..a2bacc3968 100644 --- a/packages/ai-harness/src/define.ts +++ b/packages/ai-harness/src/define.ts @@ -35,6 +35,7 @@ export interface HarnessConfig< systemPrompts?: Array tools?: ReadonlyArray middleware?: ReadonlyArray + /** When a turn stops calling the model. Defaults to `maxIterations(50)`. */ agentLoopStrategy?: AgentLoopStrategy modelOptions?: TAdapter['~types']['providerOptions'] interrupts?: ReadonlyArray> diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts index cc4fa6a623..9e5b6b473d 100644 --- a/packages/ai-harness/src/index.ts +++ b/packages/ai-harness/src/index.ts @@ -107,6 +107,7 @@ export { isExpired, loopbackLogin, refreshCredential, + startLoopbackReceiver, } from './oauth' export type { OAuthConfig } from './oauth' diff --git a/packages/ai-harness/src/oauth.ts b/packages/ai-harness/src/oauth.ts index 8928322a76..b73cf14f8e 100644 --- a/packages/ai-harness/src/oauth.ts +++ b/packages/ai-harness/src/oauth.ts @@ -158,6 +158,68 @@ export async function refreshCredential( const DONE_PAGE = 'Signed in

You are signed in. You can close this tab.

' +/** + * A one-time receiver for an OAuth redirect on `127.0.0.1` (RFC 8252). Use it + * when another library runs the OAuth flow and you only need the code back: + * register `redirectUri`, send the user to the authorization URL, then await + * `waitForCode(state)`. It answers one callback, then stops listening. + */ +export async function startLoopbackReceiver( + options: { timeoutMs?: number } = {}, +): Promise<{ + redirectUri: string + waitForCode: (state: string) => Promise + close: () => void +}> { + const { createServer } = await import('node:http') + const server = createServer() + await new Promise((resolve, reject) => { + server.once('error', reject) + server.listen(0, '127.0.0.1', () => resolve()) + }) + const address = server.address() + const port = typeof address === 'object' && address ? address.port : 0 + const redirectUri = `http://127.0.0.1:${port}/callback` + const close = () => { + server.closeAllConnections() + server.close() + } + return { + redirectUri, + close, + waitForCode: (state) => + new Promise((resolve, reject) => { + const timer = setTimeout( + () => { + close() + reject(new Error('Sign-in timed out.')) + }, + options.timeoutMs ?? 10 * 60_000, + ) + server.on('request', (req, res) => { + const url = new URL(req.url ?? '/', redirectUri) + if (url.pathname !== '/callback') { + res.writeHead(404).end() + return + } + clearTimeout(timer) + res.writeHead(200, { 'Content-Type': 'text/html' }).end(DONE_PAGE) + close() + const code = url.searchParams.get('code') + if (url.searchParams.get('state') !== state) { + reject(new Error('Sign-in failed: the state does not match.')) + } else if (!code) { + reject( + new Error( + `Sign-in failed: ${url.searchParams.get('error') ?? 'no code'}`, + ), + ) + } else resolve(code) + }) + }), + } +} + /** * Sign in through the browser with a loopback redirect (RFC 8252 + PKCE). * Listens on `127.0.0.1` on a random port, for one callback only. Calls diff --git a/packages/ai-harness/src/plugins.ts b/packages/ai-harness/src/plugins.ts index d708655a19..21cce825ce 100644 --- a/packages/ai-harness/src/plugins.ts +++ b/packages/ai-harness/src/plugins.ts @@ -62,6 +62,12 @@ export interface PluginContributions { * keep the harness adapter. The last plugin that returns one wins. */ adapter?: () => AnyTextAdapter | undefined + /** + * Tools found at run time, for example the tools of an MCP server the user + * signed in to after the session opened. Called before each chat turn. A + * name that another tool already uses is skipped. + */ + discoverTools?: () => ReadonlyArray | Promise> } /** Plugin state that survives restarts, stored in the metadata store. */ @@ -238,6 +244,10 @@ export interface MountedPlugins { /** Contributions to extension points, by point name. */ extensions: Map> adapters: Array<() => AnyTextAdapter | undefined> + discoverers: Array<{ + discover: () => ReadonlyArray | Promise> + owner: string + }> /** Who contributed what, for `session.inspect()`. */ owners: { plugins: Array<{ @@ -420,6 +430,7 @@ export async function mountPlugins( extensions.set(point, [...items]) } const adapters: Array<() => AnyTextAdapter | undefined> = [] + const discoverers: MountedPlugins['discoverers'] = [] const services = env.services ?? NO_SERVICES try { @@ -550,6 +561,12 @@ export async function mountPlugins( items.push({ value: item.value, owner: plugin.name }) } if (contributions.adapter) adapters.push(contributions.adapter) + if (contributions.discoverTools) { + discoverers.push({ + discover: contributions.discoverTools, + owner: plugin.name, + }) + } for (const agent of contributions.agents ?? []) { env.registry.add(agent, plugin.name) } @@ -602,6 +619,7 @@ export async function mountPlugins( config, extensions, adapters, + discoverers, owners: { plugins: plugins.map((plugin) => ({ name: plugin.name, diff --git a/packages/ai-harness/src/session.ts b/packages/ai-harness/src/session.ts index df0e852c08..af056aa2b5 100644 --- a/packages/ai-harness/src/session.ts +++ b/packages/ai-harness/src/session.ts @@ -6,6 +6,7 @@ import { compactForModel, convertSchemaToJsonSchema, createSubagentId, + maxIterations, runAgentStream, validateWithStandardSchema, } from '@tanstack/ai' @@ -24,6 +25,7 @@ import { import { HARNESS_EVENTS } from './types' import type { AnyChatMiddleware, + AnyTool, Interrupt, ModelMessage, RunAgentResumeItem, @@ -906,6 +908,36 @@ export class HarnessSession { }) } + /** + * Ask each plugin for the tools it found since the last turn. A plugin + * that fails (for example an MCP server that is down) is skipped, and + * clients get a `harness.plugin.warning` event. + */ + private async discoverTools( + discoverers: MountedPlugins['discoverers'], + taken: Set, + operation: OperationImpl, + ): Promise> { + const found: Array = [] + for (const { discover, owner } of discoverers) { + try { + for (const tool of await discover()) { + if (taken.has(tool.name)) continue + taken.add(tool.name) + found.push(tool) + } + } catch (error) { + operation.publish( + customEvent('harness.plugin.warning', { + plugin: owner, + message: error instanceof Error ? error.message : String(error), + }), + ) + } + } + return found + } + /** Middleware that adds queued steer messages before each model call. */ private steering(): AnyChatMiddleware { return { @@ -1017,6 +1049,16 @@ export class HarnessSession { .at(-1) const resolvePrompt = (prompt: string | (() => string)) => typeof prompt === 'function' ? prompt() : prompt + const staticTools = [ + ...(this.harness.tools ?? []), + ...(session?.tools ?? []), + ...(runPlugins?.tools ?? []), + ] + const discovered = await this.discoverTools( + [...(session?.discoverers ?? []), ...(runPlugins?.discoverers ?? [])], + new Set(staticTools.map((tool) => tool.name)), + operation, + ) const stream = chat({ adapter: picked ?? this.harness.adapter, messages: @@ -1029,11 +1071,7 @@ export class HarnessSession { .map(resolvePrompt) .filter((prompt) => prompt !== ''), ], - tools: [ - ...(this.harness.tools ?? []), - ...(session?.tools ?? []), - ...(runPlugins?.tools ?? []), - ], + tools: [...staticTools, ...discovered], middleware: [ ...bridges, withPersistence(this.persistence), @@ -1052,9 +1090,9 @@ export class HarnessSession { }, } : {}), - ...(this.harness.agentLoopStrategy - ? { agentLoopStrategy: this.harness.agentLoopStrategy } - : {}), + // chat() stops after 5 model calls by default. An agent that reads, + // searches, and calls tools needs more before it can answer. + agentLoopStrategy: this.harness.agentLoopStrategy ?? maxIterations(50), ...(this.harness.modelOptions !== undefined ? { modelOptions: this.harness.modelOptions } : {}), diff --git a/packages/ai-harness/tests/discover-tools.test.ts b/packages/ai-harness/tests/discover-tools.test.ts new file mode 100644 index 0000000000..1e7653645d --- /dev/null +++ b/packages/ai-harness/tests/discover-tools.test.ts @@ -0,0 +1,165 @@ +import { describe, expect, it } from 'vitest' +import { z } from 'zod' +import { EventType, toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + createHarnessHost, + defineHarness, + definePlugin, + startLoopbackReceiver, +} from '../src' +import { mockAdapter, text, toolCall } from './helpers' +import type { SessionEvent } from '../src' + +const echoTool = (name: string, reply: string) => + toolDefinition({ + name, + description: `The ${name} tool`, + inputSchema: z.object({}), + }).server(async () => reply) + +describe('discoverTools', () => { + it('adds tools a plugin finds at run time, skips taken names, and warns on failure', async () => { + let online = false + const connector = definePlugin({ + name: 'test/connector', + setup: () => ({ + discoverTools: () => + online + ? [ + echoTool('remote_search', 'remote result'), + echoTool('lookup', 'shadowed'), + ] + : [], + }), + }) + const broken = definePlugin({ + name: 'test/broken', + setup: () => ({ + discoverTools: async () => { + throw new Error('server is down') + }, + }), + }) + const { adapter, calls } = mockAdapter([ + () => text('nothing yet'), + () => toolCall('remote_search', {}), + () => text('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/discover', + adapter, + tools: [echoTool('lookup', 'local result')], + plugins: () => [connector, broken], + }), + { threadId: 't' }, + ) + const names = (call: number) => + (calls[call].tools ?? []).map((tool: { name: string }) => tool.name) + + const seen: Array = [] + const controller = new AbortController() + const reading = (async () => { + for await (const entry of session.events({ signal: controller.signal })) + seen.push(entry) + })() + + await session.prompt('before sign-in') + expect(names(0)).toEqual(['lookup']) + + online = true + const turn = await session.prompt('after sign-in') + expect(turn.text).toBe('done') + // The static `lookup` wins over the discovered one with the same name. + expect(names(1)).toEqual(['lookup', 'remote_search']) + expect(JSON.stringify(calls[2].messages)).toContain('remote result') + + controller.abort() + await reading + const warnings = seen.filter( + (entry) => + entry.event.type === EventType.CUSTOM && + entry.event.name === 'harness.plugin.warning', + ) + expect(warnings[0]?.event).toMatchObject({ + value: { plugin: 'test/broken', message: 'server is down' }, + }) + await host.close() + }) +}) + +describe('the default agent loop', () => { + it('keeps calling the model past five steps', async () => { + const steps = Array.from( + { length: 7 }, + (_, index) => () => toolCall('lookup', {}, `call-${index}`), + ) + const { adapter, calls } = mockAdapter([...steps, () => text('finally')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/long-turn', + adapter, + tools: [echoTool('lookup', 'fact')], + }), + { threadId: 't' }, + ) + const turn = await session.prompt('dig deep') + expect(calls).toHaveLength(8) + expect(turn.text).toBe('finally') + await host.close() + }) +}) + +describe('startLoopbackReceiver', () => { + it('returns the code for the matching state and ignores other paths', async () => { + const receiver = await startLoopbackReceiver() + expect(receiver.redirectUri).toMatch( + /^http:\/\/127\.0\.0\.1:\d+\/callback$/, + ) + const code = receiver.waitForCode('s1') + const other = await fetch( + receiver.redirectUri.replace('/callback', '/favicon.ico'), + ) + expect(other.status).toBe(404) + const done = await fetch(`${receiver.redirectUri}?code=abc&state=s1`) + expect(done.status).toBe(200) + expect(await done.text()).toContain('signed in') + expect(await code).toBe('abc') + }) + + it('refuses a callback with the wrong state', async () => { + const receiver = await startLoopbackReceiver() + const code = receiver.waitForCode('expected') + const rejected = expect(code).rejects.toThrow('the state does not match') + await fetch(`${receiver.redirectUri}?code=abc&state=other`) + await rejected + }) + + it('reports the error the server sent back', async () => { + const receiver = await startLoopbackReceiver() + const code = receiver.waitForCode('s') + const rejected = expect(code).rejects.toThrow( + 'Sign-in failed: access_denied', + ) + await fetch(`${receiver.redirectUri}?error=access_denied&state=s`) + await rejected + + const bare = await startLoopbackReceiver() + const noCode = bare.waitForCode('s') + const missing = expect(noCode).rejects.toThrow('Sign-in failed: no code') + await fetch(`${bare.redirectUri}?state=s`) + await missing + }) + + it('times out and stops listening', async () => { + const receiver = await startLoopbackReceiver({ timeoutMs: 20 }) + await expect(receiver.waitForCode('s')).rejects.toThrow( + 'Sign-in timed out.', + ) + await expect(fetch(receiver.redirectUri)).rejects.toThrow() + receiver.close() + }) +}) diff --git a/packages/ai-mcp/package.json b/packages/ai-mcp/package.json index 7989380ff3..bb338abb69 100644 --- a/packages/ai-mcp/package.json +++ b/packages/ai-mcp/package.json @@ -44,6 +44,10 @@ "./apps": { "types": "./dist/esm/apps/index.d.ts", "import": "./dist/esm/apps/index.js" + }, + "./connector": { + "types": "./dist/esm/connector.d.ts", + "import": "./dist/esm/connector.js" } }, "files": [ @@ -66,7 +70,17 @@ "@modelcontextprotocol/sdk": "^1.29.0", "@tanstack/ai": "workspace:^" }, + "peerDependencies": { + "@tanstack/ai-harness": "workspace:^" + }, + "peerDependenciesMeta": { + "@tanstack/ai-harness": { + "optional": true + } + }, "devDependencies": { + "@tanstack/ai-harness": "workspace:*", + "@tanstack/ai-persistence": "workspace:*", "@vitest/coverage-v8": "4.1.10", "jiti": "^2.4.2", "json-schema-to-typescript": "^15.0.4", diff --git a/packages/ai-mcp/src/connector.ts b/packages/ai-mcp/src/connector.ts new file mode 100644 index 0000000000..3a9e94ab94 --- /dev/null +++ b/packages/ai-mcp/src/connector.ts @@ -0,0 +1,291 @@ +import { randomBytes } from 'node:crypto' +import { auth } from '@modelcontextprotocol/sdk/client/auth.js' +import { + AuthRequiredError, + defineCommand, + definePlugin, + startLoopbackReceiver, +} from '@tanstack/ai-harness' +import { createMCPClient } from './client' +import type { + OAuthClientInformationMixed, + OAuthClientMetadata, + OAuthTokens, +} from '@modelcontextprotocol/sdk/shared/auth.js' +import type { OAuthClientProvider } from '@modelcontextprotocol/sdk/client/auth.js' +import type { AnyTool } from '@tanstack/ai' +import type { CredentialsAccess } from '@tanstack/ai-harness' +import type { MCPClient } from './client' + +export interface McpConnectorOptions { + /** A short id, for example `'notion'`. Commands are `connect:` and `disconnect:`. */ + id: string + label: string + /** The MCP server URL (Streamable HTTP), for example `https://mcp.notion.com/mcp`. */ + url: string + /** Tool name prefix. Default: the id. Tools are named `_`. */ + prefix?: string + /** OAuth scopes to ask for. Default: what the server offers. */ + scopes?: ReadonlyArray + /** The client name shown on the consent screen. Default `'TanStack AI Harness'`. */ + clientName?: string + /** + * Ask for approval before tools that can change data. Default: every tool + * the server does not mark `readOnlyHint`. + */ + needsApproval?: (tool: { + name: string + annotations?: { readOnlyHint?: boolean } + }) => boolean + /** Test hook for the OAuth and MCP requests. */ + fetch?: typeof fetch +} + +/** + * An `OAuthClientProvider` for the MCP SDK that keeps tokens (and the client + * that dynamic registration made) in the harness credential store. + */ +function credentialProvider( + id: string, + credentials: CredentialsAccess, + options: { + redirectUri: string | undefined + clientName: string + scopes: ReadonlyArray | undefined + state?: string + onRedirect?: (url: URL) => void + }, +): OAuthClientProvider { + let client: OAuthClientInformationMixed | undefined + let verifier = '' + const stored = async () => { + const credential = await credentials.get(id) + return credential?.type === 'oauth' ? credential : undefined + } + return { + get redirectUrl() { + return options.redirectUri + }, + get clientMetadata(): OAuthClientMetadata { + return { + client_name: options.clientName, + redirect_uris: options.redirectUri ? [options.redirectUri] : [], + grant_types: ['authorization_code', 'refresh_token'], + response_types: ['code'], + token_endpoint_auth_method: 'none', + ...(options.scopes?.length ? { scope: options.scopes.join(' ') } : {}), + } + }, + ...(options.state ? { state: () => options.state ?? '' } : {}), + clientInformation: async () => { + if (client) return client + // A new sign-in registers a client for its own loopback port. + if (options.onRedirect) return undefined + const saved = (await stored())?.client + return saved + ? { + client_id: saved.clientId, + ...(saved.clientSecret + ? { client_secret: saved.clientSecret } + : {}), + } + : undefined + }, + saveClientInformation: (information) => { + client = information + }, + tokens: async (): Promise => { + if (options.onRedirect) return undefined + const credential = await stored() + if (!credential) return undefined + return { + access_token: credential.accessToken, + token_type: 'Bearer', + ...(credential.refreshToken + ? { refresh_token: credential.refreshToken } + : {}), + ...(credential.expiresAt + ? { + expires_in: Math.max( + 0, + Math.floor((credential.expiresAt - Date.now()) / 1000), + ), + } + : {}), + } + }, + saveTokens: async (tokens) => { + const previous = await stored() + const information = + client ?? + (previous?.client + ? { + client_id: previous.client.clientId, + client_secret: previous.client.clientSecret, + } + : undefined) + await credentials.set(id, { + type: 'oauth', + accessToken: tokens.access_token, + ...(tokens.refresh_token + ? { refreshToken: tokens.refresh_token } + : previous?.refreshToken + ? { refreshToken: previous.refreshToken } + : {}), + ...(tokens.expires_in + ? { expiresAt: Date.now() + tokens.expires_in * 1000 } + : {}), + ...(tokens.scope ? { scopes: tokens.scope.split(' ') } : {}), + ...(information + ? { + client: { + clientId: information.client_id, + ...(information.client_secret + ? { clientSecret: information.client_secret } + : {}), + ...((options.redirectUri ?? previous?.client?.redirectUri) + ? { + redirectUri: + options.redirectUri ?? previous?.client?.redirectUri, + } + : {}), + }, + } + : {}), + }) + }, + redirectToAuthorization: (url) => { + // Outside `/connect`, a missing or revoked sign-in stops the tool. + if (!options.onRedirect) throw new AuthRequiredError(id) + options.onRedirect(url) + }, + saveCodeVerifier: (value) => { + verifier = value + }, + codeVerifier: () => verifier, + } +} + +/** + * A plugin that signs the user in to a remote MCP server (for example Notion + * or Linear) and gives the model its tools, the way Claude Code connects to + * MCP servers: + * + * - `/connect ` finds the server's OAuth settings, registers a client, + * and opens the browser (PKCE, loopback on `127.0.0.1`). + * - The tokens stay in the credential store. The model never sees them. + * - After sign-in, the next turn has the server's tools, named `_`. + * Tools that can change data ask for approval. + * + * @example + * ```ts + * const notion = mcpConnector({ id: 'notion', label: 'Notion', url: 'https://mcp.notion.com/mcp' }) + * ``` + */ +export function mcpConnector(options: McpConnectorOptions) { + const { id, label, url } = options + const prefix = options.prefix ?? id + const clientName = options.clientName ?? 'TanStack AI Harness' + return definePlugin({ + name: `connector/${id}`, + setup: async (ctx) => { + let connected = (await ctx.credentials.get(id)) !== null + let client: MCPClient | undefined + let tools: ReadonlyArray | undefined + const reset = async () => { + tools = undefined + const open = client + client = undefined + await open?.close().catch(() => {}) + } + await ctx.resources.acquire( + () => undefined, + () => reset(), + ) + + return { + prompts: [ + { + id: `connector/${id}:status`, + text: () => + connected + ? '' + : `${label} is not connected. If the user asks for ${label}, tell them to run /connect ${id}.`, + }, + ], + discoverTools: async () => { + if (!connected) return [] + if (!tools) { + client = await createMCPClient({ + transport: { + type: 'http', + url, + authProvider: credentialProvider(id, ctx.credentials, { + redirectUri: undefined, + clientName, + scopes: options.scopes, + }), + ...(options.fetch ? { fetch: options.fetch } : {}), + }, + prefix, + needsApproval: + options.needsApproval ?? + ((tool) => tool.annotations?.readOnlyHint !== true), + }) + tools = (await client.tools()) as ReadonlyArray + } + return tools + }, + commands: { + [`connect:${id}`]: defineCommand({ + description: `Sign in to ${label}`, + run: async () => { + const receiver = await startLoopbackReceiver() + const state = randomBytes(16).toString('base64url') + try { + const provider = credentialProvider(id, ctx.credentials, { + redirectUri: receiver.redirectUri, + clientName, + scopes: options.scopes, + state, + onRedirect: (authorizationUrl) => + ctx.session.authRequired({ + connector: id, + url: authorizationUrl.href, + }), + }) + const fetchFn = options.fetch ? { fetchFn: options.fetch } : {} + const first = await auth(provider, { + serverUrl: url, + ...fetchFn, + }) + if (first === 'REDIRECT') { + const code = await receiver.waitForCode(state) + await auth(provider, { + serverUrl: url, + authorizationCode: code, + ...fetchFn, + }) + } + } finally { + receiver.close() + } + connected = true + await reset() + return `Connected to ${label}.` + }, + }), + [`disconnect:${id}`]: defineCommand({ + description: `Sign out of ${label}`, + run: async () => { + await ctx.credentials.delete(id) + connected = false + await reset() + return `Disconnected from ${label}.` + }, + }), + }, + } + }, + }) +} diff --git a/packages/ai-mcp/tests/connector.test.ts b/packages/ai-mcp/tests/connector.test.ts new file mode 100644 index 0000000000..55261ca78b --- /dev/null +++ b/packages/ai-mcp/tests/connector.test.ts @@ -0,0 +1,361 @@ +import { createServer } from 'node:http' +import { afterEach, describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' +import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js' +import { EventType } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + HARNESS_EVENTS, + createHarnessHost, + defineHarness, +} from '@tanstack/ai-harness' +import { mcpConnector } from '../src/connector' +import type { IncomingMessage, Server } from 'node:http' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' + +/** + * A local MCP server behind OAuth, like Notion's or Linear's: protected + * resource metadata, authorization server metadata, dynamic client + * registration, a token endpoint, and an `echo` tool. + */ +async function startProtectedServer() { + const seen = { registrations: 0, tokenRequests: [] as Array } + let base = '' + const readBody = async (req: IncomingMessage) => { + const chunks: Array = [] + for await (const chunk of req) chunks.push(chunk as Buffer) + return Buffer.concat(chunks).toString('utf8') + } + const server: Server = createServer((req, res) => { + void (async () => { + const url = new URL(req.url ?? '/', base) + const json = (status: number, value: unknown) => { + res.writeHead(status, { 'Content-Type': 'application/json' }) + res.end(JSON.stringify(value)) + } + if (url.pathname.startsWith('/.well-known/oauth-protected-resource')) { + return json(200, { + resource: `${base}/mcp`, + authorization_servers: [base], + }) + } + if (url.pathname.startsWith('/.well-known/oauth-authorization-server')) { + return json(200, { + issuer: base, + authorization_endpoint: `${base}/authorize`, + token_endpoint: `${base}/token`, + registration_endpoint: `${base}/register`, + response_types_supported: ['code'], + grant_types_supported: ['authorization_code', 'refresh_token'], + code_challenge_methods_supported: ['S256'], + token_endpoint_auth_methods_supported: ['none'], + }) + } + if (url.pathname === '/register' && req.method === 'POST') { + seen.registrations += 1 + const metadata = JSON.parse(await readBody(req)) + return json(201, { ...metadata, client_id: 'client-1' }) + } + if (url.pathname === '/token' && req.method === 'POST') { + const form = new URLSearchParams(await readBody(req)) + seen.tokenRequests.push(form) + if ( + form.get('grant_type') === 'authorization_code' && + form.get('code') === 'code-1' + ) { + return json(200, { + access_token: 'access-1', + refresh_token: 'refresh-1', + token_type: 'Bearer', + expires_in: 3600, + }) + } + return json(400, { error: 'invalid_grant' }) + } + if (url.pathname === '/mcp') { + if (req.headers.authorization !== 'Bearer access-1') { + res.writeHead(401, { + 'WWW-Authenticate': `Bearer resource_metadata="${base}/.well-known/oauth-protected-resource"`, + }) + res.end() + return + } + if (req.method !== 'POST') { + res.writeHead(405).end() + return + } + const mcp = new McpServer({ name: 'protected', version: '1.0.0' }) + mcp.registerTool( + 'echo', + { + description: 'Echo text', + inputSchema: { text: z.string() }, + annotations: { readOnlyHint: true }, + }, + async ({ text }) => ({ + content: [{ type: 'text' as const, text: `echo: ${text}` }], + }), + ) + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: undefined, + enableJsonResponse: true, + }) + await mcp.connect(transport) + await transport.handleRequest(req, res, JSON.parse(await readBody(req))) + return + } + res.writeHead(404).end() + })() + }) + await new Promise((resolve) => + server.listen(0, '127.0.0.1', () => resolve()), + ) + const address = server.address() + base = `http://127.0.0.1:${typeof address === 'object' && address ? address.port : 0}` + return { + url: `${base}/mcp`, + seen, + close: () => + new Promise((resolve) => { + server.closeAllConnections() + server.close(() => resolve()) + }), + } +} + +/** A model that calls `tool` once, then answers with the tool result it saw. */ +function modelCalling( + tool: string, + args: Record, +): AnyTextAdapter { + let call = 0 + const now = () => Date.now() + return { + kind: 'text', + name: 'mock', + model: 'mock', + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + chatStream: (options) => + (async function* (): AsyncGenerator { + call += 1 + yield { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: now(), + } + if (call === 1) { + yield { + type: EventType.TOOL_CALL_START, + toolCallId: 'c1', + toolCallName: tool, + timestamp: now(), + } + yield { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'c1', + delta: JSON.stringify(args), + timestamp: now(), + } + yield { + type: EventType.TOOL_CALL_END, + toolCallId: 'c1', + timestamp: now(), + } + yield { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'tool_calls' } }, + } + return + } + const last = options.messages.at(-1) + const text = `saw: ${typeof last?.content === 'string' ? last.content : JSON.stringify(last?.content)}` + yield { + type: EventType.TEXT_MESSAGE_START, + messageId: `m${call}`, + role: 'assistant', + timestamp: now(), + } + yield { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: `m${call}`, + delta: text, + timestamp: now(), + } + yield { + type: EventType.TEXT_MESSAGE_END, + messageId: `m${call}`, + timestamp: now(), + } + yield { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'stop' } }, + } + })(), + } +} + +const cleanups: Array<() => unknown> = [] +afterEach(async () => { + for (const cleanup of cleanups.splice(0).reverse()) await cleanup() +}) + +describe('mcpConnector', () => { + it('signs in with OAuth discovery, registration, and PKCE, then gives the model the server tools', async () => { + const protectedServer = await startProtectedServer() + cleanups.push(() => protectedServer.close()) + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + cleanups.push(() => host.close()) + const harness = defineHarness({ + name: 'test/mcp-connector', + adapter: modelCalling('demo_echo', { text: 'hello' }), + plugins: () => [ + mcpConnector({ id: 'demo', label: 'Demo', url: protectedServer.url }), + ], + }) + const session = await host.open(harness, { + threadId: 't', + principal: { id: 'user-1' }, + }) + + // Act as the browser: approve the sign-in by calling the loopback redirect. + const reader = new AbortController() + let authorizationUrl: URL | undefined + void (async () => { + for await (const entry of session.events({ signal: reader.signal })) { + const event = entry.event + if ( + event.type === EventType.CUSTOM && + event.name === HARNESS_EVENTS.authRequired + ) { + authorizationUrl = new URL( + String((event.value as { url: string }).url), + ) + const redirect = new URL( + authorizationUrl.searchParams.get('redirect_uri') ?? '', + ) + redirect.searchParams.set('code', 'code-1') + redirect.searchParams.set( + 'state', + authorizationUrl.searchParams.get('state') ?? '', + ) + await fetch(redirect) + } + } + })() + cleanups.push(() => reader.abort()) + + await expect(session.command('connect:demo')).resolves.toBe( + 'Connected to Demo.', + ) + expect(authorizationUrl?.searchParams.get('client_id')).toBe('client-1') + expect(authorizationUrl?.searchParams.get('code_challenge_method')).toBe( + 'S256', + ) + expect( + new URL(authorizationUrl?.searchParams.get('redirect_uri') ?? '') + .hostname, + ).toBe('127.0.0.1') + expect(protectedServer.seen.registrations).toBe(1) + expect( + protectedServer.seen.tokenRequests[0]?.get('code_verifier'), + ).toBeTruthy() + + const saved = await persistence.stores.credentials.get( + { threadId: 't', userId: 'user-1' }, + 'demo', + ) + expect(saved).toMatchObject({ + type: 'oauth', + accessToken: 'access-1', + refreshToken: 'refresh-1', + client: { clientId: 'client-1' }, + }) + + const turn = await session.prompt('echo hello') + expect(turn.text).toContain('echo: hello') + // The model never sees the token. + expect( + JSON.stringify(await persistence.stores.messages.loadThread('t')), + ).not.toContain('access-1') + + await session.command('disconnect:demo') + expect( + await persistence.stores.credentials.get( + { threadId: 't', userId: 'user-1' }, + 'demo', + ), + ).toBeNull() + }) + + it('tells the model to ask for /connect before sign-in', async () => { + const seen = vi.fn() + const adapter: AnyTextAdapter = { + ...modelCalling('x', {}), + chatStream: (options) => { + seen(options.systemPrompts) + return (async function* (): AsyncGenerator { + yield { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: Date.now(), + } + yield { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: Date.now(), + metadata: { tanstack: { finishReason: 'stop' } }, + } + })() + }, + } + const host = createHarnessHost({ persistence: memoryPersistence() }) + cleanups.push(() => host.close()) + const session = await host.open( + defineHarness({ + name: 'test/mcp-connector-off', + adapter, + plugins: () => [ + mcpConnector({ + id: 'demo', + label: 'Demo', + url: 'http://127.0.0.1:9/mcp', + }), + ], + }), + { threadId: 't' }, + ) + await session.prompt('hi') + expect(JSON.stringify(seen.mock.calls[0]?.[0])).toContain( + 'run /connect demo', + ) + expect(session.commands().map((command) => command.name)).toEqual([ + 'connect:demo', + 'disconnect:demo', + ]) + }) +}) diff --git a/packages/ai-mcp/vite.config.ts b/packages/ai-mcp/vite.config.ts index 5a8fee615b..2c65c4ca38 100644 --- a/packages/ai-mcp/vite.config.ts +++ b/packages/ai-mcp/vite.config.ts @@ -25,7 +25,12 @@ const config = defineConfig({ export default mergeConfig( config, tanstackViteConfig({ - entry: ['./src/index.ts', './src/stdio.ts', './src/apps/index.ts'], + entry: [ + './src/index.ts', + './src/stdio.ts', + './src/apps/index.ts', + './src/connector.ts', + ], srcDir: './src', cjs: false, }), diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index 3785eb4240..956aef08de 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -425,6 +425,11 @@ export type Credential = refreshToken?: string expiresAt?: number scopes?: Array + /** + * The OAuth client these tokens belong to (for example one made by + * dynamic client registration). A refresh needs it. + */ + client?: { clientId: string; clientSecret?: string; redirectUri?: string } } /** diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index d8f2711c60..ea2077b4b1 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -188,12 +188,18 @@ importers: '@tanstack/ai-dashboard': specifier: workspace:* version: link:../../packages/ai-dashboard + '@tanstack/ai-grok': + specifier: workspace:* + version: link:../../packages/ai-grok '@tanstack/ai-harness': specifier: workspace:* version: link:../../packages/ai-harness '@tanstack/ai-harness-cli': specifier: workspace:* version: link:../../packages/ai-harness-cli + '@tanstack/ai-mcp': + specifier: workspace:* + version: link:../../packages/ai-mcp '@tanstack/ai-openai': specifier: workspace:* version: link:../../packages/ai-openai @@ -2753,6 +2759,12 @@ importers: specifier: workspace:^ version: link:../ai devDependencies: + '@tanstack/ai-harness': + specifier: workspace:* + version: link:../ai-harness + '@tanstack/ai-persistence': + specifier: workspace:* + version: link:../ai-persistence '@vitest/coverage-v8': specifier: 4.1.10 version: 4.1.10(vitest@4.1.11) From f479affcd432765345fd897b1b244bce02f2c6a4 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 15:19:27 +0200 Subject: [PATCH 019/286] feat(ai-code-mode, ai-harness): code mode plugin for harnesses with pluggable isolates --- .changeset/harness-p8-code-mode.md | 8 + docs/config.json | 5 + docs/harness/code-mode.md | 96 ++++++ docs/harness/dashboard.md | 2 +- docs/harness/deploy.md | 2 +- docs/harness/subagents.md | 2 +- examples/harness-cli/README.md | 1 + examples/harness-cli/package.json | 2 + examples/harness-cli/src/harness.ts | 7 + packages/ai-code-mode/package.json | 12 + packages/ai-code-mode/src/harness.ts | 104 ++++++ packages/ai-code-mode/tests/harness.test.ts | 303 ++++++++++++++++++ packages/ai-code-mode/vite.config.ts | 2 +- packages/ai-harness/src/index.ts | 8 + packages/ai-harness/src/plugins.ts | 21 ++ packages/ai-harness/src/session.ts | 48 ++- .../ai-harness/tests/prepare-tools.test.ts | 80 +++++ pnpm-lock.yaml | 12 + 18 files changed, 704 insertions(+), 11 deletions(-) create mode 100644 .changeset/harness-p8-code-mode.md create mode 100644 docs/harness/code-mode.md create mode 100644 packages/ai-code-mode/src/harness.ts create mode 100644 packages/ai-code-mode/tests/harness.test.ts create mode 100644 packages/ai-harness/tests/prepare-tools.test.ts diff --git a/.changeset/harness-p8-code-mode.md b/.changeset/harness-p8-code-mode.md new file mode 100644 index 0000000000..7c5d58f318 --- /dev/null +++ b/.changeset/harness-p8-code-mode.md @@ -0,0 +1,8 @@ +--- +'@tanstack/ai-code-mode': minor +'@tanstack/ai-harness': minor +--- + +`@tanstack/ai-code-mode/harness` adds `codeMode({ driver })`: a harness plugin that gives the model one `execute_typescript` tool. The model writes a program that calls several tools, and the program runs in the isolate of any `@tanstack/ai-isolate-*` driver. Read-only tools move into code mode, including MCP tools found after sign-in. Tools that need approval, edits, and commands stay normal tool calls. Tool names that are not JavaScript identifiers (for example `notion-search`) get a safe name inside the program. + +`@tanstack/ai-harness` adds `prepareTools` to plugins, to change the tool list of each turn. `PermissionRules` and `decidePermission` are now also exported from the package root. diff --git a/docs/config.json b/docs/config.json index 8b9dea312f..71a5f53d3a 100644 --- a/docs/config.json +++ b/docs/config.json @@ -844,6 +844,11 @@ "to": "harness/mcp", "addedAt": "2026-09-26" }, + { + "label": "Code mode in a harness", + "to": "harness/code-mode", + "addedAt": "2026-09-26" + }, { "label": "Run agents from a harness", "to": "harness/subagents", diff --git a/docs/harness/code-mode.md b/docs/harness/code-mode.md new file mode 100644 index 0000000000..c84a9e51fb --- /dev/null +++ b/docs/harness/code-mode.md @@ -0,0 +1,96 @@ +--- +title: Code mode in a harness +id: harness-code-mode +order: 9 +description: "Let the harness model write one TypeScript program that calls many tools, and run it in an isolate. Any TanStack AI isolate driver plugs in." +keywords: + - tanstack ai + - harness + - code mode + - isolate + - quickjs +--- + +Your agent has 100 tools from Notion and Linear, and it calls them one at a time. Each call is a model round trip, and every tool schema goes into every request. With code mode, the model writes one TypeScript program that calls the tools it needs, and the program runs in an isolate. The read-only tools leave the tool list and become functions in that program. + +## 1. Add the plugin + + + +react: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs +vue: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs +solid: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs +svelte: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs +preact: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs +angular: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs +octane: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs +vanilla: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs + + + +```ts group=harness-code-mode +import { defineHarness } from '@tanstack/ai-harness' +import { codeMode } from '@tanstack/ai-code-mode/harness' +import { createQuickJSIsolateDriver } from '@tanstack/ai-isolate-quickjs' +import { mcpConnector } from '@tanstack/ai-mcp/connector' +import { openaiText } from '@tanstack/ai-openai' + +const linear = mcpConnector({ + id: 'linear', + label: 'Linear', + url: 'https://mcp.linear.app/mcp', +}) + +export const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-5.6'), + plugins: () => [ + linear, + codeMode({ driver: createQuickJSIsolateDriver() }), + ], +}) +``` + +The model now has an `execute_typescript` tool. Inside the program, each moved tool is an `external_*` function, for example `external_linear_list_issues()`. + +## 2. Pick the isolate + +`driver` takes any isolate driver. Change one line to run the programs somewhere else: + +| Driver | Package | Runs in | +| --- | --- | --- | +| `createQuickJSIsolateDriver()` | `@tanstack/ai-isolate-quickjs` | WebAssembly, on any runtime | +| `createNodeIsolateDriver()` | `@tanstack/ai-isolate-node` | A V8 isolate in Node | +| `createCloudflareIsolateDriver()` | `@tanstack/ai-isolate-cloudflare` | A Cloudflare Worker | +| `createDaytonaIsolateDriver()` | `@tanstack/ai-isolate-daytona` | A Daytona sandbox | + +[Code mode isolates](../code-mode/code-mode-isolates) has the options of each driver. + +## 3. Choose which tools move + +A call inside the program does not stop for approval. So by default, a tool moves into code mode only when it is safe to run without a question: + +- The tool runs on the server and does not set `needsApproval`. +- The permission rules allow it in plan mode. File reads move. `write_file` and `run_command` stay. +- MCP tools that the server does not mark read-only ask for approval, so they stay too. + +Every other tool stays a normal tool call, with its approvals. To pick the tools yourself, pass `include`: + +```ts group=harness-code-mode +export const onlyLinear = codeMode({ + driver: createQuickJSIsolateDriver(), + include: (tool) => tool.name.startsWith('linear_'), +}) +``` + +The other options of `createCodeMode` pass through: `timeout`, `memoryLimit`, `lazyToolsConfig`, and more. + +## Tools that appear after sign-in + +The plugin chooses the tools again before every turn. When the user runs `/connect linear`, the read-only Linear tools move into code mode on the next turn. A plugin can do the same with its own tool changes: return the new list from `prepareTools`. [Use MCP servers](./mcp) shows `discoverTools`, which adds the tools first. + +## What you have now + +- One `execute_typescript` tool in place of many read-only tools. +- Programs that run in the isolate you choose. +- Approvals kept for every tool that can change something. diff --git a/docs/harness/dashboard.md b/docs/harness/dashboard.md index 3fab5e3e0a..afbd511cd7 100644 --- a/docs/harness/dashboard.md +++ b/docs/harness/dashboard.md @@ -1,7 +1,7 @@ --- title: Self-host the dashboard id: harness-dashboard -order: 11 +order: 12 description: "Watch and steer your harness sessions from a browser or a phone. Agents dial out to your dashboard server, so they need no open port." keywords: - tanstack ai diff --git a/docs/harness/deploy.md b/docs/harness/deploy.md index e75187f2a4..ea6cf77423 100644 --- a/docs/harness/deploy.md +++ b/docs/harness/deploy.md @@ -1,7 +1,7 @@ --- title: Deploy a harness id: harness-deploy -order: 10 +order: 11 description: "Run a harness in your server, as a worker process, on another machine, or as a single executable." keywords: - tanstack ai diff --git a/docs/harness/subagents.md b/docs/harness/subagents.md index 169770cb0d..78de061c97 100644 --- a/docs/harness/subagents.md +++ b/docs/harness/subagents.md @@ -1,7 +1,7 @@ --- title: Run agents from a harness id: harness-subagents -order: 9 +order: 10 description: "Start typed agents from commands and plugins, run them in groups, call a whole harness as a child, and keep the tree within limits." keywords: - tanstack ai diff --git a/examples/harness-cli/README.md b/examples/harness-cli/README.md index bd0afebc98..b87e8aa82f 100644 --- a/examples/harness-cli/README.md +++ b/examples/harness-cli/README.md @@ -27,6 +27,7 @@ Try these: - Sign-ins are kept in `~/.tanstack-harness-example/credentials.json`, so you sign in once. `/disconnect notion` deletes one. - Images use `OPENAI_API_KEY`. Videos use Grok Imagine when `XAI_API_KEY` is set, and OpenAI Sora when it is not. +- Code mode is on: read-only tools (file reads, read-only Notion and Linear tools) are `external_*` functions in one `execute_typescript` program, which runs in a QuickJS isolate. Ask: `in one program, list my Linear issues and search Notion for them`. ## Other modes diff --git a/examples/harness-cli/package.json b/examples/harness-cli/package.json index 96c40cb7d6..a9d609b2b8 100644 --- a/examples/harness-cli/package.json +++ b/examples/harness-cli/package.json @@ -12,10 +12,12 @@ "@tanstack/ai": "workspace:*", "@tanstack/ai-acp": "workspace:*", "@tanstack/ai-anthropic": "workspace:*", + "@tanstack/ai-code-mode": "workspace:*", "@tanstack/ai-dashboard": "workspace:*", "@tanstack/ai-grok": "workspace:*", "@tanstack/ai-harness": "workspace:*", "@tanstack/ai-harness-cli": "workspace:*", + "@tanstack/ai-isolate-quickjs": "workspace:*", "@tanstack/ai-mcp": "workspace:*", "@tanstack/ai-openai": "workspace:*", "@tanstack/ai-persistence": "workspace:*", diff --git a/examples/harness-cli/src/harness.ts b/examples/harness-cli/src/harness.ts index 86699642c0..bc0563abda 100644 --- a/examples/harness-cli/src/harness.ts +++ b/examples/harness-cli/src/harness.ts @@ -11,8 +11,10 @@ import { workspaceTools, } from '@tanstack/ai-harness/plugins' import { anthropicText } from '@tanstack/ai-anthropic' +import { codeMode } from '@tanstack/ai-code-mode/harness' import { mcpConnector } from '@tanstack/ai-mcp/connector' import { grokVideo } from '@tanstack/ai-grok' +import { createQuickJSIsolateDriver } from '@tanstack/ai-isolate-quickjs' import { openaiText, openaiVideo } from '@tanstack/ai-openai' import { z } from 'zod' import { imageAgent, videoAgent } from './media' @@ -167,5 +169,10 @@ export const assistant = defineHarness({ projectInstructions({ root }), compact({ adapter: main }), usage(), + // Read-only tools (file reads, read-only Notion and Linear tools) move + // behind execute_typescript, so the model can call several in one program. + // The program runs in a QuickJS isolate. Any @tanstack/ai-isolate-* driver + // works here. + codeMode({ driver: createQuickJSIsolateDriver() }), ], }) diff --git a/packages/ai-code-mode/package.json b/packages/ai-code-mode/package.json index b3370c1f6c..ce05e3b109 100644 --- a/packages/ai-code-mode/package.json +++ b/packages/ai-code-mode/package.json @@ -24,6 +24,10 @@ ".": { "types": "./dist/esm/index.d.ts", "import": "./dist/esm/index.js" + }, + "./harness": { + "types": "./dist/esm/harness.d.ts", + "import": "./dist/esm/harness.js" } }, "sideEffects": false, @@ -65,10 +69,18 @@ }, "peerDependencies": { "@tanstack/ai": "workspace:^", + "@tanstack/ai-harness": "workspace:^", "zod": "^3.0.0 || ^4.0.0" }, + "peerDependenciesMeta": { + "@tanstack/ai-harness": { + "optional": true + } + }, "devDependencies": { "@tanstack/ai": "workspace:*", + "@tanstack/ai-harness": "workspace:*", + "@tanstack/ai-persistence": "workspace:*", "@vitest/coverage-v8": "4.1.10", "zod": "^4.2.0" } diff --git a/packages/ai-code-mode/src/harness.ts b/packages/ai-code-mode/src/harness.ts new file mode 100644 index 0000000000..392959b55c --- /dev/null +++ b/packages/ai-code-mode/src/harness.ts @@ -0,0 +1,104 @@ +import { + PermissionRules, + decidePermission, + definePlugin, +} from '@tanstack/ai-harness' +import { createCodeMode } from './create-code-mode' +import type { AnyTool } from '@tanstack/ai' +import type { CodeModeTool, CodeModeToolConfig } from './types' + +export interface CodeModePluginOptions extends Omit< + CodeModeToolConfig, + 'tools' +> { + /** + * Which tools move into code mode. Default: every server tool that does + * not need approval and that the permission rules allow in plan mode (so + * no edits, no commands, nothing that asks first). + */ + include?: (tool: CodeModeTool) => boolean +} + +/** + * A name that works as a JavaScript identifier. Code mode turns each tool + * into an `external_` function, and MCP tool names often have `-`. + */ +function identifierOf(name: string): string { + const safe = name.replace(/[^A-Za-z0-9_$]/g, '_') + return /^[0-9]/.test(safe) ? `_${safe}` : safe +} + +/** A tool that runs on the server: not a client tool, and it has `execute`. */ +function isServerTool(tool: AnyTool): tool is CodeModeTool { + const side = '__toolSide' in tool ? tool.__toolSide : 'server' + return side === 'server' && typeof tool.execute === 'function' +} + +/** + * A harness plugin that gives the model an `execute_typescript` tool. The + * model writes one TypeScript program that calls several tools, and the + * program runs in the isolate of `driver` (any `@tanstack/ai-isolate-*` + * driver). The tools it calls leave the model's tool list, including tools + * that plugins find at run time, such as MCP tools after `/connect`. + * + * Calls inside the isolate do not stop for approval, so by default only + * tools that are safe to run without a question move into code mode. The + * other tools stay normal tool calls. + * + * @example + * ```ts + * import { codeMode } from '@tanstack/ai-code-mode/harness' + * import { createQuickJSIsolateDriver } from '@tanstack/ai-isolate-quickjs' + * + * defineHarness({ + * name: 'acme/agent', + * adapter, + * plugins: () => [codeMode({ driver: createQuickJSIsolateDriver() })], + * }) + * ``` + */ +export function codeMode(options: CodeModePluginOptions) { + const { include, ...config } = options + return definePlugin({ + name: 'tanstack/code-mode', + setup: (ctx) => { + const rules = ctx.collect(PermissionRules) + const safe = (tool: CodeModeTool) => + !tool.needsApproval && + decidePermission(rules, tool.name, 'plan') === 'allow' + // The system prompt for the tools of the current turn. + let prompt = '' + return { + prompts: [{ id: 'tanstack/code-mode', text: () => prompt }], + prepareTools: (tools) => { + // Two tools that map to the same identifier: the first one moves. + const byIdentifier = new Map() + for (const tool of tools.filter(isServerTool)) { + if (!(include ? include(tool) : safe(tool))) continue + const identifier = identifierOf(tool.name) + if (!byIdentifier.has(identifier)) + byIdentifier.set(identifier, tool) + } + if (byIdentifier.size === 0) { + prompt = '' + return tools + } + const created = createCodeMode({ + ...config, + tools: [...byIdentifier].map(([identifier, tool]) => + identifier === tool.name ? tool : { ...tool, name: identifier }, + ), + }) + prompt = created.systemPrompt + const moved = new Set( + [...byIdentifier.values()].map((tool) => tool.name), + ) + return [ + ...tools.filter((tool) => !moved.has(tool.name)), + ...created.tools, + ] + }, + } + }, + }) +} diff --git a/packages/ai-code-mode/tests/harness.test.ts b/packages/ai-code-mode/tests/harness.test.ts new file mode 100644 index 0000000000..8266bbf402 --- /dev/null +++ b/packages/ai-code-mode/tests/harness.test.ts @@ -0,0 +1,303 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { EventType, toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + PermissionRules, + createHarnessHost, + defineHarness, + definePlugin, +} from '@tanstack/ai-harness' +import { codeMode } from '../src/harness' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import type { ExecutionResult, IsolateDriver, ToolBinding } from '../src/types' + +const now = () => Date.now() + +/** A model whose turns are scripted. Records the options of every call. */ +function scripted(turns: Array>) { + const calls: Array = [] + const adapter: AnyTextAdapter = { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + chatStream: (options: any) => { + calls.push(options) + const chunks = turns[calls.length - 1] ?? [] + return (async function* () { + yield* chunks + })() + }, + } + return { adapter, calls } +} + +const textTurn = (content: string): Array => [ + { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, + { + type: EventType.TEXT_MESSAGE_START, + messageId: 'm', + role: 'assistant', + timestamp: now(), + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm', + delta: content, + timestamp: now(), + }, + { type: EventType.TEXT_MESSAGE_END, messageId: 'm', timestamp: now() }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'stop' } }, + }, +] + +const codeTurn = (code: string): Array => [ + { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, + { + type: EventType.TOOL_CALL_START, + toolCallId: 'call-code', + toolCallName: 'execute_typescript', + timestamp: now(), + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call-code', + delta: JSON.stringify({ typescriptCode: code }), + timestamp: now(), + }, + { type: EventType.TOOL_CALL_END, toolCallId: 'call-code', timestamp: now() }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'tool_calls' } }, + }, +] + +/** + * A stand-in isolate: it records the bindings it got and, for any program, + * calls `external_lookup` and `external_notion_search` and returns both + * results. Real drivers run the program itself. + */ +function fakeDriver() { + const seen: Array> = [] + const driver: IsolateDriver = { + createContext: async ({ bindings }) => { + seen.push(Object.keys(bindings).sort()) + const call = (name: string, args: unknown) => + (bindings[name] as ToolBinding).execute(args) + return { + execute: async (): Promise> => { + const value = { + lookup: await call('external_lookup', { q: 'x' }), + notion: await call('external_notion_search', { query: 'y' }), + } + // The caller picks T. This stand-in always returns the two results. + return { success: true, value: value as T, logs: [] } + }, + dispose: async () => {}, + } + }, + } + return { driver, seen } +} + +const lookup = vi.fn(async () => 'fact') +const notionSearch = vi.fn(async () => ({ results: ['page'] })) +const tools = { + lookup: toolDefinition({ + name: 'lookup', + description: 'Look up a fact', + inputSchema: z.object({ q: z.string() }), + }).server(lookup), + remove: toolDefinition({ + name: 'remove', + description: 'Remove a file', + needsApproval: true, + inputSchema: z.object({ path: z.string() }), + }).server(async () => 'removed'), + writeFile: toolDefinition({ + name: 'write_file', + description: 'Write a file', + inputSchema: z.object({ path: z.string() }), + }).server(async () => 'written'), + notionSearch: toolDefinition({ + name: 'notion_search', + description: 'Search Notion', + inputSchema: z.object({ query: z.string() }), + }).server(notionSearch), + askUser: toolDefinition({ + name: 'ask_user', + description: 'Ask the user in the browser', + inputSchema: z.object({}), + }).client(), +} + +// A tool plugin that marks write_file as an edit, like workspaceTools does. +const writes = definePlugin({ + name: 'test/writes', + setup: () => ({ + tools: [tools.writeFile], + contribute: [ + PermissionRules.item({ + tool: 'write_file', + decision: 'ask', + kind: 'edit', + }), + ], + }), +}) +// A connector whose tool appears at run time. +const notion = definePlugin({ + name: 'test/notion', + setup: () => ({ discoverTools: () => [tools.notionSearch] }), +}) + +describe('codeMode', () => { + it('moves safe tools, including discovered ones, behind execute_typescript', async () => { + const { driver, seen } = fakeDriver() + const { adapter, calls } = scripted([ + codeTurn('return { a: await external_lookup({ q: "x" }) }'), + textTurn('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/code-mode', + adapter, + tools: [tools.lookup, tools.remove, tools.askUser], + plugins: () => [writes, notion, codeMode({ driver })], + }), + { threadId: 't' }, + ) + const turn = await session.prompt('look it up') + expect(turn.text).toBe('done') + + const names = calls[0].tools.map((tool: { name: string }) => tool.name) + expect(names).toEqual([ + 'remove', + 'ask_user', + 'write_file', + 'execute_typescript', + ]) + const prompts = JSON.stringify(calls[0].systemPrompts) + expect(prompts).toContain('external_lookup') + expect(prompts).toContain('external_notion_search') + expect(prompts).not.toContain('external_write_file') + + expect(seen[0]).toEqual(['external_lookup', 'external_notion_search']) + expect(lookup).toHaveBeenCalledTimes(1) + expect(notionSearch).toHaveBeenCalledTimes(1) + expect(JSON.stringify(calls[1].messages)).toContain('page') + await host.close() + }) + + it('uses include to pick the tools, and leaves the list alone when none match', async () => { + const { driver } = fakeDriver() + const { adapter, calls } = scripted([textTurn('one'), textTurn('two')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const picked = await host.open( + defineHarness({ + name: 'test/code-mode-include', + adapter, + tools: [tools.lookup, tools.remove], + plugins: () => [ + codeMode({ driver, include: (tool) => tool.name === 'remove' }), + ], + }), + { threadId: 'picked' }, + ) + await picked.prompt('hi') + expect(calls[0].tools.map((tool: { name: string }) => tool.name)).toEqual([ + 'lookup', + 'execute_typescript', + ]) + + const none = await host.open( + defineHarness({ + name: 'test/code-mode-none', + adapter, + tools: [tools.remove], + plugins: () => [codeMode({ driver })], + }), + { threadId: 'none' }, + ) + await none.prompt('hi') + expect(calls[1].tools.map((tool: { name: string }) => tool.name)).toEqual([ + 'remove', + ]) + expect(JSON.stringify(calls[1].systemPrompts ?? [])).not.toContain( + 'execute_typescript', + ) + await host.close() + }) + + it('gives tools with names that are not identifiers a safe name', async () => { + const seen: Array> = [] + const driver: IsolateDriver = { + createContext: async ({ bindings }) => { + seen.push(Object.keys(bindings).sort()) + return { + execute: async (): Promise> => { + const binding = bindings['external_web_search'] as ToolBinding + const value: unknown = await binding.execute({}) + // The caller picks T. This stand-in returns the one result. + return { success: true, value: value as T, logs: [] } + }, + dispose: async () => {}, + } + }, + } + const named = (name: string, reply: string) => + toolDefinition({ + name, + description: 'A tool', + inputSchema: z.object({}), + }).server(async () => reply) + const { adapter, calls } = scripted([codeTurn('return 1'), textTurn('ok')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/code-mode-names', + adapter, + tools: [ + named('web-search', 'from web-search'), + named('web_search', 'from web_search'), + named('2fa', 'code'), + ], + plugins: () => [codeMode({ driver })], + }), + { threadId: 't' }, + ) + await session.prompt('search') + // web_search clashes with web-search after the rename, so it stays direct. + expect(calls[0].tools.map((tool: { name: string }) => tool.name)).toEqual([ + 'web_search', + 'execute_typescript', + ]) + expect(seen[0]).toEqual(['external__2fa', 'external_web_search']) + expect(JSON.stringify(calls[1].messages)).toContain('from web-search') + await host.close() + }) +}) diff --git a/packages/ai-code-mode/vite.config.ts b/packages/ai-code-mode/vite.config.ts index 77bcc2e60b..525f4daf73 100644 --- a/packages/ai-code-mode/vite.config.ts +++ b/packages/ai-code-mode/vite.config.ts @@ -29,7 +29,7 @@ const config = defineConfig({ export default mergeConfig( config, tanstackViteConfig({ - entry: ['./src/index.ts'], + entry: ['./src/index.ts', './src/harness.ts'], srcDir: './src', cjs: false, }), diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts index 9e5b6b473d..01e56f91f9 100644 --- a/packages/ai-harness/src/index.ts +++ b/packages/ai-harness/src/index.ts @@ -21,6 +21,14 @@ export type { } from './plugins' export { createExtensionPoint, createPluginEvent } from './extensions' +// Portable (no Node imports), so edge-safe plugins such as code mode can read +// the permission rules without the Node-only `./plugins` entry. +export { PermissionRules, decidePermission } from './first-party/permissions' +export type { + PermissionDecision, + PermissionMode, + PermissionRule, +} from './first-party/permissions' export type { ExtensionItem, ExtensionPoint, PluginEvent } from './extensions' export { checkConfigValue, configOption } from './config' diff --git a/packages/ai-harness/src/plugins.ts b/packages/ai-harness/src/plugins.ts index 21cce825ce..92b39e4367 100644 --- a/packages/ai-harness/src/plugins.ts +++ b/packages/ai-harness/src/plugins.ts @@ -68,6 +68,15 @@ export interface PluginContributions { * name that another tool already uses is skipped. */ discoverTools?: () => ReadonlyArray | Promise> + /** + * Change the tool list of a turn: every tool the model gets, after + * `discoverTools`. Return the new list. Runs before prompts resolve, so a + * prompt `text()` can describe the tools this returned. Code mode uses it to + * move tools behind `execute_typescript`. + */ + prepareTools?: ( + tools: ReadonlyArray, + ) => ReadonlyArray | Promise> } /** Plugin state that survives restarts, stored in the metadata store. */ @@ -248,6 +257,10 @@ export interface MountedPlugins { discover: () => ReadonlyArray | Promise> owner: string }> + preparers: Array<{ + prepare: NonNullable + owner: string + }> /** Who contributed what, for `session.inspect()`. */ owners: { plugins: Array<{ @@ -431,6 +444,7 @@ export async function mountPlugins( } const adapters: Array<() => AnyTextAdapter | undefined> = [] const discoverers: MountedPlugins['discoverers'] = [] + const preparers: MountedPlugins['preparers'] = [] const services = env.services ?? NO_SERVICES try { @@ -567,6 +581,12 @@ export async function mountPlugins( owner: plugin.name, }) } + if (contributions.prepareTools) { + preparers.push({ + prepare: contributions.prepareTools, + owner: plugin.name, + }) + } for (const agent of contributions.agents ?? []) { env.registry.add(agent, plugin.name) } @@ -620,6 +640,7 @@ export async function mountPlugins( extensions, adapters, discoverers, + preparers, owners: { plugins: plugins.map((plugin) => ({ name: plugin.name, diff --git a/packages/ai-harness/src/session.ts b/packages/ai-harness/src/session.ts index af056aa2b5..2c2352679b 100644 --- a/packages/ai-harness/src/session.ts +++ b/packages/ai-harness/src/session.ts @@ -927,17 +927,45 @@ export class HarnessSession { found.push(tool) } } catch (error) { - operation.publish( - customEvent('harness.plugin.warning', { - plugin: owner, - message: error instanceof Error ? error.message : String(error), - }), - ) + this.warn(operation, owner, error) } } return found } + /** + * Let each plugin change the tool list of this turn, in plugin order. A + * plugin that fails leaves the list as it was, with a warning event. + */ + private async prepareTools( + preparers: MountedPlugins['preparers'], + tools: Array, + operation: OperationImpl, + ): Promise> { + let prepared = tools + for (const { prepare, owner } of preparers) { + try { + prepared = [...(await prepare(prepared))] + } catch (error) { + this.warn(operation, owner, error) + } + } + return prepared + } + + private warn( + operation: OperationImpl, + plugin: string, + error: unknown, + ) { + operation.publish( + customEvent('harness.plugin.warning', { + plugin, + message: error instanceof Error ? error.message : String(error), + }), + ) + } + /** Middleware that adds queued steer messages before each model call. */ private steering(): AnyChatMiddleware { return { @@ -1059,6 +1087,12 @@ export class HarnessSession { new Set(staticTools.map((tool) => tool.name)), operation, ) + // Before the chat() options below: prompts may describe these tools. + const tools = await this.prepareTools( + [...(session?.preparers ?? []), ...(runPlugins?.preparers ?? [])], + [...staticTools, ...discovered], + operation, + ) const stream = chat({ adapter: picked ?? this.harness.adapter, messages: @@ -1071,7 +1105,7 @@ export class HarnessSession { .map(resolvePrompt) .filter((prompt) => prompt !== ''), ], - tools: [...staticTools, ...discovered], + tools, middleware: [ ...bridges, withPersistence(this.persistence), diff --git a/packages/ai-harness/tests/prepare-tools.test.ts b/packages/ai-harness/tests/prepare-tools.test.ts new file mode 100644 index 0000000000..d12fe94c02 --- /dev/null +++ b/packages/ai-harness/tests/prepare-tools.test.ts @@ -0,0 +1,80 @@ +import { describe, expect, it } from 'vitest' +import { z } from 'zod' +import { EventType, toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { mockAdapter, text } from './helpers' +import type { SessionEvent } from '../src' + +const tool = (name: string) => + toolDefinition({ + name, + description: `The ${name} tool`, + inputSchema: z.object({}), + }).server(async () => name) + +describe('prepareTools', () => { + it('lets plugins rewrite the tool list in order, before prompts resolve', async () => { + let described = '' + const bundle = definePlugin({ + name: 'test/bundle', + setup: () => ({ + prompts: [{ id: 'test/bundle', text: () => described }], + prepareTools: (tools) => { + described = `Bundled: ${tools.map((entry) => entry.name).join(', ')}` + return [tool('bundled')] + }, + }), + }) + const extra = definePlugin({ + name: 'test/extra', + setup: () => ({ + prepareTools: async (tools) => [...tools, tool('extra')], + }), + }) + const broken = definePlugin({ + name: 'test/broken', + setup: () => ({ + prepareTools: () => { + throw new Error('cannot prepare') + }, + }), + }) + const { adapter, calls } = mockAdapter([() => text('ok')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/prepare', + adapter, + tools: [tool('a'), tool('b')], + plugins: () => [bundle, broken, extra], + }), + { threadId: 't' }, + ) + const seen: Array = [] + const controller = new AbortController() + const reading = (async () => { + for await (const entry of session.events({ signal: controller.signal })) + seen.push(entry) + })() + + await session.prompt('go') + expect(calls[0].tools.map((entry: { name: string }) => entry.name)).toEqual( + ['bundled', 'extra'], + ) + expect(JSON.stringify(calls[0].systemPrompts)).toContain('Bundled: a, b') + + controller.abort() + await reading + expect( + seen.find( + (entry) => + entry.event.type === EventType.CUSTOM && + entry.event.name === 'harness.plugin.warning', + )?.event, + ).toMatchObject({ + value: { plugin: 'test/broken', message: 'cannot prepare' }, + }) + await host.close() + }) +}) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index ea2077b4b1..4bec7ae9ab 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -185,6 +185,9 @@ importers: '@tanstack/ai-anthropic': specifier: workspace:* version: link:../../packages/ai-anthropic + '@tanstack/ai-code-mode': + specifier: workspace:* + version: link:../../packages/ai-code-mode '@tanstack/ai-dashboard': specifier: workspace:* version: link:../../packages/ai-dashboard @@ -197,6 +200,9 @@ importers: '@tanstack/ai-harness-cli': specifier: workspace:* version: link:../../packages/ai-harness-cli + '@tanstack/ai-isolate-quickjs': + specifier: workspace:* + version: link:../../packages/ai-isolate-quickjs '@tanstack/ai-mcp': specifier: workspace:* version: link:../../packages/ai-mcp @@ -2254,6 +2260,12 @@ importers: '@tanstack/ai': specifier: workspace:* version: link:../ai + '@tanstack/ai-harness': + specifier: workspace:* + version: link:../ai-harness + '@tanstack/ai-persistence': + specifier: workspace:* + version: link:../ai-persistence '@vitest/coverage-v8': specifier: 4.1.10 version: 4.1.10(vitest@4.1.11) From a725cf5004341a8cbfcc868b4d685dc904ce05fa Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 15:26:40 +0200 Subject: [PATCH 020/286] fix(ai-harness): deliver session events published while a reader is busy --- packages/ai-harness/src/feed.ts | 3 ++ packages/ai-harness/tests/feed.test.ts | 51 ++++++++++++++++++++++++++ 2 files changed, 54 insertions(+) create mode 100644 packages/ai-harness/tests/feed.test.ts diff --git a/packages/ai-harness/src/feed.ts b/packages/ai-harness/src/feed.ts index 0d9e031cfb..e5f8b63081 100644 --- a/packages/ai-harness/src/feed.ts +++ b/packages/ai-harness/src/feed.ts @@ -72,6 +72,9 @@ export class SessionFeed { for (const entry of rest) yield entry return } + // Events published while the reader handled `next` woke nobody, so + // scan again before waiting. An empty scan moves `after` to the head. + if (next.length > 0) continue await this.waitForNext(options.signal) } } diff --git a/packages/ai-harness/tests/feed.test.ts b/packages/ai-harness/tests/feed.test.ts new file mode 100644 index 0000000000..c8ffa35682 --- /dev/null +++ b/packages/ai-harness/tests/feed.test.ts @@ -0,0 +1,51 @@ +import { describe, expect, it } from 'vitest' +import { EventType } from '@tanstack/ai' +import { SessionFeed } from '../src/feed' + +const custom = (name: string) => ({ + type: EventType.CUSTOM, + name, + value: {}, + timestamp: Date.now(), +}) + +describe('SessionFeed', () => { + it('delivers events published while a reader handles earlier ones', async () => { + const feed = new SessionFeed() + const controller = new AbortController() + const seen: Array = [] + const reading = (async () => { + for await (const entry of feed.read({ signal: controller.signal })) { + seen.push(entry.event.type === EventType.CUSTOM ? entry.event.name : '') + // Two more events arrive while this one is handled, and nothing after. + if (seen.length === 1) { + await Promise.resolve() + feed.publish('op', custom('second')) + feed.publish('op', custom('third')) + } + } + })() + feed.publish('op', custom('first')) + await new Promise((resolve) => setTimeout(resolve, 20)) + expect(seen).toEqual(['first', 'second', 'third']) + controller.abort() + await reading + }) + + it('skips filtered events and replays from a cursor', async () => { + const feed = new SessionFeed() + feed.publish('a', custom('one')) + feed.publish('b', custom('two')) + feed.publish('a', custom('three')) + feed.close() + const names: Array = [] + for await (const entry of feed.read({ + from: '1', + filter: (item) => item.operationId === 'a', + })) { + names.push(entry.event.type === EventType.CUSTOM ? entry.event.name : '') + } + expect(names).toEqual(['three']) + expect(feed.head()).toBe('3') + }) +}) From f9a2cc6bd757d90613b1bfd53dbd04920dc5d3a0 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 15:28:38 +0200 Subject: [PATCH 021/286] test(ai-harness): type the feed test events --- packages/ai-harness/tests/feed.test.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/ai-harness/tests/feed.test.ts b/packages/ai-harness/tests/feed.test.ts index c8ffa35682..ec2032bcab 100644 --- a/packages/ai-harness/tests/feed.test.ts +++ b/packages/ai-harness/tests/feed.test.ts @@ -1,8 +1,9 @@ import { describe, expect, it } from 'vitest' import { EventType } from '@tanstack/ai' import { SessionFeed } from '../src/feed' +import type { StreamChunk } from '@tanstack/ai' -const custom = (name: string) => ({ +const custom = (name: string): StreamChunk => ({ type: EventType.CUSTOM, name, value: {}, From 23158adb660283d22e411d7f9ecd9b8a78798545 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 15:29:11 +0200 Subject: [PATCH 022/286] fix(ai-harness): compact reads text content parts; cover OAuth, workspace tools, and plugin session APIs --- .../src/first-party/session-tools.ts | 4 +- packages/ai-harness/tests/oauth-edges.test.ts | 334 ++++++++++ .../ai-harness/tests/session-plugins.test.ts | 595 ++++++++++++++++++ .../ai-harness/tests/workspace-tools.test.ts | 141 +++++ 4 files changed, 1071 insertions(+), 3 deletions(-) create mode 100644 packages/ai-harness/tests/oauth-edges.test.ts create mode 100644 packages/ai-harness/tests/session-plugins.test.ts create mode 100644 packages/ai-harness/tests/workspace-tools.test.ts diff --git a/packages/ai-harness/src/first-party/session-tools.ts b/packages/ai-harness/src/first-party/session-tools.ts index 02dbe6fb6c..219bfe5180 100644 --- a/packages/ai-harness/src/first-party/session-tools.ts +++ b/packages/ai-harness/src/first-party/session-tools.ts @@ -10,9 +10,7 @@ function textOf(message: ModelMessage): string { if (typeof message.content === 'string') return message.content if (Array.isArray(message.content)) { return message.content - .map((part) => - 'text' in part && typeof part.text === 'string' ? part.text : '', - ) + .map((part) => (part.type === 'text' ? part.content : '')) .join('') } return '' diff --git a/packages/ai-harness/tests/oauth-edges.test.ts b/packages/ai-harness/tests/oauth-edges.test.ts new file mode 100644 index 0000000000..825b69f045 --- /dev/null +++ b/packages/ai-harness/tests/oauth-edges.test.ts @@ -0,0 +1,334 @@ +import { describe, expect, it } from 'vitest' +import { toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + AuthRequiredError, + buildAuthorizationUrl, + createHarnessHost, + defineHarness, + deviceLogin, + exchangeCode, + isExpired, + loopbackLogin, + oauthConnector, + refreshCredential, + scrubSecrets, +} from '../src' +import { mockAdapter, text, toolCall } from './helpers' +import type { OAuthConfig } from '../src' + +const oauth: OAuthConfig = { + authorizationUrl: 'https://auth.example/authorize', + tokenUrl: 'https://auth.example/token', + deviceUrl: 'https://auth.example/device', + clientId: 'client-1', +} + +/** A fetch that answers each call with the next body, and records the forms. */ +function replies(bodies: Array, status = 200) { + const forms: Array = [] + const fake: typeof fetch = async (_input, init) => { + forms.push(new URLSearchParams(String(init?.body ?? ''))) + const body = bodies.shift() + return new Response( + typeof body === 'string' ? body : JSON.stringify(body), + { + status, + }, + ) + } + return { fake, forms } +} + +describe('token requests', () => { + it('sends the client secret and keeps the old refresh token', async () => { + const { fake, forms } = replies([ + { access_token: 'a2', expires_in: 60, scope: 'repo,read:user' }, + ]) + const refreshed = await refreshCredential( + { ...oauth, clientSecret: 'shh' }, + { type: 'oauth', accessToken: 'a1', refreshToken: 'r1' }, + fake, + ) + expect(forms[0]?.get('client_secret')).toBe('shh') + expect(refreshed).toMatchObject({ + type: 'oauth', + accessToken: 'a2', + refreshToken: 'r1', + scopes: ['repo', 'read:user'], + }) + expect(refreshed.type === 'oauth' && refreshed.expiresAt).toBeGreaterThan( + Date.now(), + ) + }) + + it('refuses a credential without a refresh token', async () => { + await expect( + refreshCredential(oauth, { type: 'oauth', accessToken: 'a' }), + ).rejects.toThrow('no refresh token') + await expect( + refreshCredential(oauth, { type: 'api_key', value: 'k' }), + ).rejects.toThrow('no refresh token') + }) + + it('reports bodies that are not token objects', async () => { + const options = { code: 'c', verifier: 'v', redirectUri: 'http://x' } + await expect( + exchangeCode(oauth, { ...options, fetch: replies(['"text"'], 500).fake }), + ).rejects.toThrow('OAuth token request failed (500).') + await expect( + exchangeCode(oauth, { ...options, fetch: replies(['not json']).fake }), + ).rejects.toThrow('no access_token') + }) + + it('knows when a credential expires', () => { + expect(isExpired({ type: 'api_key', value: 'k' })).toBe(false) + expect(isExpired({ type: 'oauth', accessToken: 'a' })).toBe(false) + expect( + isExpired({ + type: 'oauth', + accessToken: 'a', + expiresAt: Date.now() + 30_000, + }), + ).toBe(true) + expect( + isExpired( + { type: 'oauth', accessToken: 'a', expiresAt: Date.now() + 30_000 }, + 0, + ), + ).toBe(false) + }) + + it('leaves the scope out of the URL when there is none', () => { + const url = new URL( + buildAuthorizationUrl(oauth, { + redirectUri: 'http://127.0.0.1:1/callback', + state: 's', + challenge: 'c', + }), + ) + expect(url.searchParams.has('scope')).toBe(false) + }) +}) + +describe('loopbackLogin edges', () => { + async function login( + callback: (redirect: URL, state: string) => string, + options: { timeoutMs?: number } = {}, + ) { + const { fake } = replies([{ access_token: 'token' }]) + return loopbackLogin(oauth, { + fetch: fake, + ...options, + onUrl: (value) => { + const url = new URL(value) + const redirect = new URL(url.searchParams.get('redirect_uri') ?? '') + const state = url.searchParams.get('state') ?? '' + void (async () => { + // A request to another path does not end the sign-in. + await fetch(new URL('/favicon.ico', redirect)) + await fetch(callback(redirect, state)) + })() + }, + }) + } + + it('reports the error from the server and a missing code', async () => { + await expect( + login( + (redirect, state) => `${redirect}?state=${state}&error=access_denied`, + ), + ).rejects.toThrow('Sign-in failed: access_denied') + await expect( + login((redirect, state) => `${redirect}?state=${state}`), + ).rejects.toThrow('Sign-in failed: no code') + }) + + it('times out when nobody comes back', async () => { + await expect( + loopbackLogin(oauth, { onUrl: () => {}, timeoutMs: 10 }), + ).rejects.toThrow('Sign-in timed out.') + }) +}) + +describe('deviceLogin edges', () => { + const noSleep = async () => {} + + it('needs a device endpoint and a device code', async () => { + await expect( + deviceLogin({ ...oauth, deviceUrl: undefined }, { onCode: () => {} }), + ).rejects.toThrow('no device endpoint') + await expect( + deviceLogin(oauth, { + onCode: () => {}, + fetch: replies([{ user_code: 'X' }]).fake, + sleep: noSleep, + }), + ).rejects.toThrow('no device code') + }) + + it('slows down when asked, sends scopes, and falls back to the authorization URL', async () => { + const waits: Array = [] + const shown: Array = [] + const { fake, forms } = replies([ + { device_code: 'd', user_code: 'CODE' }, + { error: 'slow_down' }, + { error: 'authorization_pending' }, + { access_token: 'granted' }, + ]) + const credential = await deviceLogin( + { ...oauth, scopes: ['repo'] }, + { + onCode: (info) => + shown.push(`${info.userCode} ${info.verificationUri}`), + fetch: fake, + sleep: async (ms) => { + waits.push(ms) + }, + }, + ) + expect(forms[0]?.get('scope')).toBe('repo') + expect(shown).toEqual(['CODE https://auth.example/authorize']) + expect(waits).toEqual([5000, 10000, 10000]) + expect(credential).toMatchObject({ accessToken: 'granted' }) + }) + + it('stops when the device code expires', async () => { + const { fake } = replies([ + { device_code: 'd', user_code: 'C', expires_in: 0, interval: 0 }, + ]) + await expect( + deviceLogin(oauth, { onCode: () => {}, fetch: fake, sleep: noSleep }), + ).rejects.toThrow('Device sign-in expired.') + }) +}) + +describe('oauthConnector edges', () => { + function connector(fetchFn: typeof fetch, seen: Array) { + return oauthConnector({ + id: 'svc', + label: 'Service', + oauth, + fetch: fetchFn, + tools: (token) => [ + toolDefinition({ name: 'whoami', description: 'Who am I' }).server( + async () => { + seen.push(await token()) + return 'me' + }, + ), + ], + }) + } + + it('refreshes an expired token, and passes an API key through', async () => { + const persistence = memoryPersistence() + await persistence.stores.credentials.set( + { threadId: 't1', userId: 'u1' }, + 'svc', + { + type: 'oauth', + accessToken: 'old', + refreshToken: 'r', + expiresAt: Date.now() - 1, + }, + ) + await persistence.stores.credentials.set( + { threadId: 't2', userId: 'u2' }, + 'svc', + { + type: 'api_key', + value: 'key-1', + }, + ) + const seen: Array = [] + const { fake } = replies([{ access_token: 'fresh', expires_in: 3600 }]) + const { adapter } = mockAdapter([ + () => toolCall('whoami', {}), + () => text('done'), + () => toolCall('whoami', {}), + () => text('done'), + ]) + const host = createHarnessHost({ persistence }) + const harness = defineHarness({ + name: 'test/refresh', + adapter, + plugins: () => [connector(fake, seen)], + }) + await ( + await host.open(harness, { threadId: 't1', principal: { id: 'u1' } }) + ).prompt('who') + await ( + await host.open(harness, { threadId: 't2', principal: { id: 'u2' } }) + ).prompt('who') + expect(seen).toEqual(['fresh', 'key-1']) + expect( + await persistence.stores.credentials.get( + { threadId: 't1', userId: 'u1' }, + 'svc', + ), + ).toMatchObject({ accessToken: 'fresh', refreshToken: 'r' }) + await host.close() + }) + + it('signs in through the loopback and signs out', async () => { + const persistence = memoryPersistence() + const { fake } = replies([{ access_token: 'loop-token' }]) + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ + name: 'test/loopback-connector', + adapter: mockAdapter([]).adapter, + plugins: () => [connector(fake, [])], + }), + { threadId: 't' }, + ) + const controller = new AbortController() + const signIn = (async () => { + for await (const entry of session.events({ signal: controller.signal })) { + const value = + entry.event.type === 'CUSTOM' ? entry.event.value : undefined + if ( + typeof value === 'object' && + value !== null && + 'url' in value && + typeof value.url === 'string' + ) { + const url = new URL(value.url) + const redirect = url.searchParams.get('redirect_uri') ?? '' + const state = url.searchParams.get('state') ?? '' + await fetch(`${redirect}?code=c&state=${state}`) + return + } + } + })() + expect(await session.command('connect:svc')).toBe('Connected to Service.') + controller.abort() + await signIn + expect( + await persistence.stores.credentials.get({ threadId: 't' }, 'svc'), + ).toMatchObject({ accessToken: 'loop-token' }) + expect(await session.command('disconnect:svc')).toBe( + 'Disconnected from Service.', + ) + expect( + await persistence.stores.credentials.get({ threadId: 't' }, 'svc'), + ).toBeNull() + await host.close() + }) +}) + +describe('auth helpers', () => { + it('describes the sign-in in the error, with and without a URL', () => { + expect(new AuthRequiredError('svc').message).toContain('/connect svc') + const withUrl = new AuthRequiredError('svc', 'https://x/login') + expect(withUrl.message).toContain('https://x/login') + expect(withUrl.url).toBe('https://x/login') + }) + + it('keeps short strings when scrubbing', () => { + expect( + scrubSecrets('pin 12345 and token abcdef', ['12345', 'abcdef']), + ).toBe('pin 12345 and token [redacted]') + }) +}) diff --git a/packages/ai-harness/tests/session-plugins.test.ts b/packages/ai-harness/tests/session-plugins.test.ts new file mode 100644 index 0000000000..cafc4d4b2c --- /dev/null +++ b/packages/ai-harness/tests/session-plugins.test.ts @@ -0,0 +1,595 @@ +import { mkdtemp, rm, writeFile, mkdir } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { EventType, toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + checkConfigValue, + configOption, + createHarnessHost, + createPluginEvent, + defineCommand, + defineHarness, + definePlugin, +} from '../src' +import { + compact, + fileCommands, + modelPicker, + projectInstructions, + todos, + usage, +} from '../src/first-party' +import { gate, mockAdapter, text, toolCall } from './helpers' +import type { SessionEvent } from '../src' + +const settings = definePlugin({ + name: 'test/settings', + setup: () => ({ + config: { + level: configOption.select({ options: ['low', 'high'], default: 'low' }), + verbose: configOption.boolean({ default: false }), + label: configOption.text({ default: 'none' }), + retries: configOption.number({ default: 1, min: 0, max: 5 }), + }, + }), +}) + +describe('config values', () => { + it('checks every option type and its limits', () => { + const number = configOption.number({ default: 1, min: 0, max: 5 }) + expect(() => checkConfigValue('n', number, 'x')).toThrow('a number') + expect(() => checkConfigValue('n', number, Number.NaN)).toThrow('a number') + expect(() => checkConfigValue('n', number, -1)).toThrow('>= 0') + expect(() => checkConfigValue('n', number, 6)).toThrow('<= 5') + expect(checkConfigValue('n', configOption.number({ default: 1 }), 99)).toBe( + 99, + ) + expect(() => + checkConfigValue('b', configOption.boolean({ default: false }), 'yes'), + ).toThrow('true or false') + expect(() => + checkConfigValue('t', configOption.text({ default: '' }), 1), + ).toThrow('expects text') + expect(() => configOption.select({ options: ['a'], default: 'b' })).toThrow( + 'is not an option', + ) + }) + + it('rejects unknown keys and bad values, and loads saved values in a new session', async () => { + const persistence = memoryPersistence() + await persistence.stores.metadata.set('harness:config', 't', { + verbose: 'not a boolean', + gone: 'old key', + }) + const { adapter } = mockAdapter([]) + const host = createHarnessHost({ persistence }) + const harness = defineHarness({ + name: 'test/config', + adapter, + plugins: () => [settings], + }) + const session = await host.open(harness, { threadId: 't' }) + // The saved values were invalid or unknown, so the defaults apply. + expect(session.config().verbose?.value).toBe(false) + expect(await session.setConfig('missing', 1)).toMatchObject({ + status: 'rejected', + reason: 'unknown_config', + }) + expect(await session.setConfig('retries', 9)).toMatchObject({ + status: 'rejected', + reason: expect.stringContaining('<= 5'), + }) + expect((await session.setConfig('level', 'high')).status).toBe('accepted') + expect((await session.setConfig('label', 'mine')).status).toBe('accepted') + await host.close() + + const again = createHarnessHost({ persistence }) + const reopened = await again.open(harness, { threadId: 't' }) + expect(reopened.config()).toMatchObject({ + level: { value: 'high', owner: 'test/settings' }, + label: { value: 'mine' }, + retries: { value: 1 }, + }) + await again.close() + }) +}) + +describe('commands', () => { + it('lists commands with input schemas and reports failures', async () => { + const release = gate() + const tools = definePlugin({ + name: 'test/commands', + setup: () => ({ + commands: { + greet: defineCommand({ + description: 'Greet someone', + input: z.object({ name: z.string() }), + run: (input) => `hello ${input.name}`, + }), + boom: defineCommand({ + description: 'Throws', + run: () => { + throw new Error('command broke') + }, + }), + wait: defineCommand({ + description: 'Waits until stopped', + run: (_input, ctx) => + new Promise((_resolve, reject) => { + ctx.signal.addEventListener('abort', () => + reject(new Error('stopped')), + ) + release.open() + }), + }), + }, + }), + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/cmds', + adapter: mockAdapter([]).adapter, + plugins: () => [tools], + }), + { threadId: 't' }, + ) + const listed = session.commands() + expect(listed.find((command) => command.name === 'greet')).toMatchObject({ + owner: 'test/commands', + input: { type: 'object' }, + }) + expect( + listed.find((command) => command.name === 'boom')?.input, + ).toBeUndefined() + + expect(await session.command('greet', { name: 'Ada' })).toBe('hello Ada') + await expect(session.command('greet', { name: 1 })).rejects.toThrow( + 'Input validation failed for command greet', + ) + await expect(session.command('nope')).rejects.toThrow( + 'Unknown command: nope', + ) + await expect(session.command('boom')).rejects.toThrow('command broke') + + const waiting = session.command('wait') + await release.opened + await waiting.cancel() + await waiting.then( + () => undefined, + () => undefined, + ) + expect(waiting.status()).toBe('cancelled') + await host.close() + }) +}) + +describe('questions', () => { + it('asks with a schema, checks answers, and rejects open questions on close', async () => { + const asker = definePlugin({ + name: 'test/asker', + setup: (ctx) => ({ + commands: { + pick: defineCommand({ + description: 'Pick a number', + run: async () => { + const answer = await ctx.session.ask({ + message: 'How many?', + schema: z.object({ count: z.number() }), + }) + return answer.count * 2 + }, + }), + open: defineCommand({ + description: 'Asks and never gets an answer', + run: () => ctx.session.ask({ message: 'Still there?' }), + }), + }, + }), + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/questions', + adapter: mockAdapter([]).adapter, + plugins: () => [asker], + }), + { threadId: 't' }, + ) + const picking = session.command('pick') + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + const [question] = session.snapshot().pendingQuestions + expect(question).toMatchObject({ + message: 'How many?', + schema: { type: 'object' }, + }) + expect(await session.answer('q-unknown', 1)).toMatchObject({ + status: 'rejected', + reason: 'unknown_question', + }) + expect( + await session.answer(question?.questionId ?? '', { count: 'two' }), + ).toMatchObject({ + status: 'rejected', + reason: expect.stringContaining('Invalid answer'), + }) + expect( + (await session.answer(question?.questionId ?? '', { count: 2 })).status, + ).toBe('accepted') + expect(await picking).toBe(4) + + const open = session.command('open') + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + await session.close() + await expect(open).rejects.toThrow('Session closed.') + await host.close() + }) +}) + +describe('plugin state and events', () => { + const counter = createPluginEvent('test/count') + + function statePlugin(onValue: (value: number) => void) { + return definePlugin({ + name: 'test/state', + setup: (ctx) => { + const state = ctx.state({ count: 0 }) + const stop = ctx.on(counter, () => { + throw new Error('a listener that throws is ignored') + }) + ctx.on(counter, onValue) + return { + commands: { + bump: defineCommand({ + description: 'Add one', + run: async () => { + const next = await state.update((value) => ({ + count: value.count + 1, + })) + ctx.emit(counter, next.count) + stop() + return (await state.get()).count + }, + }), + }, + } + }, + }) + } + + it('keeps state with a plain metadata store and with none', async () => { + const base = memoryPersistence() + const plain = { + ...base, + stores: { + ...base.stores, + metadata: { + get: (namespace: string, key: string) => + base.stores.metadata.get(namespace, key), + set: (namespace: string, key: string, value: unknown) => + base.stores.metadata.set(namespace, key, value), + }, + }, + } + const none = { ...base, stores: { ...base.stores, metadata: undefined } } + for (const persistence of [plain, none]) { + const seen: Array = [] + const host = createHarnessHost({ persistence: persistence as never }) + const session = await host.open( + defineHarness({ + name: 'test/state-store', + adapter: mockAdapter([]).adapter, + plugins: () => [statePlugin((value) => seen.push(value))], + }), + { threadId: 't' }, + ) + expect(await session.command('bump')).toBe(1) + expect(await session.command('bump')).toBe(2) + expect(seen).toEqual([1, 2]) + await host.close() + } + }) + + it('gives up after five conflicting updates', async () => { + const base = memoryPersistence() + const conflicting = { + ...base, + stores: { + ...base.stores, + metadata: { + get: async () => null, + set: async () => {}, + getVersioned: async () => ({ value: { count: 0 }, revision: 'r1' }), + setIf: async () => ({ ok: false }), + }, + }, + } + const host = createHarnessHost({ persistence: conflicting as never }) + const session = await host.open( + defineHarness({ + name: 'test/state-conflict', + adapter: mockAdapter([]).adapter, + plugins: () => [statePlugin(() => {})], + }), + { threadId: 't' }, + ) + await expect(session.command('bump')).rejects.toThrow( + 'state update conflicted 5 times', + ) + await host.close() + }) +}) + +describe('the plugin session API and inspect', () => { + it('reads and replaces the transcript, prompts, and describes the session', async () => { + let api: any + const probe = definePlugin({ + name: 'test/probe', + setup: (ctx) => { + api = ctx.session + return { + tools: [ + toolDefinition({ name: 'ping', description: 'Ping' }).server( + async () => 'pong', + ), + ], + prompts: ['Be brief.'], + } + }, + }) + const { adapter, calls } = mockAdapter([ + () => text('first'), + () => text('from the plugin prompt'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/api', + adapter, + plugins: () => [probe, settings], + }), + { threadId: 't', principal: { id: 'u1' } }, + ) + expect(api.threadId).toBe('t') + expect(api.principal).toEqual({ id: 'u1' }) + await session.prompt('hello') + expect( + (await api.transcript()).map((m: { content: string }) => m.content), + ).toEqual(['hello', 'first']) + await api.replaceTranscript([{ role: 'user', content: 'replaced' }]) + expect(await api.transcript()).toEqual([ + { role: 'user', content: 'replaced' }, + ]) + api.prompt('queued by a plugin') + await vi.waitFor(() => expect(calls).toHaveLength(2)) + await vi.waitFor(() => expect(api.snapshot().status).toBe('idle')) + expect(JSON.stringify(calls[1].messages)).toContain('replaced') + + const inspected = session.inspect() + expect(inspected.plugins.map((plugin) => plugin.name)).toEqual([ + 'test/probe', + 'test/settings', + ]) + expect(inspected.tools).toEqual( + expect.arrayContaining([expect.objectContaining({ name: 'ping' })]), + ) + expect(inspected.config.map((entry) => entry.key)).toContain('retries') + await host.close() + }) +}) + +describe('first-party plugin edges', () => { + let dir: string + beforeAll(async () => { + dir = await mkdtemp(join(tmpdir(), 'harness-files-')) + await mkdir(join(dir, 'commands')) + await writeFile( + join(dir, 'commands', 'review.md'), + '---\ndescription: Review a file\n---\nReview $ARGUMENTS carefully.', + ) + await writeFile(join(dir, 'commands', 'plain.md'), 'Say hi to $ARGUMENTS.') + await writeFile(join(dir, 'commands', 'skip.txt'), 'not a command') + await writeFile(join(dir, 'RULES.md'), 'Use tabs.') + await writeFile(join(dir, 'EMPTY.md'), ' ') + }) + afterAll(async () => { + await rm(dir, { recursive: true, force: true }) + }) + + it('turns markdown files into commands and reads chosen instruction files', async () => { + const { adapter, calls } = mockAdapter([ + () => text('a'), + () => text('b'), + () => text('c'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/files', + adapter, + plugins: () => [ + fileCommands({ dir: join(dir, 'commands') }), + projectInstructions({ + root: dir, + files: ['RULES.md', 'EMPTY.md', 'NONE.md'], + }), + ], + }), + { threadId: 't' }, + ) + const described = session.commands() + expect( + described.map((command) => [command.name, command.description]).sort(), + ).toEqual([ + ['plain', 'Run plain.md'], + ['review', 'Review a file'], + ]) + expect(await session.command('review', 'src/a.ts')).toBe('Sent /review.') + await vi.waitFor(() => expect(calls).toHaveLength(1)) + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + await session.command('plain', { who: 'Ada' }) + await vi.waitFor(() => expect(calls).toHaveLength(2)) + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + await session.command('plain') + await vi.waitFor(() => expect(calls).toHaveLength(3)) + + // A missing folder gives no commands. + const empty = await host.open( + defineHarness({ + name: 'test/files-missing', + adapter, + plugins: () => [fileCommands({ dir: join(dir, 'missing') })], + }), + { threadId: 'empty' }, + ) + expect(empty.commands()).toEqual([]) + + const said = calls.map((call) => JSON.stringify(call.messages.at(-1))) + expect(said[0]).toContain('Review src/a.ts carefully.') + expect(said[1]).toContain('Say hi to {\\"who\\":\\"Ada\\"}.') + expect(said[2]).toContain('Say hi to .') + const prompts = JSON.stringify(calls[0].systemPrompts) + expect(prompts).toContain('Project instructions from RULES.md') + expect(prompts).not.toContain('EMPTY.md') + await host.close() + }) + + it('checks todo input and shows an empty list', async () => { + const { adapter, calls } = mockAdapter([ + () => toolCall('todo_write', { todos: [{ text: 'x', status: 'later' }] }), + () => text('bad input'), + () => toolCall('todo_write', { todos: [] }, 'call-2'), + () => text('cleared'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ name: 'test/todos', adapter, plugins: () => [todos()] }), + { threadId: 't' }, + ) + await session.prompt('plan') + expect(JSON.stringify(calls[1].messages)).toContain( + 'Each todo needs text and a status.', + ) + await session.prompt('clear') + expect(JSON.stringify(calls[3].messages)).toContain( + 'The todo list is empty.', + ) + expect(await session.command('todos')).toBe('The todo list is empty.') + await host.close() + }) + + it('needs a model choice and answers every /model form', async () => { + expect(() => modelPicker({ choices: {} })).toThrow('at least one choice') + const fast = mockAdapter([]).adapter + const smart = mockAdapter([]).adapter + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/models', + adapter: smart, + plugins: () => [modelPicker({ choices: { smart, fast } })], + }), + { threadId: 't' }, + ) + expect(await session.command('model')).toBe( + 'Model: smart. Choices: smart, fast.', + ) + expect(await session.command('model', 'slow')).toBe( + 'Unknown model "slow". Choices: smart, fast.', + ) + expect(await session.command('model', { name: 'fast' })).toBe( + 'Model: fast. It applies at the next turn.', + ) + await host.close() + }) + + it('keeps short conversations, keeps the last messages, and counts usage', async () => { + const summarizer = mockAdapter([() => text('the summary')]) + const withUsage = (content: string) => + text(content).map((chunk) => + chunk.type === EventType.RUN_FINISHED + ? { + ...chunk, + usage: { promptTokens: 3, completionTokens: 2, totalTokens: 5 }, + } + : chunk, + ) + const { adapter } = mockAdapter([ + () => withUsage('one'), + () => withUsage('two'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/compact', + adapter, + plugins: () => [ + compact({ adapter: summarizer.adapter, keepLast: 1 }), + usage(), + ], + }), + { threadId: 't' }, + ) + expect(await session.command('compact')).toBe( + 'The conversation is already short.', + ) + await session.prompt([{ type: 'text', content: 'first question' }] as never) + await session.prompt('second question') + expect(await session.command('compact')).toBe( + 'Compacted 3 messages into a summary.', + ) + expect(JSON.stringify(summarizer.calls[0].messages)).toContain( + 'first question', + ) + expect(await session.command('usage')).toBe( + '2 model calls, 6 input tokens, 4 output tokens, 10 total.', + ) + await host.close() + }) +}) + +describe('session events', () => { + it('publishes config changes', async () => { + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/events', + adapter: mockAdapter([]).adapter, + plugins: () => [settings], + }), + { threadId: 't' }, + ) + const seen: Array = [] + const controller = new AbortController() + const reading = (async () => { + for await (const entry of session.events({ + from: '0', + signal: controller.signal, + })) + seen.push(entry) + })() + await session.setConfig('verbose', true) + await vi.waitFor(() => + expect( + seen.some( + (entry) => + entry.event.type === EventType.CUSTOM && + entry.event.value !== null && + typeof entry.event.value === 'object' && + 'key' in entry.event.value && + entry.event.value.key === 'verbose', + ), + ).toBe(true), + ) + controller.abort() + await reading + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/workspace-tools.test.ts b/packages/ai-harness/tests/workspace-tools.test.ts new file mode 100644 index 0000000000..c3e1f71f7f --- /dev/null +++ b/packages/ai-harness/tests/workspace-tools.test.ts @@ -0,0 +1,141 @@ +import { mkdir, mkdtemp, readFile, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { workspaceTools } from '../src/first-party' +import type { AnyTool } from '@tanstack/ai' + +let root: string +let tools: Map + +const run = (name: string, args: unknown): Promise => { + const tool = tools.get(name) + if (!tool?.execute) throw new Error(`No tool ${name}`) + return Promise.resolve(tool.execute(args)) +} + +beforeAll(async () => { + root = await mkdtemp(join(tmpdir(), 'harness-workspace-')) + await mkdir(join(root, 'src', 'deep'), { recursive: true }) + await mkdir(join(root, 'node_modules', 'pkg'), { recursive: true }) + await writeFile(join(root, 'src', 'a.ts'), 'one\ntwo\nthree\nfour\n') + await writeFile(join(root, 'src', 'deep', 'b.ts'), 'const two = 2\n') + await writeFile(join(root, 'notes.md'), 'two words\n') + await writeFile(join(root, 'node_modules', 'pkg', 'index.js'), 'two\n') + await writeFile(join(root, 'big.txt'), `two\n${'x'.repeat(1_100_000)}`) + const plugin = workspaceTools({ root, bashTimeoutMs: 20_000 }) + const contributions = await plugin.setup() + tools = new Map( + (contributions.tools ?? []).map((tool): [string, AnyTool] => [ + tool.name, + tool, + ]), + ) +}) + +afterAll(async () => { + await rm(root, { recursive: true, force: true }) +}) + +describe('read_file', () => { + it('reads a window of numbered lines', async () => { + expect( + await run('read_file', { path: 'src/a.ts', offset: 2, limit: 2 }), + ).toBe('2\ttwo\n3\tthree') + expect( + await run('read_file', { path: 'src/a.ts', offset: -4, limit: 1 }), + ).toBe('1\tone') + }) + + it('refuses paths outside the workspace and missing arguments', async () => { + await expect(run('read_file', { path: '../outside.txt' })).rejects.toThrow( + 'outside the workspace', + ) + await expect( + run('read_file', { path: join(tmpdir(), 'elsewhere.txt') }), + ).rejects.toThrow('outside the workspace') + await expect(run('read_file', {})).rejects.toThrow( + 'Argument "path" must be a string.', + ) + await expect(run('read_file', null)).rejects.toThrow('must be a string') + }) + + it('cuts very long output', async () => { + await writeFile(join(root, 'long.txt'), 'y'.repeat(25_000)) + const text = await run('read_file', { path: 'long.txt' }) + expect(text).toContain('more characters]') + expect(text.length).toBeLessThan(21_000) + }) +}) + +describe('write_file and edit_file', () => { + it('creates folders and files', async () => { + expect( + await run('write_file', { path: 'new/dir/c.txt', content: 'hi' }), + ).toBe('Wrote new/dir/c.txt.') + expect(await readFile(join(root, 'new', 'dir', 'c.txt'), 'utf8')).toBe('hi') + await expect(run('write_file', { path: 'x.txt' })).rejects.toThrow( + 'Argument "content" must be a string.', + ) + }) + + it('edits one match, refuses ambiguous or missing text, and replaces all when asked', async () => { + await writeFile(join(root, 'e.txt'), 'a b a') + await expect( + run('edit_file', { path: 'e.txt', old: 'zzz', new: 'q' }), + ).rejects.toThrow('not in the file') + await expect( + run('edit_file', { path: 'e.txt', old: 'a', new: 'q' }), + ).rejects.toThrow('appears 2 times') + expect( + await run('edit_file', { + path: 'e.txt', + old: 'a', + new: 'q', + replaceAll: true, + }), + ).toBe('Edited e.txt (2 changes).') + expect(await run('edit_file', { path: 'e.txt', old: 'b', new: 'c' })).toBe( + 'Edited e.txt (1 change).', + ) + expect(await readFile(join(root, 'e.txt'), 'utf8')).toBe('q c q') + }) +}) + +describe('list_files and grep', () => { + it('lists files, skips dependency folders, and filters by glob', async () => { + const all = await run('list_files', {}) + expect(all).toContain('src/deep/b.ts') + expect(all).not.toContain('node_modules') + expect(await run('list_files', { pattern: 'src/**/*.ts' })).toBe( + 'src/a.ts\nsrc/deep/b.ts', + ) + expect(await run('list_files', { pattern: '*.none' })).toBe('No files.') + }) + + it('finds lines, skips big files, and filters by glob', async () => { + const hits = await run('grep', { pattern: 'two' }) + expect(hits).toContain('src/a.ts:2: two') + expect(hits).toContain('notes.md:1: two words') + expect(hits).not.toContain('big.txt') + expect(hits).not.toContain('node_modules') + expect(await run('grep', { pattern: 'two', glob: '**/*.md' })).toBe( + 'notes.md:1: two words', + ) + expect(await run('grep', { pattern: 'nothing-like-this' })).toBe( + 'No matches.', + ) + }) +}) + +describe('bash', () => { + it('reports the exit code and stderr', async () => { + const ok = await run('bash', { + command: `node -e "console.log('out'); console.error('warn')"`, + }) + expect(ok).toMatch(/^exit code: 0\nout/) + expect(ok).toContain('stderr:\nwarn') + const failed = await run('bash', { command: `node -e "process.exit(3)"` }) + expect(failed).toMatch(/^exit code: 3/) + }) +}) From 7672d22b180ae2693f78ebce3d64717c73c4435b Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 15:34:45 +0200 Subject: [PATCH 023/286] test(ai-harness): cover plugin agent lookups and services outside a session --- .../ai-harness/tests/plugin-agents.test.ts | 125 ++++++++++++++++++ 1 file changed, 125 insertions(+) create mode 100644 packages/ai-harness/tests/plugin-agents.test.ts diff --git a/packages/ai-harness/tests/plugin-agents.test.ts b/packages/ai-harness/tests/plugin-agents.test.ts new file mode 100644 index 0000000000..2ce1ec81da --- /dev/null +++ b/packages/ai-harness/tests/plugin-agents.test.ts @@ -0,0 +1,125 @@ +import { describe, expect, it } from 'vitest' +import { z } from 'zod' +import { defineAgent } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + createHarnessHost, + createPluginEvent, + defineCommand, + defineHarness, + definePlugin, +} from '../src' +import { AgentRegistry } from '../src/agents' +import { mountPlugins } from '../src/plugins' +import { mockAdapter } from './helpers' + +const painter = defineAgent({ + name: 'painter', + description: 'Paints', + produces: 'image', + inputSchema: z.object({ subject: z.string() }), + run: async (ctx) => `a painting of ${ctx.input.subject}`, +}) + +describe('ctx.agents inside a session', () => { + it('lists, gets, finds, and starts agents', async () => { + const director = definePlugin({ + name: 'test/director', + setup: (ctx) => ({ + commands: { + survey: defineCommand({ + description: 'Look at the agents', + run: () => ({ + names: ctx.agents.list().map((agent) => agent.name), + byName: ctx.agents.get('painter')?.description, + byOutput: ctx.agents.find({ produces: 'image' })?.name, + none: ctx.agents.find({ produces: 'video' }), + }), + }), + paint: defineCommand({ + description: 'Start the painter by name, then by value', + run: async () => [ + await ctx.agents.start('painter', { subject: 'rain' }), + await ctx.agents.start( + painter, + { subject: 'sun' }, + { wake: false }, + ), + ], + }), + }, + }), + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/agents', + adapter: mockAdapter([]).adapter, + agents: [painter], + plugins: () => [director], + }), + { threadId: 't' }, + ) + expect(await session.command('survey')).toEqual({ + names: ['painter'], + byName: 'Paints', + byOutput: 'painter', + none: undefined, + }) + expect(await session.command('paint')).toEqual([ + 'a painting of rain', + 'a painting of sun', + ]) + await host.close() + }) +}) + +describe('plugin services outside a session', () => { + it('ignores events and config, and refuses session-only services', async () => { + const ping = createPluginEvent('test/ping') + let checks: Record = {} + const probe = definePlugin({ + name: 'test/probe', + setup: (ctx) => { + const stop = ctx.on(ping, () => {}) + stop() + ctx.emit(ping, 1) + const refused = (call: () => unknown) => { + try { + call() + return 'ran' + } catch (error) { + return error instanceof Error ? error.message : String(error) + } + } + checks = { + config: ctx.config.get('anything'), + state: refused(() => ctx.state({})), + credentials: refused(() => ctx.credentials.get('x')), + session: refused(() => ctx.session.snapshot()), + run: refused(() => ctx.agents.run('painter')), + start: refused(() => ctx.agents.start('painter')), + group: refused(() => ctx.agents.group({}, async () => 1)), + } + return {} + }, + }) + await mountPlugins([probe], { + threadId: 't', + registry: new AgentRegistry(), + harnessTools: [], + harnessProvides: [], + }) + expect(checks.config).toBeUndefined() + for (const key of [ + 'state', + 'credentials', + 'session', + 'run', + 'start', + 'group', + ]) { + expect(checks[key]).toContain('only available inside a harness session') + } + }) +}) From d8d21525292ff5d6065b01bfc958a1a10b762990 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 15:38:25 +0200 Subject: [PATCH 024/286] test(ai-harness-cli): cover line commands, questions, and the session view --- packages/ai-harness-cli/package.json | 3 +- .../ai-harness-cli/tests/commands.test.ts | 306 ++++++++++++++++++ pnpm-lock.yaml | 3 + 3 files changed, 311 insertions(+), 1 deletion(-) create mode 100644 packages/ai-harness-cli/tests/commands.test.ts diff --git a/packages/ai-harness-cli/package.json b/packages/ai-harness-cli/package.json index e9224933f2..f1d78da10d 100644 --- a/packages/ai-harness-cli/package.json +++ b/packages/ai-harness-cli/package.json @@ -87,6 +87,7 @@ "@tanstack/ai-persistence": "workspace:*", "@tanstack/ai-sandbox": "workspace:*", "@types/react": "^19.2.7", - "@vitest/coverage-v8": "4.1.10" + "@vitest/coverage-v8": "4.1.10", + "zod": "^4.2.0" } } diff --git a/packages/ai-harness-cli/tests/commands.test.ts b/packages/ai-harness-cli/tests/commands.test.ts new file mode 100644 index 0000000000..00fb3735c7 --- /dev/null +++ b/packages/ai-harness-cli/tests/commands.test.ts @@ -0,0 +1,306 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { EventType, defineAgent } from '@tanstack/ai' +import { + HARNESS_EVENTS, + configOption, + createHarnessHost, + defineCommand, + defineHarness, + definePlugin, +} from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { HELP_TEXT, handleLine, parseAnswer } from '../src/commands' +import { applyEvent } from '../src/session-view' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import type { SessionEvent } from '@tanstack/ai-harness' + +/** A model that answers `echo: `, or waits for cancel on "wait". */ +function model(): AnyTextAdapter { + return { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + chatStream: (options: any) => + (async function* (): AsyncGenerator { + const said = String(options.messages.at(-1)?.content ?? '') + const now = Date.now() + yield { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: now, + } + if (said === 'wait') { + const signal: AbortSignal | undefined = + options.abortController?.signal ?? options.request?.signal + await new Promise((resolve) => { + if (!signal || signal.aborted) return resolve() + signal.addEventListener('abort', () => resolve(), { once: true }) + }) + return + } + yield { + type: EventType.TEXT_MESSAGE_START, + messageId: 'm', + role: 'assistant', + timestamp: now, + } + yield { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm', + delta: `echo: ${said}`, + timestamp: now, + } + yield { + type: EventType.TEXT_MESSAGE_END, + messageId: 'm', + timestamp: now, + } + yield { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now, + metadata: { tanstack: { finishReason: 'stop' } }, + } + })(), + } +} + +const tools = definePlugin({ + name: 'test/tools', + setup: (ctx) => ({ + config: { + level: configOption.select({ options: ['low', 'high'], default: 'low' }), + }, + commands: { + quiet: defineCommand({ + description: 'Returns nothing', + run: () => undefined, + }), + stats: defineCommand({ + description: 'Returns an object', + run: () => ({ ok: true }), + }), + broken: defineCommand({ + description: 'Throws', + run: () => { + throw new Error('it broke') + }, + }), + echo: defineCommand({ + description: 'Echoes', + run: (input: unknown) => input, + }), + confirm: defineCommand({ + description: 'Asks a yes or no question', + run: async () => { + const answer = await ctx.session.ask({ + message: 'Sure?', + schema: z.boolean(), + }) + return answer ? 'confirmed' : 'declined' + }, + }), + 'connect:svc': defineCommand({ + description: 'Sign in', + run: () => 'Connected to Svc.', + }), + 'disconnect:svc': defineCommand({ + description: 'Sign out', + run: () => { + throw new Error('not signed in') + }, + }), + }, + }), +}) + +const counter = defineAgent({ + name: 'counter', + description: 'Counts letters', + inputSchema: z.object({ word: z.string() }), + run: async (ctx) => ctx.input.word.length, +}) + +async function open(plugins = true) { + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/cli-commands', + adapter: model(), + ...(plugins ? { agents: [counter], plugins: () => [tools] } : {}), + }), + { threadId: 't' }, + ) + const say = async (line: string) => { + const result = await handleLine(session, line) + return result.type === 'notice' ? result.text : result.type + } + return { host, session, say } +} + +describe('handleLine', () => { + it('answers built-in commands for a bare harness', async () => { + const { host, say } = await open(false) + expect(await say(' ')).toBe('') + expect(await say('/help')).toBe(HELP_TEXT) + expect(await say('/config')).toBe('This harness has no settings.') + expect(await say('/agents')).toBe('This harness has no agents.') + expect(await say('/cancel')).toBe('Nothing is running.') + expect(await say('/status')).toBe( + 'Status: idle. Running: nothing. Queued turns: 0.', + ) + expect(await say('/nope')).toBe('Unknown command: /nope. Type /help.') + expect(await say('/exit')).toBe('exit') + expect(await say('/quit')).toBe('exit') + await host.close() + }) + + it('lists plugin commands, settings, and agents, and changes settings', async () => { + const { host, say } = await open() + expect(await say('/help')).toContain('Plugin commands:\n') + expect(await say('/help')).toContain('/quiet Returns nothing') + expect(await say('/config')).toBe(' level = "low"') + expect(await say('/config level high')).toBe('level changed.') + expect(await say('/config level extreme')).toMatch(/^Not changed: /) + expect(await say('/agents')).toBe(' counter: Counts letters') + await host.close() + }) + + it('runs plugin commands and shows their results or errors', async () => { + const { host, say } = await open() + expect(await say('/quiet')).toBe('Done.') + expect(await say('/stats')).toBe('{\n "ok": true\n}') + expect(await say('/broken')).toBe('/broken failed: it broke') + expect(await say('/echo {"a":1}')).toBe('{\n "a": 1\n}') + expect(await say('/echo plain words')).toBe('plain words') + expect(await say('/connect svc')).toBe('Connected to Svc.') + expect(await say('/disconnect svc')).toBe( + '/disconnect:svc failed: not signed in', + ) + await host.close() + }) + + it('starts agents in the background and checks their input', async () => { + const { host, say } = await open() + expect(await say('/agent')).toBe('Unknown agent: (none)') + expect(await say('/agent ghost')).toBe('Unknown agent: ghost') + expect(await say('/agent counter {bad')).toContain('must be JSON') + expect(await say('/agent counter {"word":"four"}')).toBe( + 'Started counter in the background.', + ) + await host.close() + }) + + it('answers questions, including a rejected answer, then shows the result later', async () => { + const { host, session } = await open() + const asked = await handleLine(session, '/confirm') + expect(asked).toMatchObject({ type: 'notice', text: '' }) + const later = asked.type === 'notice' ? asked.later : undefined + expect(later).toBeDefined() + expect(session.snapshot().pendingQuestions).toHaveLength(1) + expect(await handleLine(session, 'maybe')).toMatchObject({ + text: expect.stringMatching(/^Answer again: /), + }) + expect(await handleLine(session, 'y')).toMatchObject({ text: '' }) + expect(await later).toBe('confirmed') + await host.close() + }) + + it('prompts when idle, steers while running, and cancels', async () => { + const { host, session, say } = await open() + expect(await say('wait')).toBe('sent') + await vi.waitFor(() => expect(session.snapshot().status).toBe('running')) + expect(await say('/status')).toMatch(/^Status: running\. Running: chat/) + expect(await say('go faster')).toBe('sent') + expect(await say('/cancel')).toBe('Cancelled.') + await host.close() + }) +}) + +describe('parseAnswer', () => { + it('reads yes and no for boolean questions, JSON, and plain text', () => { + const boolean = { type: 'boolean' } + expect(parseAnswer('Yes', boolean)).toBe(true) + expect(parseAnswer(' n ', boolean)).toBe(false) + expect(parseAnswer('yes', { type: 'string' })).toBe('yes') + expect(parseAnswer('{"n":2}', undefined)).toEqual({ n: 2 }) + expect(parseAnswer('', undefined)).toBeUndefined() + }) +}) + +describe('applyEvent', () => { + const entry = (event: StreamChunk, operationId = 'op'): SessionEvent => ({ + cursor: '1', + operationId, + event, + }) + const custom = (name: string, value: unknown): StreamChunk => ({ + type: EventType.CUSTOM, + name, + value, + timestamp: 1, + }) + + it('shows questions, sign-ins, resumes, errors, and child agents', () => { + const events: Array = [ + custom(HARNESS_EVENTS.question, { message: 'Sure?' }), + custom(HARNESS_EVENTS.authRequired, { + connector: 'gh', + url: 'https://gh.example/device', + userCode: 'ABCD', + }), + custom(HARNESS_EVENTS.authRequired, { connector: 'svc' }), + custom(HARNESS_EVENTS.operationResumed, {}), + { type: EventType.RUN_ERROR, message: 'model down', timestamp: 1 }, + { + type: EventType.SUBAGENT_STARTED, + subagentRunId: 'child', + name: 'painter', + timestamp: 1, + } as StreamChunk, + ] + const shown = events.reduce( + (entries, event) => applyEvent(entries, entry(event)), + [] as ReturnType, + ) + expect(shown.map((item) => item.text)).toEqual([ + '? Sure?', + 'Sign in to gh. Open https://gh.example/device and enter the code ABCD.', + 'Sign in to svc. Run /connect svc.', + 'Resumed a turn that a crash stopped.', + 'Error: model down', + 'agent painter started', + ]) + // Other child events and unknown custom events change nothing. + const unchanged = applyEvent( + shown, + entry({ + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'x', + delta: 'child text', + subagentRunId: 'child', + timestamp: 1, + } as StreamChunk), + ) + expect(unchanged).toBe(shown) + expect(applyEvent(shown, entry(custom('other.event', {})))).toBe(shown) + }) +}) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 4a914a15c6..320e1d7e5f 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2571,6 +2571,9 @@ importers: '@vitest/coverage-v8': specifier: 4.1.10 version: 4.1.10(vitest@4.1.11) + zod: + specifier: ^4.2.0 + version: 4.3.6 packages/ai-isolate-cloudflare: dependencies: From b68dd4776275bd9bbef52b7a501ea9da7649dfee Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 16:08:33 +0200 Subject: [PATCH 025/286] test(ai-harness): cover remote harnessText and worker edge frames --- .../ai-harness/tests/remote-worker.test.ts | 193 ++++++++++++++++++ 1 file changed, 193 insertions(+) create mode 100644 packages/ai-harness/tests/remote-worker.test.ts diff --git a/packages/ai-harness/tests/remote-worker.test.ts b/packages/ai-harness/tests/remote-worker.test.ts new file mode 100644 index 0000000000..8fa68281e9 --- /dev/null +++ b/packages/ai-harness/tests/remote-worker.test.ts @@ -0,0 +1,193 @@ +import { PassThrough } from 'node:stream' +import { describe, expect, it, vi } from 'vitest' +import { EventType, chat } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + createHarnessHandler, + createHarnessHost, + defineHarness, + harnessText, +} from '../src' +import { runHarnessWorker } from '../src/worker' +import { mockAdapter, text } from './helpers' +import type { StreamChunk } from '@tanstack/ai' + +async function collect(stream: AsyncIterable) { + const chunks: Array = [] + for await (const chunk of stream) chunks.push(chunk) + return chunks +} + +const textOf = (chunks: Array) => + chunks + .map((chunk) => + chunk.type === EventType.TEXT_MESSAGE_CONTENT ? chunk.delta : '', + ) + .join('') + +describe('harnessText for a remote harness', () => { + it('sends the token and the last user message, and streams the answer', async () => { + const { adapter, calls } = mockAdapter([() => text('remote answer')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const handler = createHarnessHandler({ + host, + harness: defineHarness({ name: 'test/remote', adapter }), + authorize: (request) => + request.headers.get('authorization') === 'Bearer secret' + ? { id: 'u' } + : null, + }) + const remote = harnessText({ + url: 'http://remote.test/api/harness/', + token: 'secret', + fetch: (input, init) => handler(new Request(input, init)), + }) + expect(remote.model).toBe('http://remote.test/api/harness') + const chunks = await collect( + chat({ + adapter: remote, + messages: [ + { role: 'user', content: 'older question' }, + { role: 'assistant', content: 'older answer' }, + { role: 'user', content: 'new question' }, + ], + threadId: 'outer', + }) as AsyncIterable, + ) + expect(textOf(chunks)).toBe('remote answer') + expect(JSON.stringify(calls[0].messages)).toContain('new question') + expect(JSON.stringify(calls[0].messages)).not.toContain('older question') + await expect(remote.structuredOutput({} as never)).rejects.toThrow( + 'does not support structured output', + ) + await host.close() + }) + + it('reports a refused request', async () => { + const remote = harnessText({ + url: 'http://remote.test', + fetch: async () => new Response('unauthorized', { status: 401 }), + }) + await expect( + collect( + remote.chatStream({ + model: 'x', + messages: [{ role: 'user', content: 'hi' }], + } as never), + ), + ).rejects.toThrow('Remote harness failed (401): unauthorized') + }) + + it('skips child and unknown data, and ends on a run error', async () => { + const sse = [ + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm', + delta: 'kept', + timestamp: 1, + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'c', + delta: 'child', + subagentRunId: 'x', + timestamp: 1, + }, + 'not an object', + { no: 'type' }, + { type: EventType.RUN_ERROR, message: 'remote broke', timestamp: 1 }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm', + delta: 'after the error', + timestamp: 1, + }, + ] + .map((data) => `data: ${JSON.stringify(data)}\n\n`) + .join('') + const remote = harnessText({ + url: 'http://remote.test', + fetch: async () => new Response(`: comment\n\n${sse}`), + }) + const chunks = await collect( + remote.chatStream({ + model: 'x', + runId: 'run-1', + threadId: 'thread-1', + messages: [], + } as never), + ) + expect(textOf(chunks)).toBe('kept') + expect(chunks.at(-1)).toMatchObject({ + type: EventType.RUN_ERROR, + message: 'remote broke', + }) + }) + + it('finishes cleanly when the response has no body', async () => { + const remote = harnessText({ + url: 'http://remote.test', + fetch: async () => new Response(null, { status: 200 }), + }) + const chunks = await collect( + remote.chatStream({ model: 'x', messages: [] } as never), + ) + expect(chunks.map((chunk) => chunk.type)).toEqual([ + EventType.RUN_STARTED, + EventType.RUN_FINISHED, + ]) + }) +}) + +describe('runHarnessWorker edges', () => { + it('answers bad frames, a second subscribe, and snapshots, and resumes from a cursor', async () => { + const { adapter } = mockAdapter([() => text('worker answer')]) + const input = new PassThrough() + const frames: Array = [] + const running = runHarnessWorker( + defineHarness({ name: 'test/worker-edges', adapter }), + { + input, + output: { write: (line: string) => frames.push(JSON.parse(line)) }, + persistence: memoryPersistence(), + }, + ) + const send = (frame: unknown) => + input.write( + `${typeof frame === 'string' ? frame : JSON.stringify(frame)}\n`, + ) + + send('') + send({ type: 'harness.snapshot' }) + await vi.waitFor(() => + expect(frames.at(-1)).toEqual({ + type: 'harness.error', + message: 'Send harness.subscribe first.', + }), + ) + send('{not json') + await vi.waitFor(() => expect(frames).toHaveLength(2)) + expect(frames[1].type).toBe('harness.error') + + send({ type: 'harness.subscribe', threadId: 'w', from: '0' }) + await vi.waitFor(() => + expect(frames.some((frame) => frame.type === 'harness.hello')).toBe(true), + ) + send({ type: 'harness.subscribe', threadId: 'w' }) + await vi.waitFor(() => + expect(frames.at(-1)).toEqual({ + type: 'harness.error', + message: 'Already subscribed.', + }), + ) + send({ type: 'harness.snapshot' }) + await vi.waitFor(() => + expect(frames.at(-1)).toMatchObject({ + type: 'harness.snapshot', + snapshot: { threadId: 'w', status: 'idle' }, + }), + ) + input.end() + await running + }) +}) From ed8a768202c2b25c6ff32154440114f19fd5e815 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 16:10:47 +0200 Subject: [PATCH 026/286] test(ai-harness-cli): cover --dashboard pairing and saved-token reconnect --- .../ai-harness-cli/tests/dashboard.test.ts | 118 ++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 packages/ai-harness-cli/tests/dashboard.test.ts diff --git a/packages/ai-harness-cli/tests/dashboard.test.ts b/packages/ai-harness-cli/tests/dashboard.test.ts new file mode 100644 index 0000000000..62d5f0d2b4 --- /dev/null +++ b/packages/ai-harness-cli/tests/dashboard.test.ts @@ -0,0 +1,118 @@ +import { describe, expect, it, vi } from 'vitest' +import { EventType } from '@tanstack/ai' +import { startDashboard } from '@tanstack/ai-dashboard' +import { defineHarness } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { runCli } from '../src' +import type { AnyTextAdapter } from '@tanstack/ai' + +const adapter: AnyTextAdapter = { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + chatStream: () => + (async function* () { + yield { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: 1, + } + })(), +} + +/** + * Stop a CLI that waits for Ctrl+C by calling the listeners it added, not by + * sending a real signal to the test process. + */ +function stopper() { + const before = { + SIGINT: new Set(process.listeners('SIGINT')), + SIGTERM: new Set(process.listeners('SIGTERM')), + } + return () => { + for (const signal of ['SIGINT', 'SIGTERM'] as const) { + for (const listener of process.listeners(signal)) { + if (before[signal].has(listener)) continue + process.removeListener(signal, listener) + if (signal === 'SIGINT') (listener as () => void)() + } + } + } +} + +describe('--dashboard', () => { + it('pairs on first use, then reconnects with the saved token', async () => { + const dashboard = await startDashboard({ port: 0 }) + const owner = { + Authorization: `Bearer ${dashboard.ownerToken}`, + 'Content-Type': 'application/json', + } + const harness = defineHarness({ name: 'test/dashboard-cli', adapter }) + try { + let log = '' + const stderr = { + write: (text: string) => { + log += text + const code = /with the code (\S+)\./.exec(text)?.[1] + if (code) { + void fetch(`${dashboard.url}/api/pair/approve`, { + method: 'POST', + headers: owner, + body: JSON.stringify({ code }), + }) + } + return true + }, + } + const stop = stopper() + const first = runCli(harness, { + argv: ['--dashboard', dashboard.url], + env: {}, + stderr, + stdout: { write: () => true }, + persistence: memoryPersistence(), + }) + await vi.waitFor(() => expect(log).toContain('Connected to'), { + timeout: 5000, + }) + stop() + expect(await first).toBe(0) + const token = /HARNESS_DASHBOARD_TOKEN=(\S+) /.exec(log)?.[1] + expect(token).toBeDefined() + + log = '' + const stopAgain = stopper() + const second = runCli(harness, { + argv: ['--dashboard', dashboard.url], + env: { HARNESS_DASHBOARD_TOKEN: token ?? '' }, + stderr, + stdout: { write: () => true }, + persistence: memoryPersistence(), + }) + await vi.waitFor(() => expect(log).toContain('Connected to'), { + timeout: 5000, + }) + expect(log).not.toContain('Pair this host') + stopAgain() + expect(await second).toBe(0) + } finally { + await dashboard.close() + } + }) +}) From f72f17c83e969845464fd2f6380daa462da6f6bd Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 16:15:15 +0200 Subject: [PATCH 027/286] fix(ai-mcp): refresh expired connector tokens outside /connect; cover saved, refreshed, and revoked sign-ins --- packages/ai-mcp/src/connector.ts | 10 +- packages/ai-mcp/tests/connector-edges.test.ts | 277 ++++++++++++++++++ packages/ai-mcp/tests/connector.test.ts | 118 +------- packages/ai-mcp/tests/protected-server.ts | 133 +++++++++ 4 files changed, 421 insertions(+), 117 deletions(-) create mode 100644 packages/ai-mcp/tests/connector-edges.test.ts create mode 100644 packages/ai-mcp/tests/protected-server.ts diff --git a/packages/ai-mcp/src/connector.ts b/packages/ai-mcp/src/connector.ts index 3a9e94ab94..13d9327327 100644 --- a/packages/ai-mcp/src/connector.ts +++ b/packages/ai-mcp/src/connector.ts @@ -216,12 +216,20 @@ export function mcpConnector(options: McpConnectorOptions) { discoverTools: async () => { if (!connected) return [] if (!tools) { + // The SDK only refreshes tokens for a client with a redirect URL. + // Without one it treats the client as machine-to-machine. A new + // browser sign-in still happens only through /connect. + const saved = await ctx.credentials.get(id) + const redirectUri = + (saved?.type === 'oauth' + ? saved.client?.redirectUri + : undefined) ?? 'http://127.0.0.1/callback' client = await createMCPClient({ transport: { type: 'http', url, authProvider: credentialProvider(id, ctx.credentials, { - redirectUri: undefined, + redirectUri, clientName, scopes: options.scopes, }), diff --git a/packages/ai-mcp/tests/connector-edges.test.ts b/packages/ai-mcp/tests/connector-edges.test.ts new file mode 100644 index 0000000000..af6565cfbc --- /dev/null +++ b/packages/ai-mcp/tests/connector-edges.test.ts @@ -0,0 +1,277 @@ +import { afterEach, describe, expect, it } from 'vitest' +import { EventType } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + HARNESS_EVENTS, + createHarnessHost, + defineHarness, +} from '@tanstack/ai-harness' +import { mcpConnector } from '../src/connector' +import { startProtectedServer } from './protected-server' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import type { SessionEvent } from '@tanstack/ai-harness' + +/** A model that records the tools it gets, calls `tool` once, then stops. */ +function recorder(tool?: string) { + const calls: Array = [] + const now = () => Date.now() + const adapter: AnyTextAdapter = { + kind: 'text', + name: 'mock', + model: 'mock', + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + chatStream: (options) => { + calls.push(options) + const first = calls.length === 1 && tool !== undefined + return (async function* (): AsyncGenerator { + yield { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: now(), + } + if (first) { + yield { + type: EventType.TOOL_CALL_START, + toolCallId: 'c1', + toolCallName: tool, + timestamp: now(), + } + yield { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'c1', + delta: '{"text":"hi"}', + timestamp: now(), + } + yield { + type: EventType.TOOL_CALL_END, + toolCallId: 'c1', + timestamp: now(), + } + } + yield { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { + tanstack: { finishReason: first ? 'tool_calls' : 'stop' }, + }, + } + })() + }, + } + const toolNames = (index: number) => + (calls[index]?.tools ?? []).map((entry: { name: string }) => entry.name) + return { adapter, calls, toolNames } +} + +const scope = { threadId: 't', userId: 'user-1' } +const cleanups: Array<() => unknown> = [] +afterEach(async () => { + for (const cleanup of cleanups.splice(0).reverse()) await cleanup() +}) + +async function setup( + credential?: Parameters< + ReturnType['stores']['credentials']['set'] + >[2], +) { + const server = await startProtectedServer() + cleanups.push(() => server.close()) + const persistence = memoryPersistence() + if (credential) + await persistence.stores.credentials.set(scope, 'demo', credential) + const host = createHarnessHost({ persistence }) + cleanups.push(() => host.close()) + return { server, persistence, host } +} + +describe('mcpConnector with a saved sign-in', () => { + it('uses the saved token and client after a restart, without /connect', async () => { + const { server, host } = await setup({ + type: 'oauth', + accessToken: 'access-1', + refreshToken: 'refresh-1', + expiresAt: Date.now() + 3_600_000, + client: { + clientId: 'client-1', + clientSecret: 'secret-1', + redirectUri: 'http://127.0.0.1:1/callback', + }, + }) + const model = recorder('demo_echo') + const session = await host.open( + defineHarness({ + name: 'test/saved', + adapter: model.adapter, + plugins: () => [ + mcpConnector({ id: 'demo', label: 'Demo', url: server.url }), + ], + }), + { threadId: 't', principal: { id: 'user-1' } }, + ) + await session.prompt('echo') + expect(model.toolNames(0)).toEqual(['demo_echo']) + expect(JSON.stringify(model.calls[1].messages)).toContain('echo: hi') + // The status prompt is empty once connected. + expect(JSON.stringify(model.calls[0].systemPrompts ?? [])).not.toContain( + 'is not connected', + ) + expect(server.seen.registrations).toHaveLength(0) + }) + + it('refreshes an expired token and keeps the old refresh token and client', async () => { + const { server, persistence, host } = await setup({ + type: 'oauth', + accessToken: 'stale', + refreshToken: 'refresh-1', + expiresAt: Date.now() - 1000, + client: { clientId: 'client-1' }, + }) + const model = recorder('demo_echo') + const session = await host.open( + defineHarness({ + name: 'test/refresh', + adapter: model.adapter, + plugins: () => [ + mcpConnector({ id: 'demo', label: 'Demo', url: server.url }), + ], + }), + { threadId: 't', principal: { id: 'user-1' } }, + ) + await session.prompt('echo') + expect(model.toolNames(0)).toEqual(['demo_echo']) + expect( + server.seen.tokenRequests.some( + (form) => form.get('grant_type') === 'refresh_token', + ), + ).toBe(true) + expect( + await persistence.stores.credentials.get(scope, 'demo'), + ).toMatchObject({ + accessToken: 'access-2', + refreshToken: 'refresh-1', + client: { clientId: 'client-1' }, + }) + }) + + it('drops the tools with a warning when the sign-in no longer works', async () => { + const { server, host } = await setup({ + type: 'oauth', + accessToken: 'stale', + refreshToken: 'revoked', + }) + const model = recorder() + const session = await host.open( + defineHarness({ + name: 'test/revoked', + adapter: model.adapter, + plugins: () => [ + mcpConnector({ id: 'demo', label: 'Demo', url: server.url }), + ], + }), + { threadId: 't', principal: { id: 'user-1' } }, + ) + const seen: Array = [] + const controller = new AbortController() + const reading = (async () => { + for await (const entry of session.events({ + from: '0', + signal: controller.signal, + })) + seen.push(entry) + })() + await session.prompt('hi') + controller.abort() + await reading + expect(model.toolNames(0)).toEqual([]) + const warning = seen.find( + (entry) => + entry.event.type === EventType.CUSTOM && + entry.event.name === 'harness.plugin.warning', + ) + expect(warning?.event).toMatchObject({ + value: { plugin: 'connector/demo' }, + }) + }) +}) + +describe('mcpConnector options', () => { + it('asks for scopes, shows the client name, and uses the prefix and approval rule', async () => { + const { server, persistence, host } = await setup() + const model = recorder() + const session = await host.open( + defineHarness({ + name: 'test/options', + adapter: model.adapter, + plugins: () => [ + mcpConnector({ + id: 'demo', + label: 'Demo', + url: server.url, + prefix: 'd', + scopes: ['read'], + clientName: 'Acme Agent', + needsApproval: () => true, + fetch: (input, init) => fetch(input, init), + }), + ], + }), + { threadId: 't', principal: { id: 'user-1' } }, + ) + const controller = new AbortController() + let scopeParam: string | null = null + void (async () => { + for await (const entry of session.events({ signal: controller.signal })) { + const event = entry.event + if ( + event.type !== EventType.CUSTOM || + event.name !== HARNESS_EVENTS.authRequired + ) + continue + const url = new URL(String((event.value as { url: string }).url)) + scopeParam = url.searchParams.get('scope') + const redirect = new URL(url.searchParams.get('redirect_uri') ?? '') + redirect.searchParams.set('code', 'code-1') + redirect.searchParams.set('state', url.searchParams.get('state') ?? '') + await fetch(redirect) + } + })() + cleanups.push(() => controller.abort()) + + expect(await session.command('connect:demo')).toBe('Connected to Demo.') + expect(scopeParam).toBe('read') + expect(server.seen.registrations[0]).toMatchObject({ + client_name: 'Acme Agent', + scope: 'read', + }) + expect( + await persistence.stores.credentials.get(scope, 'demo'), + ).toMatchObject({ + scopes: ['read', 'write'], + client: { + clientId: 'client-1', + redirectUri: expect.stringContaining('127.0.0.1'), + }, + }) + + await session.prompt('hi') + expect(model.toolNames(0)).toEqual(['d_echo']) + expect(model.calls[0].tools[0].needsApproval).toBe(true) + }) +}) diff --git a/packages/ai-mcp/tests/connector.test.ts b/packages/ai-mcp/tests/connector.test.ts index 55261ca78b..5e299a4115 100644 --- a/packages/ai-mcp/tests/connector.test.ts +++ b/packages/ai-mcp/tests/connector.test.ts @@ -1,8 +1,4 @@ -import { createServer } from 'node:http' import { afterEach, describe, expect, it, vi } from 'vitest' -import { z } from 'zod' -import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' -import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js' import { EventType } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' import { @@ -11,119 +7,9 @@ import { defineHarness, } from '@tanstack/ai-harness' import { mcpConnector } from '../src/connector' -import type { IncomingMessage, Server } from 'node:http' +import { startProtectedServer } from './protected-server' import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' -/** - * A local MCP server behind OAuth, like Notion's or Linear's: protected - * resource metadata, authorization server metadata, dynamic client - * registration, a token endpoint, and an `echo` tool. - */ -async function startProtectedServer() { - const seen = { registrations: 0, tokenRequests: [] as Array } - let base = '' - const readBody = async (req: IncomingMessage) => { - const chunks: Array = [] - for await (const chunk of req) chunks.push(chunk as Buffer) - return Buffer.concat(chunks).toString('utf8') - } - const server: Server = createServer((req, res) => { - void (async () => { - const url = new URL(req.url ?? '/', base) - const json = (status: number, value: unknown) => { - res.writeHead(status, { 'Content-Type': 'application/json' }) - res.end(JSON.stringify(value)) - } - if (url.pathname.startsWith('/.well-known/oauth-protected-resource')) { - return json(200, { - resource: `${base}/mcp`, - authorization_servers: [base], - }) - } - if (url.pathname.startsWith('/.well-known/oauth-authorization-server')) { - return json(200, { - issuer: base, - authorization_endpoint: `${base}/authorize`, - token_endpoint: `${base}/token`, - registration_endpoint: `${base}/register`, - response_types_supported: ['code'], - grant_types_supported: ['authorization_code', 'refresh_token'], - code_challenge_methods_supported: ['S256'], - token_endpoint_auth_methods_supported: ['none'], - }) - } - if (url.pathname === '/register' && req.method === 'POST') { - seen.registrations += 1 - const metadata = JSON.parse(await readBody(req)) - return json(201, { ...metadata, client_id: 'client-1' }) - } - if (url.pathname === '/token' && req.method === 'POST') { - const form = new URLSearchParams(await readBody(req)) - seen.tokenRequests.push(form) - if ( - form.get('grant_type') === 'authorization_code' && - form.get('code') === 'code-1' - ) { - return json(200, { - access_token: 'access-1', - refresh_token: 'refresh-1', - token_type: 'Bearer', - expires_in: 3600, - }) - } - return json(400, { error: 'invalid_grant' }) - } - if (url.pathname === '/mcp') { - if (req.headers.authorization !== 'Bearer access-1') { - res.writeHead(401, { - 'WWW-Authenticate': `Bearer resource_metadata="${base}/.well-known/oauth-protected-resource"`, - }) - res.end() - return - } - if (req.method !== 'POST') { - res.writeHead(405).end() - return - } - const mcp = new McpServer({ name: 'protected', version: '1.0.0' }) - mcp.registerTool( - 'echo', - { - description: 'Echo text', - inputSchema: { text: z.string() }, - annotations: { readOnlyHint: true }, - }, - async ({ text }) => ({ - content: [{ type: 'text' as const, text: `echo: ${text}` }], - }), - ) - const transport = new StreamableHTTPServerTransport({ - sessionIdGenerator: undefined, - enableJsonResponse: true, - }) - await mcp.connect(transport) - await transport.handleRequest(req, res, JSON.parse(await readBody(req))) - return - } - res.writeHead(404).end() - })() - }) - await new Promise((resolve) => - server.listen(0, '127.0.0.1', () => resolve()), - ) - const address = server.address() - base = `http://127.0.0.1:${typeof address === 'object' && address ? address.port : 0}` - return { - url: `${base}/mcp`, - seen, - close: () => - new Promise((resolve) => { - server.closeAllConnections() - server.close(() => resolve()) - }), - } -} - /** A model that calls `tool` once, then answers with the tool result it saw. */ function modelCalling( tool: string, @@ -278,7 +164,7 @@ describe('mcpConnector', () => { new URL(authorizationUrl?.searchParams.get('redirect_uri') ?? '') .hostname, ).toBe('127.0.0.1') - expect(protectedServer.seen.registrations).toBe(1) + expect(protectedServer.seen.registrations).toHaveLength(1) expect( protectedServer.seen.tokenRequests[0]?.get('code_verifier'), ).toBeTruthy() diff --git a/packages/ai-mcp/tests/protected-server.ts b/packages/ai-mcp/tests/protected-server.ts new file mode 100644 index 0000000000..6d225bcdb3 --- /dev/null +++ b/packages/ai-mcp/tests/protected-server.ts @@ -0,0 +1,133 @@ +import { createServer } from 'node:http' +import { z } from 'zod' +import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' +import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js' +import type { IncomingMessage, Server } from 'node:http' + +/** + * A local MCP server behind OAuth, like Notion's or Linear's: protected + * resource metadata, authorization server metadata, dynamic client + * registration, a token endpoint, and an `echo` tool. + * + * - The code `code-1` gives `access-1` and `refresh-1`. + * - The refresh token `refresh-1` gives `access-2` (and no new refresh token). + * - `revoke(token)` makes the MCP endpoint refuse a token. + */ +export async function startProtectedServer() { + const seen = { + registrations: [] as Array>, + tokenRequests: [] as Array, + } + const valid = new Set(['access-1']) + let base = '' + const readBody = async (req: IncomingMessage) => { + const chunks: Array = [] + for await (const chunk of req) chunks.push(chunk as Buffer) + return Buffer.concat(chunks).toString('utf8') + } + const server: Server = createServer((req, res) => { + void (async () => { + const url = new URL(req.url ?? '/', base) + const json = (status: number, value: unknown) => { + res.writeHead(status, { 'Content-Type': 'application/json' }) + res.end(JSON.stringify(value)) + } + if (url.pathname.startsWith('/.well-known/oauth-protected-resource')) { + return json(200, { + resource: `${base}/mcp`, + authorization_servers: [base], + }) + } + if (url.pathname.startsWith('/.well-known/oauth-authorization-server')) { + return json(200, { + issuer: base, + authorization_endpoint: `${base}/authorize`, + token_endpoint: `${base}/token`, + registration_endpoint: `${base}/register`, + response_types_supported: ['code'], + grant_types_supported: ['authorization_code', 'refresh_token'], + code_challenge_methods_supported: ['S256'], + token_endpoint_auth_methods_supported: ['none'], + }) + } + if (url.pathname === '/register' && req.method === 'POST') { + const metadata = JSON.parse(await readBody(req)) + seen.registrations.push(metadata) + return json(201, { ...metadata, client_id: 'client-1' }) + } + if (url.pathname === '/token' && req.method === 'POST') { + const form = new URLSearchParams(await readBody(req)) + seen.tokenRequests.push(form) + if ( + form.get('grant_type') === 'authorization_code' && + form.get('code') === 'code-1' + ) { + return json(200, { + access_token: 'access-1', + refresh_token: 'refresh-1', + token_type: 'Bearer', + expires_in: 3600, + scope: 'read write', + }) + } + if ( + form.get('grant_type') === 'refresh_token' && + form.get('refresh_token') === 'refresh-1' + ) { + valid.add('access-2') + return json(200, { access_token: 'access-2', token_type: 'Bearer' }) + } + return json(400, { error: 'invalid_grant' }) + } + if (url.pathname === '/mcp') { + const token = (req.headers.authorization ?? '').replace(/^Bearer /, '') + if (!valid.has(token)) { + res.writeHead(401, { + 'WWW-Authenticate': `Bearer resource_metadata="${base}/.well-known/oauth-protected-resource"`, + }) + res.end() + return + } + if (req.method !== 'POST') { + res.writeHead(405).end() + return + } + const mcp = new McpServer({ name: 'protected', version: '1.0.0' }) + mcp.registerTool( + 'echo', + { + description: 'Echo text', + inputSchema: { text: z.string() }, + annotations: { readOnlyHint: true }, + }, + async ({ text }) => ({ + content: [{ type: 'text' as const, text: `echo: ${text}` }], + }), + ) + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: undefined, + enableJsonResponse: true, + }) + await mcp.connect(transport) + await transport.handleRequest(req, res, JSON.parse(await readBody(req))) + return + } + res.writeHead(404).end() + })() + }) + await new Promise((resolve) => + server.listen(0, '127.0.0.1', () => resolve()), + ) + const address = server.address() + base = `http://127.0.0.1:${typeof address === 'object' && address ? address.port : 0}` + return { + url: `${base}/mcp`, + seen, + revoke: (token: string) => valid.delete(token), + close: () => + new Promise((resolve) => { + server.closeAllConnections() + server.close(() => resolve()) + }), + } +} From ec2abc4f12d26e31ed9d85b834ae9f5c7c0a1d3c Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 16:42:28 +0200 Subject: [PATCH 028/286] test(ai-mcp, ai-opencode): cover stored connector credentials and the opencode server start paths --- packages/ai-mcp/tests/connector-edges.test.ts | 93 ++++++++++++ .../ai-opencode/tests/sandbox-server.test.ts | 141 ++++++++++++++++++ 2 files changed, 234 insertions(+) create mode 100644 packages/ai-opencode/tests/sandbox-server.test.ts diff --git a/packages/ai-mcp/tests/connector-edges.test.ts b/packages/ai-mcp/tests/connector-edges.test.ts index af6565cfbc..7aa3b5e7c9 100644 --- a/packages/ai-mcp/tests/connector-edges.test.ts +++ b/packages/ai-mcp/tests/connector-edges.test.ts @@ -275,3 +275,96 @@ describe('mcpConnector options', () => { expect(model.calls[0].tools[0].needsApproval).toBe(true) }) }) + +describe('mcpConnector stored credential shapes', () => { + function connectorHarness(url: string, model: ReturnType) { + return defineHarness({ + name: 'test/shapes', + adapter: model.adapter, + plugins: () => [mcpConnector({ id: 'demo', label: 'Demo', url })], + }) + } + + it('works with a bare token, reuses the tools, and signs out before any tool call', async () => { + const { server, persistence, host } = await setup({ + type: 'oauth', + accessToken: 'access-1', + }) + const model = recorder() + const session = await host.open(connectorHarness(server.url, model), { + threadId: 't', + principal: { id: 'user-1' }, + }) + await session.prompt('one') + await session.prompt('two') + expect(model.toolNames(0)).toEqual(['demo_echo']) + expect(model.toolNames(1)).toEqual(['demo_echo']) + + const fresh = await host.open(connectorHarness(server.url, recorder()), { + threadId: 't2', + principal: { id: 'user-1' }, + }) + expect(await fresh.command('disconnect:demo')).toBe( + 'Disconnected from Demo.', + ) + expect(await persistence.stores.credentials.get(scope, 'demo')).toBeNull() + }) + + it('refreshes a token saved without a client, and keeps a client secret', async () => { + const bare = await setup({ + type: 'oauth', + accessToken: 'stale', + refreshToken: 'refresh-1', + }) + const model = recorder() + const session = await bare.host.open( + connectorHarness(bare.server.url, model), + { + threadId: 't', + principal: { id: 'user-1' }, + }, + ) + await session.prompt('hi') + expect(model.toolNames(0)).toEqual(['demo_echo']) + const saved = await bare.persistence.stores.credentials.get(scope, 'demo') + expect(saved).toMatchObject({ + accessToken: 'access-2', + refreshToken: 'refresh-1', + }) + // No saved client: the SDK registered one before it refreshed. + expect(bare.server.seen.registrations).toHaveLength(1) + expect(saved).toMatchObject({ client: { clientId: 'client-1' } }) + + const withSecret = await setup({ + type: 'oauth', + accessToken: 'stale', + refreshToken: 'refresh-1', + client: { clientId: 'client-1', clientSecret: 'secret-1' }, + }) + const other = await withSecret.host.open( + connectorHarness(withSecret.server.url, recorder()), + { threadId: 't', principal: { id: 'user-1' } }, + ) + await other.prompt('hi') + expect( + await withSecret.persistence.stores.credentials.get(scope, 'demo'), + ).toMatchObject({ + accessToken: 'access-2', + client: { clientId: 'client-1', clientSecret: 'secret-1' }, + }) + }) + + it('treats a saved API key as no sign-in for the MCP server', async () => { + const { server, host } = await setup({ + type: 'api_key', + value: 'not-oauth', + }) + const model = recorder() + const session = await host.open(connectorHarness(server.url, model), { + threadId: 't', + principal: { id: 'user-1' }, + }) + await session.prompt('hi') + expect(model.toolNames(0)).toEqual([]) + }) +}) diff --git a/packages/ai-opencode/tests/sandbox-server.test.ts b/packages/ai-opencode/tests/sandbox-server.test.ts new file mode 100644 index 0000000000..583c013186 --- /dev/null +++ b/packages/ai-opencode/tests/sandbox-server.test.ts @@ -0,0 +1,141 @@ +import { describe, expect, it } from 'vitest' +import { startOpencodeServerInSandbox } from '../src/process/sandbox-server' +import type { SandboxHandle, SpawnHandle } from '@tanstack/ai-sandbox' + +async function* chunks(values: Array): AsyncIterable { + for (const value of values) { + await Promise.resolve() + yield value + } +} + +async function* failing(message: string): AsyncIterable { + await Promise.resolve() + throw new Error(message) +} + +/** An iterable that never yields, like a server that prints nothing. */ +function silent(): AsyncIterable { + return { + [Symbol.asyncIterator]: () => ({ next: () => new Promise(() => {}) }), + } +} + +function sandboxWith( + streams: { stdout: AsyncIterable; stderr: AsyncIterable }, + channel: { url: string; headers?: Record; token?: string }, +): { sandbox: SandboxHandle; spawned: Array } { + const spawned: Array = [] + const handle: SpawnHandle = { + pid: 1, + ...streams, + stdin: { write: () => Promise.resolve(), end: () => Promise.resolve() }, + wait: () => Promise.resolve(0), + kill: () => Promise.resolve(), + } + const sandbox: SandboxHandle = { + id: 'sbx', + provider: 'mock', + capabilities: {} as SandboxHandle['capabilities'], + fs: {} as SandboxHandle['fs'], + git: {} as SandboxHandle['git'], + process: { + exec: () => Promise.reject(new Error('unused')), + spawn: (command, options) => { + spawned.push({ command, options }) + return Promise.resolve(handle) + }, + }, + ports: { connect: () => Promise.resolve(channel) }, + env: { set: () => Promise.resolve() }, + destroy: () => Promise.resolve(), + } + return { sandbox, spawned } +} + +const ready = 'opencode server listening on http://0.0.0.0:4096\n' + +describe('startOpencodeServerInSandbox edges', () => { + it('passes the channel headers through, with the host, env, and signal', async () => { + const { sandbox, spawned } = sandboxWith( + { stdout: chunks([ready]), stderr: chunks(['warming up\n']) }, + { url: 'https://preview.test', headers: { 'x-preview-token': 't' } }, + ) + const controller = new AbortController() + const server = await startOpencodeServerInSandbox(sandbox, { + port: 4096, + hostname: '127.0.0.1', + cwd: '/workspace', + env: { OPENCODE_CONFIG_CONTENT: '{}' }, + signal: controller.signal, + }) + expect(server.headers).toEqual({ 'x-preview-token': 't' }) + expect(spawned[0]).toMatchObject({ + command: 'opencode serve --hostname=127.0.0.1 --port=4096', + options: { cwd: '/workspace', env: { OPENCODE_CONFIG_CONTENT: '{}' } }, + }) + }) + + it('turns a channel token into a bearer header', async () => { + const { sandbox } = sandboxWith( + { stdout: chunks([ready]), stderr: chunks([]) }, + { url: 'https://preview.test', token: 'tok' }, + ) + const server = await startOpencodeServerInSandbox(sandbox, { + port: 4096, + cwd: '/workspace', + }) + expect(server.headers).toEqual({ Authorization: 'Bearer tok' }) + }) + + it('times out with the server output in the error', async () => { + const { sandbox } = sandboxWith( + { stdout: silent(), stderr: chunks(['address already in use\n']) }, + { url: 'http://127.0.0.1:4096' }, + ) + await expect( + startOpencodeServerInSandbox(sandbox, { + port: 4096, + cwd: '/workspace', + timeoutMs: 30, + }), + ).rejects.toThrow( + 'opencode serve did not become ready within 30ms: address already in use', + ) + }) + + it('times out without output, and reports an exit with no output', async () => { + const quiet = sandboxWith( + { stdout: silent(), stderr: chunks([]) }, + { url: 'http://127.0.0.1:4096' }, + ) + await expect( + startOpencodeServerInSandbox(quiet.sandbox, { + port: 4096, + cwd: '/workspace', + timeoutMs: 20, + }), + ).rejects.toThrow(/^opencode serve did not become ready within 20ms$/) + + const exited = sandboxWith( + { stdout: chunks([]), stderr: chunks([]) }, + { url: 'http://127.0.0.1:4096' }, + ) + await expect( + startOpencodeServerInSandbox(exited.sandbox, { + port: 4096, + cwd: '/workspace', + }), + ).rejects.toThrow('exited before becoming ready (no output)') + }) + + it('rejects when stdout fails, and ignores a failing stderr', async () => { + const { sandbox } = sandboxWith( + { stdout: failing('stdout closed'), stderr: failing('stderr closed') }, + { url: 'http://127.0.0.1:4096' }, + ) + await expect( + startOpencodeServerInSandbox(sandbox, { port: 4096, cwd: '/workspace' }), + ).rejects.toThrow('stdout closed') + }) +}) From a65d4bad66b0e30daf2f3e1fe03003b46d4f988d Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 26 Sep 2026 16:42:58 +0200 Subject: [PATCH 029/286] test(ai-opencode): use a plain failing iterable --- packages/ai-opencode/tests/sandbox-server.test.ts | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/packages/ai-opencode/tests/sandbox-server.test.ts b/packages/ai-opencode/tests/sandbox-server.test.ts index 583c013186..c4a5ba9e94 100644 --- a/packages/ai-opencode/tests/sandbox-server.test.ts +++ b/packages/ai-opencode/tests/sandbox-server.test.ts @@ -9,9 +9,13 @@ async function* chunks(values: Array): AsyncIterable { } } -async function* failing(message: string): AsyncIterable { - await Promise.resolve() - throw new Error(message) +/** An iterable whose first read fails. */ +function failing(message: string): AsyncIterable { + return { + [Symbol.asyncIterator]: () => ({ + next: () => Promise.reject(new Error(message)), + }), + } } /** An iterable that never yields, like a server that prints nothing. */ From 9409b06db836e9d897569266d17747bc6f9fadef Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 28 Sep 2026 11:23:13 +0200 Subject: [PATCH 030/286] feat(ai-sandbox, ai-harness): codingAgents plugin to delegate to Claude Code, Codex, and other coding agents; show child agent work in the CLI, ACP, and dashboard --- .changeset/harness-p9-coding-agents.md | 13 + docs/config.json | 11 +- docs/harness/coding-agent.md | 4 + docs/harness/coding-agents.md | 122 ++++++ docs/harness/dashboard.md | 2 +- docs/harness/deploy.md | 2 +- docs/harness/plugins.md | 1 + docs/harness/subagents.md | 2 +- examples/harness-cli/.env.example | 4 + examples/harness-cli/.gitignore | 2 + examples/harness-cli/README.md | 10 + examples/harness-cli/package.json | 4 + examples/harness-cli/src/harness.ts | 55 +++ packages/ai-acp/src/agent/index.ts | 7 +- packages/ai-acp/tests/child-updates.test.ts | 62 +++ packages/ai-dashboard/src/ui.ts | 31 +- .../ai-dashboard/tests/child-agents.test.ts | 143 +++++++ packages/ai-harness-cli/src/interactive.tsx | 7 + packages/ai-harness-cli/src/lines.ts | 4 +- packages/ai-harness-cli/src/session-view.ts | 85 +++- .../ai-harness-cli/tests/child-view.test.ts | 226 +++++++++++ .../ai-harness-cli/tests/commands.test.ts | 11 +- packages/ai-harness/src/plugins.ts | 14 + packages/ai-harness/src/session.ts | 12 +- .../ai-harness/tests/plugin-subagents.test.ts | 74 ++++ packages/ai-sandbox/package.json | 9 + packages/ai-sandbox/src/harness.ts | 210 ++++++++++ .../ai-sandbox/tests/coding-agents.test.ts | 377 ++++++++++++++++++ packages/ai-sandbox/vite.config.ts | 7 +- pnpm-lock.yaml | 15 + 30 files changed, 1502 insertions(+), 24 deletions(-) create mode 100644 .changeset/harness-p9-coding-agents.md create mode 100644 docs/harness/coding-agents.md create mode 100644 packages/ai-acp/tests/child-updates.test.ts create mode 100644 packages/ai-dashboard/tests/child-agents.test.ts create mode 100644 packages/ai-harness-cli/tests/child-view.test.ts create mode 100644 packages/ai-harness/tests/plugin-subagents.test.ts create mode 100644 packages/ai-sandbox/src/harness.ts create mode 100644 packages/ai-sandbox/tests/coding-agents.test.ts diff --git a/.changeset/harness-p9-coding-agents.md b/.changeset/harness-p9-coding-agents.md new file mode 100644 index 0000000000..62a57e4b40 --- /dev/null +++ b/.changeset/harness-p9-coding-agents.md @@ -0,0 +1,13 @@ +--- +'@tanstack/ai-sandbox': minor +'@tanstack/ai-harness': minor +'@tanstack/ai-harness-cli': minor +'@tanstack/ai-acp': minor +'@tanstack/ai-dashboard': minor +--- + +`@tanstack/ai-sandbox/harness` adds `codingAgents({ sandbox, agents, workspace? })`: a harness plugin that gives the lead model one tool per coding agent (Claude Code, Codex, Grok Build, or any ACP agent). Each agent runs in the sandbox, keeps its own session per thread (also after a restart), and starts read-only in the harness `plan` mode. `workspace: 'shared'` (default) runs one agent at a time in one sandbox per thread. `'per-agent'` gives each agent its own sandbox. `/fresh [agent]` starts new sessions. + +`@tanstack/ai-harness` plugins can contribute `subagents`: agents the model can call as tools. + +Child agent work is now visible: the CLI shows each child's tool calls and a finish line with the start of its answer, ACP editors get the child's tool calls, and the dashboard shows a block per child. diff --git a/docs/config.json b/docs/config.json index 71a5f53d3a..db8c43e1bc 100644 --- a/docs/config.json +++ b/docs/config.json @@ -827,12 +827,14 @@ { "label": "Write a plugin", "to": "harness/plugins", - "addedAt": "2026-09-26" + "addedAt": "2026-09-26", + "updatedAt": "2026-09-28" }, { "label": "Build a coding agent", "to": "harness/coding-agent", - "addedAt": "2026-09-26" + "addedAt": "2026-09-26", + "updatedAt": "2026-09-28" }, { "label": "Auth and connectors", @@ -849,6 +851,11 @@ "to": "harness/code-mode", "addedAt": "2026-09-26" }, + { + "label": "Delegate to coding agents", + "to": "harness/coding-agents", + "addedAt": "2026-09-28" + }, { "label": "Run agents from a harness", "to": "harness/subagents", diff --git a/docs/harness/coding-agent.md b/docs/harness/coding-agent.md index 8a45100343..3e1c2663ec 100644 --- a/docs/harness/coding-agent.md +++ b/docs/harness/coding-agent.md @@ -90,6 +90,10 @@ A trailing `*` matches every tool that starts with the text. The last matching r The workspace tools run on your machine with your permissions. Run code you do not trust in a sandbox. +## Hand work to Claude Code or Codex + +Your agent can also give tasks to coding agents you already use. [Delegate to coding agents](./coding-agents) shows how. + ## What you have now - A terminal coding agent with file tools, approvals, modes, a todo list, and a model picker. diff --git a/docs/harness/coding-agents.md b/docs/harness/coding-agents.md new file mode 100644 index 0000000000..7f61e47424 --- /dev/null +++ b/docs/harness/coding-agents.md @@ -0,0 +1,122 @@ +--- +title: Delegate to coding agents +id: harness-coding-agents +order: 10 +description: "Let a harness hand coding work to Claude Code, Codex, Grok Build, or any ACP agent. Each one works in a sandbox and keeps its own session." +keywords: + - tanstack ai + - harness + - claude code + - codex + - grok build + - sandbox + - subagents +--- + +Your lead agent plans the work, but you want Claude Code or Codex to do the edits, because they are good at code and you already use them. `codingAgents` gives the lead model one tool per coding agent. Each agent works in a sandbox, keeps its own session between calls, and its tool calls show up in your UI. + +## 1. Add the plugin + + + +react: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex +vue: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex +solid: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex +svelte: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex +preact: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex +angular: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex +octane: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex +vanilla: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex + + + +```ts group=harness-coding-agents +import { defineHarness } from '@tanstack/ai-harness' +import { permissions } from '@tanstack/ai-harness/plugins' +import { claudeCodeText } from '@tanstack/ai-claude-code' +import { codexText } from '@tanstack/ai-codex' +import { openaiText } from '@tanstack/ai-openai' +import { defineSandbox, defineWorkspace, localSource } from '@tanstack/ai-sandbox' +import { codingAgents } from '@tanstack/ai-sandbox/harness' +import { localProcessSandbox } from '@tanstack/ai-sandbox-local-process' + +const repo = '/path/to/your/repo' + +export const lead = defineHarness({ + name: 'acme/lead', + adapter: openaiText('gpt-5.6'), + plugins: () => [ + permissions(), + codingAgents({ + sandbox: defineSandbox({ + id: 'repo', + provider: localProcessSandbox({ dir: repo }), + workspace: defineWorkspace({ source: localSource(repo) }), + }), + agents: { + claude_code: { + adapter: claudeCodeText('claude-opus-4-8', { permissionMode: 'acceptEdits' }), + description: 'Larger changes, refactors, and reviews', + }, + codex: { + adapter: codexText('gpt-5.3-codex', { sandboxMode: 'workspace-write' }), + description: 'Quick fixes and tests', + }, + }, + }), + ], +}) +``` + +The lead model now has a `claude_code` tool and a `codex` tool. Each tool takes one `task`: the whole job, in words. The agent sees only that text and the files in its sandbox. + +## 2. Ask for work + +1. Start the harness, for example with the CLI. +2. Ask the lead: `have claude_code add a test for the date parser, then have codex fix what fails`. +3. Watch the child work. The CLI shows each agent's tool calls, then one line with the start of its answer: + +```text +[agent claude_code started] +[claude_code: tool Write] +[agent claude_code finished: Added parse-date.test.ts with three cases.] +``` + +The dashboard shows the same work in a block under the lead's message. An ACP editor shows the child's tool calls next to the lead's own. + +## Sessions + +Each agent keeps its own session per harness thread. The next task for `claude_code` resumes the same Claude Code session, so it still knows the files it read. The session ids live in the plugin state, so they also survive a restart of the host. + +To start over, run `/fresh claude_code`, or `/fresh` for every agent. + +## Workspaces + +`workspace` chooses where the agents work: + +| Value | Where each agent works | At the same time | +| --- | --- | --- | +| `'shared'` (default) | One sandbox per harness thread | One agent at a time | +| `'per-agent'` | A sandbox per agent | Yes | + +Use `'shared'` when the agents build on each other's changes. Use `'per-agent'` when they work on separate copies and you merge the results. + +## Plan mode + +When the `permissions()` plugin is in `plan` mode, the agents start read-only: + +- Claude Code gets `permissionMode: 'plan'`. +- Codex gets `sandboxMode: 'read-only'`. +- Any other agent gets its `planModelOptions`, for example `{ permissionMode: 'default' }` for an ACP agent that asks before each edit. + +## Other agents and sandboxes + +- `adapter` takes any coding-agent adapter: `grokBuildText` from `@tanstack/ai-grok-build`, or `acpCompatibleText` from `@tanstack/ai-acp` for any ACP agent. +- `sandbox` takes any sandbox provider. [Sandbox providers](../sandbox/providers) lists them. +- `modelOptions` on an agent is added to every call, for example a fixed `permissionMode`. + +## What you have now + +- A lead model that hands tasks to Claude Code and Codex. +- One sandbox per thread, and one saved session per agent. +- Child tool calls in the CLI, the dashboard, and ACP editors. diff --git a/docs/harness/dashboard.md b/docs/harness/dashboard.md index afbd511cd7..7ed93a5de4 100644 --- a/docs/harness/dashboard.md +++ b/docs/harness/dashboard.md @@ -1,7 +1,7 @@ --- title: Self-host the dashboard id: harness-dashboard -order: 12 +order: 13 description: "Watch and steer your harness sessions from a browser or a phone. Agents dial out to your dashboard server, so they need no open port." keywords: - tanstack ai diff --git a/docs/harness/deploy.md b/docs/harness/deploy.md index ea6cf77423..04f94332f8 100644 --- a/docs/harness/deploy.md +++ b/docs/harness/deploy.md @@ -1,7 +1,7 @@ --- title: Deploy a harness id: harness-deploy -order: 11 +order: 12 description: "Run a harness in your server, as a worker process, on another machine, or as a single executable." keywords: - tanstack ai diff --git a/docs/harness/plugins.md b/docs/harness/plugins.md index 4e01346a91..f05cd1d3bc 100644 --- a/docs/harness/plugins.md +++ b/docs/harness/plugins.md @@ -33,6 +33,7 @@ Add it with `plugins: () => [today]` in `defineHarness`. `setup` runs once per s - `middleware`: chat middleware, the same type as `chat({ middleware })`. - `generationMiddleware`: middleware for the activities agents call. - `agents`: agents added to `session.agents`. +- `subagents`: agents the model can call as tools. They are also added to `session.agents`. [Delegate to coding agents](./coding-agents) uses them. - `commands`: user actions, see below. - `config`: session settings, see below. - `contribute`: items for another plugin's extension point. diff --git a/docs/harness/subagents.md b/docs/harness/subagents.md index 78de061c97..9e17ba8306 100644 --- a/docs/harness/subagents.md +++ b/docs/harness/subagents.md @@ -1,7 +1,7 @@ --- title: Run agents from a harness id: harness-subagents -order: 10 +order: 11 description: "Start typed agents from commands and plugins, run them in groups, call a whole harness as a child, and keep the tree within limits." keywords: - tanstack ai diff --git a/examples/harness-cli/.env.example b/examples/harness-cli/.env.example index bb89717ef3..38405cac10 100644 --- a/examples/harness-cli/.env.example +++ b/examples/harness-cli/.env.example @@ -5,3 +5,7 @@ OPENAI_API_KEY= ANTHROPIC_API_KEY= # Optional: make videos with Grok Imagine instead of OpenAI Sora. XAI_API_KEY= +# Optional: hand coding work to Claude Code and Codex (uses claude login and codex login). +CODING_AGENTS= +CODEX_MODEL= +CODEX_SANDBOX_MODE= diff --git a/examples/harness-cli/.gitignore b/examples/harness-cli/.gitignore index ec9309579e..d3d132daa7 100644 --- a/examples/harness-cli/.gitignore +++ b/examples/harness-cli/.gitignore @@ -1,2 +1,4 @@ .env playground/media/ +# Marker the sandbox writes into the workspace. +playground/.tanstack-projected-* diff --git a/examples/harness-cli/README.md b/examples/harness-cli/README.md index b87e8aa82f..65345557bb 100644 --- a/examples/harness-cli/README.md +++ b/examples/harness-cli/README.md @@ -29,6 +29,16 @@ Try these: - Images use `OPENAI_API_KEY`. Videos use Grok Imagine when `XAI_API_KEY` is set, and OpenAI Sora when it is not. - Code mode is on: read-only tools (file reads, read-only Notion and Linear tools) are `external_*` functions in one `execute_typescript` program, which runs in a QuickJS isolate. Ask: `in one program, list my Linear issues and search Notion for them`. +## Hand work to Claude Code and Codex + +1. Sign in to the CLIs once: `claude login` and `codex login`. +2. Start with `CODING_AGENTS=1`. If you use Codex with a ChatGPT login, also set `CODEX_MODEL` to the model in `~/.codex/config.toml`. +3. Ask: `have claude_code create notes.md with one line, then have codex add a second line`. + +- Both agents work in `./playground` with your own logins. The API keys are removed from their processes. +- The CLI shows each agent's tool calls and a finish line. `/fresh` starts new agent sessions. +- On Windows, the Codex sandbox can block the folder (Access is denied). Then set `CODEX_SANDBOX_MODE=danger-full-access`, only for a folder you trust. + ## Other modes - One prompt for scripts and CI: `pnpm --filter harness-cli-example start -p "list the files"` diff --git a/examples/harness-cli/package.json b/examples/harness-cli/package.json index a9d609b2b8..f85a495f63 100644 --- a/examples/harness-cli/package.json +++ b/examples/harness-cli/package.json @@ -12,7 +12,9 @@ "@tanstack/ai": "workspace:*", "@tanstack/ai-acp": "workspace:*", "@tanstack/ai-anthropic": "workspace:*", + "@tanstack/ai-claude-code": "workspace:*", "@tanstack/ai-code-mode": "workspace:*", + "@tanstack/ai-codex": "workspace:*", "@tanstack/ai-dashboard": "workspace:*", "@tanstack/ai-grok": "workspace:*", "@tanstack/ai-harness": "workspace:*", @@ -21,6 +23,8 @@ "@tanstack/ai-mcp": "workspace:*", "@tanstack/ai-openai": "workspace:*", "@tanstack/ai-persistence": "workspace:*", + "@tanstack/ai-sandbox": "workspace:*", + "@tanstack/ai-sandbox-local-process": "workspace:*", "zod": "^4.2.0" }, "devDependencies": { diff --git a/examples/harness-cli/src/harness.ts b/examples/harness-cli/src/harness.ts index bc0563abda..ce06031439 100644 --- a/examples/harness-cli/src/harness.ts +++ b/examples/harness-cli/src/harness.ts @@ -11,11 +11,20 @@ import { workspaceTools, } from '@tanstack/ai-harness/plugins' import { anthropicText } from '@tanstack/ai-anthropic' +import { claudeCodeText } from '@tanstack/ai-claude-code' import { codeMode } from '@tanstack/ai-code-mode/harness' +import { codexText } from '@tanstack/ai-codex' import { mcpConnector } from '@tanstack/ai-mcp/connector' import { grokVideo } from '@tanstack/ai-grok' import { createQuickJSIsolateDriver } from '@tanstack/ai-isolate-quickjs' import { openaiText, openaiVideo } from '@tanstack/ai-openai' +import { + defineSandbox, + defineWorkspace, + localSource, +} from '@tanstack/ai-sandbox' +import { codingAgents } from '@tanstack/ai-sandbox/harness' +import { localProcessSandbox } from '@tanstack/ai-sandbox-local-process' import { z } from 'zod' import { imageAgent, videoAgent } from './media' import type { AnyTextAdapter } from '@tanstack/ai' @@ -149,6 +158,51 @@ const linear = mcpConnector({ url: 'https://mcp.linear.app/mcp', }) +// With CODING_AGENTS=1, the lead model can hand coding work to Claude Code and +// Codex. They work in ./playground with your own `claude login` and +// `codex login`, so the API keys are removed from their processes. +const coding = + process.env.CODING_AGENTS === '1' + ? [ + codingAgents({ + sandbox: defineSandbox({ + id: 'playground', + provider: localProcessSandbox({ + dir: root, + scrubEnv: ['ANTHROPIC_API_KEY', 'OPENAI_API_KEY'], + }), + workspace: defineWorkspace({ source: localSource(root) }), + }), + agents: { + claude_code: { + adapter: claudeCodeText('claude-opus-4-8', { + authMode: 'host', + permissionMode: 'acceptEdits', + }), + description: + 'Claude Code. Larger changes, refactors, and reviews in ./playground.', + }, + codex: { + // A ChatGPT login supports only some models. Set CODEX_MODEL to the + // model in ~/.codex/config.toml. + adapter: codexText(process.env.CODEX_MODEL || 'gpt-5.3-codex', { + authMode: 'host', + // On Windows, the Codex sandbox can block the folder (Access is + // denied). Then set CODEX_SANDBOX_MODE=danger-full-access, only for + // a folder you trust. + sandboxMode: + process.env.CODEX_SANDBOX_MODE === 'danger-full-access' + ? 'danger-full-access' + : 'workspace-write', + approvalPolicy: 'never', + }), + description: 'Codex. Quick fixes and tests in ./playground.', + }, + }, + }), + ] + : [] + export const assistant = defineHarness({ name: 'example/coder', description: 'A small coding agent that works in ./playground', @@ -174,5 +228,6 @@ export const assistant = defineHarness({ // The program runs in a QuickJS isolate. Any @tanstack/ai-isolate-* driver // works here. codeMode({ driver: createQuickJSIsolateDriver() }), + ...coding, ], }) diff --git a/packages/ai-acp/src/agent/index.ts b/packages/ai-acp/src/agent/index.ts index a9d3fc1af0..c36cb9ac8e 100644 --- a/packages/ai-acp/src/agent/index.ts +++ b/packages/ai-acp/src/agent/index.ts @@ -40,9 +40,11 @@ function promptText( /** One AG-UI chunk as an ACP session update, or `undefined` to skip it. */ export function toSessionUpdate(chunk: StreamChunk): SessionUpdate | undefined { - // Child agent work stays inside the harness. ACP shows the main turn. - if ('subagentRunId' in chunk && chunk.subagentRunId) return undefined + // A child agent's text and thoughts stay out of the main message. Its tool + // calls (the edits and commands it runs) show up like the lead's own. + const fromChild = 'subagentRunId' in chunk && Boolean(chunk.subagentRunId) if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) { + if (fromChild) return undefined return { sessionUpdate: 'agent_message_chunk', messageId: chunk.messageId, @@ -50,6 +52,7 @@ export function toSessionUpdate(chunk: StreamChunk): SessionUpdate | undefined { } } if (chunk.type === EventType.REASONING_MESSAGE_CONTENT) { + if (fromChild) return undefined return { sessionUpdate: 'agent_thought_chunk', messageId: chunk.messageId, diff --git a/packages/ai-acp/tests/child-updates.test.ts b/packages/ai-acp/tests/child-updates.test.ts new file mode 100644 index 0000000000..fd67ee7526 --- /dev/null +++ b/packages/ai-acp/tests/child-updates.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, it } from 'vitest' +import { EventType } from '@tanstack/ai' +import { toSessionUpdate } from '../src/agent' +import type { StreamChunk } from '@tanstack/ai' + +const child = (event: Record): StreamChunk => + ({ subagentRunId: 'child-1', timestamp: 1, ...event }) as StreamChunk + +describe('toSessionUpdate for child agents', () => { + it('shows child tool calls like the lead tool calls', () => { + expect( + toSessionUpdate( + child({ + type: EventType.TOOL_CALL_START, + toolCallId: 'edit-1', + toolCallName: 'Edit', + }), + ), + ).toEqual({ + sessionUpdate: 'tool_call_update', + toolCallId: 'edit-1', + name: 'Edit', + title: 'Edit', + status: 'in_progress', + }) + expect( + toSessionUpdate( + child({ + type: EventType.TOOL_CALL_RESULT, + messageId: 'm', + toolCallId: 'edit-1', + content: 'ok', + }), + ), + ).toEqual({ + sessionUpdate: 'tool_call_update', + toolCallId: 'edit-1', + status: 'completed', + }) + }) + + it('keeps child text and thoughts out of the main message', () => { + expect( + toSessionUpdate( + child({ + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm', + delta: 'hi', + }), + ), + ).toBeUndefined() + expect( + toSessionUpdate( + child({ + type: EventType.REASONING_MESSAGE_CONTENT, + messageId: 'r', + delta: 'hmm', + }), + ), + ).toBeUndefined() + }) +}) diff --git a/packages/ai-dashboard/src/ui.ts b/packages/ai-dashboard/src/ui.ts index 4f0b2fd6ae..76d2e0c840 100644 --- a/packages/ai-dashboard/src/ui.ts +++ b/packages/ai-dashboard/src/ui.ts @@ -89,7 +89,7 @@ function open(session) { function apply(frame) { if (frame.type !== 'harness.event' || !current) return const event = frame.event - if (event.subagentRunId) return + if (event.subagentRunId) { applyChild(event); return } if (event.type === 'TEXT_MESSAGE_CONTENT') { const last = current.messages[current.messages.length - 1] if (last && last.kind === 'assistant' && last.op === frame.operationId) last.text += event.delta @@ -109,6 +109,32 @@ function apply(frame) { } } +// A child agent shows as one block under the message it came from: its name, +// status, tool calls, and text. +function applyChild(event) { + const children = current.children || (current.children = new Map()) + if (event.type === 'SUBAGENT_STARTED') { + const node = { kind: 'agent', name: event.name || 'agent', status: 'working', tools: [], text: '' } + children.set(event.subagentRunId, node) + current.messages.push(node) + return + } + const node = children.get(event.subagentRunId) + if (!node) return + if (event.type === 'TEXT_MESSAGE_CONTENT') node.text += event.delta + else if (event.type === 'TOOL_CALL_START') node.tools.push(event.toolCallName) + else if (event.type === 'SUBAGENT_FINISHED') node.status = 'done' + else if (event.type === 'SUBAGENT_ERROR') { node.status = 'failed'; node.text += (node.text ? '\\n' : '') + event.message } +} + +function messageNode(message) { + if (message.kind !== 'agent') return $('div', { class: 'msg ' + message.kind }, message.text) + return $('div', { class: 'msg agent ' + message.status }, + $('strong', {}, 'agent ' + message.name + ' (' + message.status + ')'), + ...(message.tools.length ? [$('div', { class: 'tools' }, message.tools.map((tool) => 'tool ' + tool).join(', '))] : []), + ...(message.text ? [$('div', {}, message.text)] : [])) +} + async function send(input) { if (input.op === 'prompt' || input.op === 'steer') current.messages.push({ kind: 'user', text: input.message }) const receipt = await api(current.path + '/input', { method: 'POST', body: JSON.stringify({ input }) }) @@ -119,7 +145,7 @@ async function send(input) { function draw() { const main = document.getElementById('main') if (!current) { main.replaceChildren($('p', { class: 'muted' }, 'Pick a session.')); return } - const log = $('div', { class: 'log' }, ...current.messages.map((message) => $('div', { class: 'msg ' + message.kind }, message.text))) + const log = $('div', { class: 'log' }, ...current.messages.map(messageNode)) const actions = [] if (current.interrupts.length) { const decide = (approved) => { const resume = current.interrupts.map((interrupt) => ({ interruptId: interrupt.id, status: 'resolved', payload: approved })); current.interrupts = []; send({ op: 'resolve', resume }) } @@ -183,6 +209,7 @@ export const DASHBOARD_HTML = ` .msg { white-space: pre-wrap; padding: 10px 12px; border-radius: 10px; background: var(--panel); max-width: 80ch; } .msg.user { align-self: flex-end; background: #1d3557; } .msg.tool, .msg.notice { color: var(--muted); background: transparent; padding: 2px 12px; } .msg.error { border: 1px solid #e5484d; } + .msg.agent { border-left: 3px solid #7c5cff; display: flex; flex-direction: column; gap: 4px; } .msg.agent.failed { border-left-color: #e5484d; } .msg.agent .tools { color: var(--muted); font-size: 0.9em; } code { background: var(--panel); padding: 2px 6px; border-radius: 6px; } @media (max-width: 720px) { #app { grid-template-columns: 1fr; } aside { border-right: 0; border-bottom: 1px solid var(--line); } } diff --git a/packages/ai-dashboard/tests/child-agents.test.ts b/packages/ai-dashboard/tests/child-agents.test.ts new file mode 100644 index 0000000000..c5fcb36485 --- /dev/null +++ b/packages/ai-dashboard/tests/child-agents.test.ts @@ -0,0 +1,143 @@ +import { runInContext, createContext } from 'node:vm' +import { describe, expect, it } from 'vitest' +import { DASHBOARD_HTML } from '../src/ui' + +/** Just enough DOM for the page script: nodes with children and text. */ +class FakeNode { + children: Array = [] + className = '' + attributes: Record = {} + constructor( + readonly tag: string, + readonly text = '', + ) {} + append(...children: Array) { + this.children.push(...children) + } + replaceChildren(...children: Array) { + this.children = children + } + setAttribute(key: string, value: string) { + this.attributes[key] = value + } + get textContent(): string { + return this.text + this.children.map((child) => child.textContent).join('') + } +} + +function loadPage() { + const script = /' + const record = await upload('user-1-a', page, 'text/html') + const { path } = await signedPath(record) + + const response = await fetchSigned(path) + + expect(Object.fromEntries(response.headers)).toEqual({ + 'accept-ranges': 'bytes', + 'content-length': String(page.length), + 'content-security-policy': "default-src 'none'; sandbox", + 'content-type': 'text/html', + 'x-content-type-options': 'nosniff', + }) + await host.close() + }) +}) + +describe('POST run with media parts', () => { + // AG-UI text parts carry `text`. Media parts have the TanStack shape. + const question = { type: 'text', text: 'What is this?' } + // 'aGVsbG8=' is base64 for "hello", the bytes of every upload here. + const data: ContentPart = { + type: 'image', + source: { type: 'data', value: 'aGVsbG8=', mimeType: 'image/png' }, + } + + const cases: Array<{ + label: string + part: (record: MediaRecord) => ContentPart + }> = [ + { label: 'a harness-media URL, and keeps the URL', part: mediaPart }, + { label: 'a data part, and keeps it as sent', part: () => data }, + ] + + it.each(cases)( + 'gives the model the bytes of $label in the transcript', + async ({ part }) => { + const { host, call, upload, calls } = setup() + const record = await upload('user-1-run') + + const response = await call('run', { + method: 'POST', + body: runBody('user-1-run', [question, part(record)]), + headers: { 'content-type': 'application/json' }, + }) + await response.text() + const transcript = await ( + await call('transcript?threadId=user-1-run') + ).json() + + expect(calls[0].messages[0].content).toEqual([ + { type: 'text', content: 'What is this?' }, + data, + ]) + expect(transcript[0].content).toEqual([ + { type: 'text', content: 'What is this?' }, + part(record), + ]) + await host.close() + }, + ) +}) diff --git a/packages/ai-harness/tests/resume-edges.test.ts b/packages/ai-harness/tests/resume-edges.test.ts index 7c9367ec33..685de95ffa 100644 --- a/packages/ai-harness/tests/resume-edges.test.ts +++ b/packages/ai-harness/tests/resume-edges.test.ts @@ -8,7 +8,7 @@ import { findCrashedRuns, repairTranscript, } from '../src/resume' -import { messageTexts, mockAdapter, text, untilAborted } from './helpers' +import { mockAdapter, text, untilAborted } from './helpers' import type { RunRecord, StreamChunk } from '@tanstack/ai' import type { HarnessPersistence } from '../src/host' @@ -162,7 +162,7 @@ describe('crash recovery edges', () => { }) describe('harnessText edges', () => { - it('uses a memory host by default and reads text content parts', async () => { + it('uses a memory host by default and keeps the content parts', async () => { const inner = mockAdapter([() => text('parts answer'), () => text('empty')]) const studio = defineHarness({ name: 'test/parts', adapter: inner.adapter }) const model = harnessText(studio) @@ -173,22 +173,15 @@ describe('harnessText edges', () => { } return out } - expect( - await collect([ - { - role: 'user', - content: [ - { type: 'text', content: 'one ' }, - { - type: 'image', - source: { type: 'url', value: 'https://x/y.png' }, - }, - { type: 'text', content: 'two' }, - ], - }, - ]), - ).toBe('parts answer') - expect(messageTexts(inner.calls[0])).toEqual(['one two']) + const parts = [ + { type: 'text', content: 'one ' }, + { type: 'image', source: { type: 'url', value: 'https://x/y.png' } }, + { type: 'text', content: 'two' }, + ] + expect(await collect([{ role: 'user', content: parts }])).toBe( + 'parts answer', + ) + expect(inner.calls[0].messages[0].content).toEqual(parts) await expect(model.structuredOutput({} as never)).rejects.toThrow( 'does not support structured output', ) From 116f1def149c8167ae79a8acaefdc4b6a1cc75f4 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 29 Sep 2026 12:32:33 +0200 Subject: [PATCH 061/286] feat(ai-mcp): harness MCP media (attachments, media in results, media links), and resource templates that read their URI createHarnessMcpServer takes filePaths for path attachments, the chat tool takes path, url, and data attachments, and results carry the media a turn made: small images and audio inline, other files as harness-media:// links. resourceDefinition({ uriTemplate, argsSchema }) now gives read the parsed template variables and the URI, and a read result can set its own mimeType. Claude-Session: https://claude.ai/code/session_01APYv1qshKyjPPpkFyRZhfZ --- .changeset/mcp-resource-template-args.md | 5 + docs/config.json | 3 +- docs/mcp/server-content.md | 12 +- packages/ai-mcp/src/direct-client.ts | 10 +- packages/ai-mcp/src/harness.ts | 377 +++++++++++++- packages/ai-mcp/src/server/create-server.ts | 31 +- packages/ai-mcp/src/server/definitions.ts | 41 +- packages/ai-mcp/tests/harness-media.test.ts | 479 ++++++++++++++++++ .../ai-mcp/tests/server/create-server.test.ts | 27 + .../ai-mcp/tests/server/definitions.test.ts | 13 + 10 files changed, 974 insertions(+), 24 deletions(-) create mode 100644 .changeset/mcp-resource-template-args.md create mode 100644 packages/ai-mcp/tests/harness-media.test.ts diff --git a/.changeset/mcp-resource-template-args.md b/.changeset/mcp-resource-template-args.md new file mode 100644 index 0000000000..b5126989b0 --- /dev/null +++ b/.changeset/mcp-resource-template-args.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-mcp': minor +--- + +A resource with a `uriTemplate` can read the URI the host asked for. `resourceDefinition({ uriTemplate, argsSchema })` parses the template variables, and `read` gets them and the URI. A read result `{ text | blob, mimeType }` sets the MIME type of that answer. diff --git a/docs/config.json b/docs/config.json index d74f1aeba1..b43bb68369 100644 --- a/docs/config.json +++ b/docs/config.json @@ -530,7 +530,8 @@ { "label": "Serve Resources and Prompts", "to": "mcp/server-content", - "addedAt": "2026-09-22" + "addedAt": "2026-09-22", + "updatedAt": "2026-09-29" }, { "label": "Ask for Input", diff --git a/docs/mcp/server-content.md b/docs/mcp/server-content.md index 13d84c56f9..bf8af07df4 100644 --- a/docs/mcp/server-content.md +++ b/docs/mcp/server-content.md @@ -41,7 +41,8 @@ const file = resourceDefinition({ name: 'file', mimeType: 'text/plain', uriTemplate: 'file:///{path}', -}).read(async () => ({ text: 'file body' })) + argsSchema: z.object({ path: z.string() }), +}).read(async ({ path }) => ({ text: `The body of ${path}` })) const summarize = promptDefinition({ name: 'summarize', @@ -72,7 +73,12 @@ If the resource has no `uri` and no `uriTemplate`, `resourceDefinition` throws ` If you pass `uri` and `uriTemplate`, the server uses `uri`. -`read` takes no arguments. `read` returns `{ text }` for a text document. For a binary document, `read` returns `{ blob }` with a base64 string. +`read` returns `{ text }` for a text document. For a binary document, `read` returns `{ blob }` with a base64 string. Add `mimeType` to that object when one template serves files of different types. + +For a `uriTemplate`, `read` gets the variables of the URI the host asked for, and the URI itself: + +- `argsSchema.parse` runs first, so `read` gets the parsed variables. For `file:///notes.md`, `path` is `notes.md`. +- Without `argsSchema`, `read` gets the variables as strings. ## Prompts @@ -89,4 +95,6 @@ If `role` is not `user` or `assistant`, the server sends that message as `user`. The host reads `file:///readme.md`. The `text` is `# Hello`. +The host reads `file:///notes.md`. The `text` is `The body of notes.md`. + The host starts from the `summarize` prompt with topic `weather`. The `role` is `user`. The `content` string is `weather`. diff --git a/packages/ai-mcp/src/direct-client.ts b/packages/ai-mcp/src/direct-client.ts index 7719839a08..cb0233bbf2 100644 --- a/packages/ai-mcp/src/direct-client.ts +++ b/packages/ai-mcp/src/direct-client.ts @@ -113,7 +113,7 @@ type ListedTool = { type ListedResource = { uri?: string - read: () => unknown + read: (variables: Record>, uri: URL) => unknown } type ListedPrompt = { @@ -179,10 +179,14 @@ export function directMCPClient( if (resource === undefined) { throw new Error(`The MCP server has no resource ${uri}.`) } - const read = resource.read as () => + const read = resource.read as ( + variables: Record>, + uri: URL, + ) => | ResourceContents> | Promise>> - return read() + // A fixed uri has no template variables. + return read({}, new URL(uri)) }, async getPrompt>( diff --git a/packages/ai-mcp/src/harness.ts b/packages/ai-mcp/src/harness.ts index 11a330212d..eacf38174c 100644 --- a/packages/ai-mcp/src/harness.ts +++ b/packages/ai-mcp/src/harness.ts @@ -4,18 +4,26 @@ import { readUnopenedInterruptBinding, toolDefinition, } from '@tanstack/ai' -import { HARNESS_EVENTS } from '@tanstack/ai-harness' +import { + HARNESS_EVENTS, + isMediaRecord, + kindOf, + mediaPart, +} from '@tanstack/ai-harness' // The inferred return type names `MCPHandleOptions`. Without an import of // `./server/index`, the .d.ts emit writes `./server.js`, which does not resolve. -import { createMCPServer } from './server/index' -import type { Interrupt, JSONSchema } from '@tanstack/ai' +import { createMCPServer, resourceDefinition } from './server/index' +import { toCallToolResult } from './server/tasks' +import type { ContentPart, Interrupt, JSONSchema } from '@tanstack/ai' import type { AnyHarness, Cursor, HarnessHost, HarnessSession, + MediaRecord, Operation, } from '@tanstack/ai-harness' +import type { ContentBlock } from '@modelcontextprotocol/server' import type { MCPToolContext } from './server/context' /** Options for {@link createHarnessMcpServer}. */ @@ -38,6 +46,15 @@ export interface HarnessMcpServerOptions { * client answers them all with one `resolve`. */ approvals?: 'ask' | 'auto' + /** + * The folders that a `path` attachment of `chat` can read from. A path must + * resolve, after every symlink, to a file inside one of them. Leave it out, + * and every `path` attachment is refused. + * + * The stdio CLI (`--mcp`) passes the working folder. An HTTP server should + * pass none, because a remote client must not read files on the server. + */ + filePaths?: ReadonlyArray /** The MCP server name. Default: the harness name. */ name?: string /** The MCP server version. Default `'1.0.0'`. */ @@ -50,12 +67,101 @@ type Decision = { interruptId: string; approved?: boolean; payload?: unknown } /** The answer to one interrupt, for `session.resolve`. */ type Answer = { interruptId: string; payload: unknown } +/** One file of a `chat` call, from the client. */ +type Attachment = + | { path: string; name?: string } + | { url: string; mimeType?: string; name?: string } + | { data: string; mimeType: string; name?: string } + +/** The URI of a media resource is `harness-media:///`. */ +const MEDIA_URI_PREFIX = 'harness-media://' + +/** The biggest image or audio file that a result carries inline. */ +const INLINE_MEDIA_MAX_BYTES = 5 * 1024 * 1024 + +// ponytail: common types only. Send any other file as data with its mimeType. +const mimeTypes = new Map([ + ['png', 'image/png'], + ['jpg', 'image/jpeg'], + ['jpeg', 'image/jpeg'], + ['gif', 'image/gif'], + ['webp', 'image/webp'], + ['mp3', 'audio/mpeg'], + ['wav', 'audio/wav'], + ['ogg', 'audio/ogg'], + ['m4a', 'audio/mp4'], + ['flac', 'audio/flac'], + ['mp4', 'video/mp4'], + ['webm', 'video/webm'], + ['mov', 'video/quicktime'], + ['pdf', 'application/pdf'], + ['txt', 'text/plain'], + ['md', 'text/markdown'], + ['csv', 'text/csv'], + ['html', 'text/html'], +]) + const threadIdSchema: JSONSchema = { type: 'string', description: 'The conversation id. Leave it out to use the default conversation.', } +const nameSchema: JSONSchema = { + type: 'string', + description: 'The file name. Optional.', +} + +const attachmentsSchema: JSONSchema = { + type: 'array', + description: 'Files to send with the message.', + items: { + anyOf: [ + { + type: 'object', + properties: { + path: { + type: 'string', + description: + 'A file path on the machine of the server. The file must be in a folder that the server allows.', + }, + name: nameSchema, + }, + required: ['path'], + }, + { + type: 'object', + properties: { + url: { + type: 'string', + description: + 'A URL of the file. The model reads it. The server does not download it.', + }, + mimeType: { + type: 'string', + description: + 'The MIME type, for example image/png. Needed when the URL does not end in a known file extension.', + }, + name: nameSchema, + }, + required: ['url'], + }, + { + type: 'object', + properties: { + data: { type: 'string', description: 'The file bytes in base64.' }, + mimeType: { + type: 'string', + description: 'The MIME type, for example image/png.', + }, + name: nameSchema, + }, + required: ['data', 'mimeType'], + }, + ], + }, +} + /** * Serves a harness as an MCP server, so any MCP client (Claude Code, Claude * Desktop, Cursor, another agent) can use it. @@ -70,10 +176,17 @@ const threadIdSchema: JSONSchema = { * `generic`. `approve` and `reject` answer approvals only. `resolve` answers * every kind: `approved` for an approval, and `payload` for the others. * + * `chat` takes `attachments`: `{ path }` (only inside `filePaths`), + * `{ url, mimeType? }` (the model reads the URL, the server does not fetch + * it), or `{ data, mimeType }` (base64). A result lists the media that the + * work made in `media`. An image or audio file up to 5 MB comes back inline. + * Any other file comes back as a `resource_link` to + * `harness-media:///`, which `resources/read` returns. + * * The result is the server from `createMCPServer`. Mount `server.fetch` on * an HTTP route, or pass the server to `serveMCPStdio`. * - * @param options - The host, the harness, the default thread, and the approval mode + * @param options - The host, the harness, the default thread, the approval mode, and the folders for `path` attachments * * @example * ```ts @@ -85,6 +198,7 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { const { host, harness } = options const defaultThread = options.threadId ?? 'main' const approvals = options.approvals ?? 'ask' + const filePaths = options.filePaths ?? [] // Work that stopped for a question, by question id. `answer` continues it. const waiting = new Map>() @@ -124,7 +238,7 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { } // Waits for the work. A chat turn that stops for approvals continues while - // the approval mode decides them. + // the approval mode decides them. The result gets the media of the work. async function finish( session: HarnessSession, work: Operation, @@ -133,23 +247,29 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { ) { let operation = work let from = workFrom + // The media of every operation that this call waited for. + const media: Array = [] while (true) { const outcome = await settle(session, operation, from) if (!outcome.done) { waiting.set(outcome.questionId, operation) return { status: 'waiting', ...pending(session) } } - if (operation.kind !== 'chat') return outcome.value + media.push(...(await mediaOf(operation))) + if (operation.kind !== 'chat') { + return withMedia(session, outcome.value, media) + } const interrupts = session.snapshot().pendingInterrupts const stopped = operation.status() === 'interrupted' && interrupts.length > 0 const answers = stopped ? await decide(interrupts, context) : undefined if (answers === undefined) { - return { + const result = { status: operation.status(), text: turnText(outcome.value), ...pending(session), } + return withMedia(session, result, media) } from = session.snapshot().cursor operation = (await resolveAll(session, answers)).turn @@ -188,18 +308,20 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { const chat = toolDefinition({ name: 'chat', description: - 'Send a message to the harness and wait for its answer. While a turn runs, the message waits in the queue. The result has the answer text, the status, and the interrupts and questions that wait for you.', + 'Send a message to the harness and wait for its answer. While a turn runs, the message waits in the queue. The result has the answer text, the status, the interrupts and questions that wait for you, and the media that the turn made.', inputSchema: withThreadId({ type: 'object', properties: { message: { type: 'string', description: 'The message.' }, + attachments: attachmentsSchema, }, required: ['message'], }), }).server(async (args, ctx) => { const session = await open(args) + const message = await userInputOf(session, args, filePaths) const from = session.snapshot().cursor - const turn = session.prompt(textArg(args, 'message'), { busy: 'queue' }) + const turn = session.prompt(message, { busy: 'queue' }) return finish(session, turn, from, ctx.context) }) @@ -373,9 +495,27 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { }), ) + // The bytes of a `resource_link` from a result. The link names its thread, + // and a thread reads only its own media. + const mediaResource = resourceDefinition({ + uriTemplate: `${MEDIA_URI_PREFIX}{threadId}/{id}`, + name: 'media', + mimeType: 'application/octet-stream', + argsSchema: { parse: mediaAddress }, + }).read(async ({ threadId, id }) => { + const target = await host.open(harness, { threadId }) + const record = await target.getMedia(id) + if (record === null) { + throw new Error(`Media ${id} was not found in thread ${threadId}.`) + } + const blob = toBase64(await target.loadMedia(id)) + return { blob, mimeType: record.mimeType } + }) + return createMCPServer({ name: options.name ?? harness.name, version: options.version ?? '1.0.0', + resources: [mediaResource], tools: [ chat, steer, @@ -471,6 +611,214 @@ function pending(session: HarnessSession) { } } +/** + * The message of a `chat` call: the text alone, or the text and one part per + * attachment. A file goes into the media store of the thread first. A + * `MediaError` (too big, a type the harness does not take) stops the call. + */ +async function userInputOf( + session: HarnessSession, + args: unknown, + filePaths: ReadonlyArray, +) { + const message = textArg(args, 'message') + const attachments = attachmentsOf(args) + if (attachments.length === 0) return message + const parts: Array = [{ type: 'text', content: message }] + for (const attachment of attachments) { + parts.push(await attachmentPart(session, attachment, filePaths)) + } + return parts +} + +function attachmentsOf(args: unknown) { + const list = + isRecord(args) && Array.isArray(args.attachments) ? args.attachments : [] + const attachments = list.filter(isAttachment) + if (attachments.length !== list.length) { + throw new Error('Each attachment needs path, url, or data with mimeType.') + } + return attachments +} + +async function attachmentPart( + session: HarnessSession, + attachment: Attachment, + filePaths: ReadonlyArray, +) { + if ('path' in attachment) { + const file = await readAllowedFile(attachment.path, filePaths) + const name = attachment.name ?? file.name + return mediaPart( + await session.putMedia(file.bytes, { mimeType: file.mimeType, name }), + ) + } + if ('url' in attachment) return urlPart(attachment) + const bytes = Uint8Array.from(atob(attachment.data), (char) => + char.charCodeAt(0), + ) + const name = attachment.name ?? 'attachment' + return mediaPart( + await session.putMedia(bytes, { mimeType: attachment.mimeType, name }), + ) +} + +/** + * The bytes of `path` when its real path is inside one of `folders`. The real + * path follows every symlink and `..`, so neither can leave a folder. A + * missing file gets the same refusal, so a client cannot probe for files. + */ +async function readAllowedFile(path: string, folders: ReadonlyArray) { + if (folders.length === 0) { + throw new Error( + 'This server does not read files. Send the file as data or url.', + ) + } + // Loaded here, not at the top: the module must also load where node:fs is + // missing (an edge worker), and only a path attachment needs it. + const { readFile, realpath } = await import('node:fs/promises') + const { basename, isAbsolute, relative, sep } = await import('node:path') + const refusal = `Cannot read ${path}. It must be a file inside the folders this server may read.` + const real = await realpath(path).catch(() => undefined) + if (real === undefined) throw new Error(refusal) + const roots = await Promise.all(folders.map((folder) => realpath(folder))) + const isInside = roots.some((root) => { + const rest = relative(root, real) + const isOut = rest === '..' || rest.startsWith(`..${sep}`) + return rest !== '' && !isOut && !isAbsolute(rest) + }) + if (!isInside) throw new Error(refusal) + const name = basename(path) + const mimeType = mimeTypeOf(name) + if (mimeType === undefined) { + throw new Error( + `Unknown file type: ${name}. Send the file as data with its mimeType.`, + ) + } + // ponytail: reads the whole file, and the store refuses it after that when + // it is over `media.maxBytes`. Stream it when big local files matter. + return { bytes: await readFile(real), mimeType, name } +} + +/** + * A part that sends `attachment.url` to the model as it is. The server does + * not fetch it: a remote client must not make the server call an internal + * address. The kind comes from `mimeType`, or from the extension of the path. + */ +function urlPart(attachment: { url: string; mimeType?: string }) { + const { url } = attachment + const mimeType = attachment.mimeType ?? mimeTypeOf(new URL(url).pathname) + const kind = mimeType === undefined ? undefined : kindOf(mimeType) + if (mimeType === undefined || kind === undefined) { + throw new Error(`Cannot tell the kind of ${url}. Add its mimeType.`) + } + const part: ContentPart = { + type: kind, + source: { type: 'url', value: url, mimeType }, + } + return part +} + +/** The MIME type of a file name or URL path by its extension, or `undefined`. */ +function mimeTypeOf(path: string) { + const extension = /\.([^./\\]+)$/.exec(path)?.[1] + return extension === undefined + ? undefined + : mimeTypes.get(extension.toLowerCase()) +} + +/** The media records that the settled `operation` published, in order. */ +async function mediaOf(operation: Operation) { + const records: Array = [] + // The operation is settled, so its stream replays its events, then ends. + const events = operation.stream() + for await (const event of events) { + const isMedia = + event.type === EventType.CUSTOM && event.name === HARNESS_EVENTS.media + if (isMedia && isMediaRecord(event.value)) records.push(event.value) + } + return records +} + +/** + * The tool result of `value` with its media. Without media, `value` stays as + * it is. An object result also lists the media in `media`. Each file then + * follows the JSON text as its own content block. + */ +async function withMedia( + session: HarnessSession, + value: unknown, + media: Array, +) { + if (media.length === 0) return value + const listed = media.map((record) => ({ + id: record.id, + kind: record.kind, + name: record.name, + mimeType: record.mimeType, + size: record.size, + uri: mediaUri(record), + })) + const blocks = await Promise.all( + media.map((record) => mediaBlock(session, record)), + ) + const result = toCallToolResult( + isRecord(value) ? { ...value, media: listed } : value, + ) + return { ...result, content: [...result.content, ...blocks] } +} + +/** Image and audio up to 5 MB go inline. Anything else goes as a link. */ +async function mediaBlock(session: HarnessSession, record: MediaRecord) { + const { kind } = record + const isSmall = record.size <= INLINE_MEDIA_MAX_BYTES + if (isSmall && (kind === 'image' || kind === 'audio')) { + const data = toBase64(await session.loadMedia(record.id)) + const inline: ContentBlock = { type: kind, data, mimeType: record.mimeType } + return inline + } + const link: ContentBlock = { + type: 'resource_link', + uri: mediaUri(record), + name: record.name, + mimeType: record.mimeType, + } + return link +} + +function mediaUri(record: MediaRecord) { + const threadId = encodeURIComponent(record.threadId) + return `${MEDIA_URI_PREFIX}${threadId}/${encodeURIComponent(record.id)}` +} + +/** + * The thread and media id of a `harness-media:///` URI, from + * the template variables. `mediaUri` encodes both, so they are decoded here. + */ +function mediaAddress(variables: unknown) { + const threadId = isRecord(variables) ? variables.threadId : undefined + const id = isRecord(variables) ? variables.id : undefined + if (typeof threadId !== 'string' || typeof id !== 'string') { + throw new Error('A media URI is harness-media:///.') + } + return { + threadId: decodeURIComponent(threadId), + id: decodeURIComponent(id), + } +} + +// ponytail: the same as `toBase64` in `@tanstack/ai-harness/src/media.ts`, +// which the package does not export. `btoa`, not `Buffer`, for edge workers. +function toBase64(bytes: Uint8Array) { + const chunk = 0x8000 + let binary = '' + // Chunks keep `fromCharCode` under the engine's argument limit. + for (let offset = 0; offset < bytes.length; offset += chunk) { + binary += String.fromCharCode(...bytes.subarray(offset, offset + chunk)) + } + return btoa(binary) +} + /** * How the client answers `interrupt`. `chat()` puts a binding on each * interrupt. The binding marks a tool approval (answer with `approved`) and a @@ -620,6 +968,17 @@ function isDecision(value: unknown): value is Decision { ) } +function isAttachment(value: unknown): value is Attachment { + if (!isRecord(value) || !isOptionalString(value.name)) return false + if (typeof value.path === 'string') return true + if (typeof value.url === 'string') return isOptionalString(value.mimeType) + return typeof value.data === 'string' && typeof value.mimeType === 'string' +} + +function isOptionalString(value: unknown) { + return value === undefined || typeof value === 'string' +} + // `session.describe()` gives each command input as JSON Schema. function isJsonSchema(value: unknown): value is JSONSchema { return isRecord(value) diff --git a/packages/ai-mcp/src/server/create-server.ts b/packages/ai-mcp/src/server/create-server.ts index 8d63123f62..ab64d63c4c 100644 --- a/packages/ai-mcp/src/server/create-server.ts +++ b/packages/ai-mcp/src/server/create-server.ts @@ -59,7 +59,8 @@ type McpResource = { mimeType: string uri?: string uriTemplate?: string - read: () => unknown + /** Gets the variables of a `uriTemplate` (none for a `uri`) and the URI. */ + read: (variables: Record>, uri: URL) => unknown } // A method type is bivariant. A function property is strict, so a prompt @@ -641,20 +642,38 @@ function isTextBlock(value: unknown): value is { type: 'text'; text: string } { function registerServerResource(server: McpServer, resource: McpResource) { const metadata = { mimeType: resource.mimeType } - const read = async (uri: URL) => - resourceContents(uri.href, resource.mimeType, await resource.read()) + const read = async ( + uri: URL, + variables: Record>, + ) => + resourceContents( + uri.href, + resource.mimeType, + await resource.read(variables, uri), + ) if (resource.uri !== undefined) { - server.registerResource(resource.name, resource.uri, metadata, read) + // A fixed uri has no variables. The SDK's second argument is the request. + server.registerResource(resource.name, resource.uri, metadata, (uri) => + read(uri, {}), + ) return } if (resource.uriTemplate === undefined) return const template = new ResourceTemplate(resource.uriTemplate, { list: undefined, }) - server.registerResource(resource.name, template, metadata, read) + // A template gets the variables the SDK matched in the asked-for URI. + server.registerResource(resource.name, template, metadata, (uri, variables) => + read(uri, variables), + ) } -function resourceContents(uri: string, mimeType: string, body: unknown) { +function resourceContents(uri: string, fallback: string, body: unknown) { + // A body can name its own MIME type: one template can serve many types. + const mimeType = + isRecord(body) && typeof body.mimeType === 'string' + ? body.mimeType + : fallback if (isRecord(body) && typeof body.text === 'string') { return { contents: [{ uri, mimeType, text: body.text }] } } diff --git a/packages/ai-mcp/src/server/definitions.ts b/packages/ai-mcp/src/server/definitions.ts index c24763f628..a006c2a50d 100644 --- a/packages/ai-mcp/src/server/definitions.ts +++ b/packages/ai-mcp/src/server/definitions.ts @@ -7,6 +7,15 @@ type PromptArgsSchema = { parse: (input: unknown) => TArgs } +/** The variables of a `uriTemplate`, as the MCP SDK matches them. */ +type TemplateVariables = Record> + +type ResourceArgsOf = TConfig extends { + argsSchema: PromptArgsSchema +} + ? TArgs + : TemplateVariables + /** * Builds a resource definition for the MCP server. * @@ -14,7 +23,12 @@ type PromptArgsSchema = { * If `uri` and `uriTemplate` are both missing, this function throws a TypeError. * Call `.read` with a function that returns the resource contents. * - * @param config - The resource `name`, `mimeType`, and `uri` or `uriTemplate`. + * For a `uriTemplate`, the read function gets the variables of the URI the + * client asked for, and the URI itself. Pass `argsSchema` to parse the + * variables first. A body `{ text | blob, mimeType }` sets the MIME type of + * that answer, for a template whose files have different types. + * + * @param config - The resource `name`, `mimeType`, `uri` or `uriTemplate`, and an optional `argsSchema`. * @throws {TypeError} When `uri` and `uriTemplate` are both missing. * * @example @@ -24,6 +38,13 @@ type PromptArgsSchema = { * name: 'readme', * mimeType: 'text/markdown', * }).read(async () => ({ text: '# Hello' })) + * + * const user = resourceDefinition({ + * uriTemplate: 'users://{id}', + * name: 'user', + * mimeType: 'application/json', + * argsSchema: z.object({ id: z.string() }), + * }).read(async ({ id }) => ({ text: JSON.stringify(await loadUser(id)) })) * ``` */ export function resourceDefinition< @@ -32,6 +53,7 @@ export function resourceDefinition< mimeType: string uri?: string uriTemplate?: string + argsSchema?: PromptArgsSchema }, >(config: TConfig) { const hasUri = config.uri !== undefined @@ -44,10 +66,23 @@ export function resourceDefinition< return { ...config, - read(readContents: () => TContents | Promise) { + read( + readContents: ( + args: ResourceArgsOf, + uri: URL | undefined, + ) => TContents | Promise, + ) { return { ...config, - read: readContents, + async read(variables: TemplateVariables = {}, uri?: URL) { + // Without `argsSchema`, the args are the variables as matched. + // `ResourceArgsOf` picks the same branch from the config type, which + // TypeScript cannot follow through the runtime check. + const args = ( + config.argsSchema ? config.argsSchema.parse(variables) : variables + ) as ResourceArgsOf + return readContents(args, uri) + }, } }, } diff --git a/packages/ai-mcp/tests/harness-media.test.ts b/packages/ai-mcp/tests/harness-media.test.ts new file mode 100644 index 0000000000..4212406253 --- /dev/null +++ b/packages/ai-mcp/tests/harness-media.test.ts @@ -0,0 +1,479 @@ +import { mkdir, mkdtemp, rm, symlink, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join, sep } from 'node:path' +import { + Client, + StreamableHTTPClientTransport, +} from '@modelcontextprotocol/client' +import { EventType, defineAgent } from '@tanstack/ai' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { afterEach, describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { createHarnessMcpServer } from '../src/harness' +import type { + AnyTextAdapter, + ImageAdapter, + ModelMessage, + StreamChunk, +} from '@tanstack/ai' +import type { AnyHarness } from '@tanstack/ai-harness' +import type { HarnessMcpServerOptions } from '../src/harness' + +const serverUrl = new URL('https://harness.example.com/mcp') +const now = () => Date.now() +/** The bytes "hello", and their base64. */ +const hello = new TextEncoder().encode('hello') +const helloBase64 = 'aGVsbG8=' +/** One byte over the 5 MB inline limit. */ +const bigImage = new Uint8Array(5 * 1024 * 1024 + 1) + +const mediaListShape = z.object({ + media: z.array(z.object({ id: z.string(), uri: z.string() })), +}) +const linkShape = z.object({ uri: z.string() }) + +function textTurn(text: string): Array { + return [ + { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: now(), + }, + { + type: EventType.TEXT_MESSAGE_START, + messageId: `m-${text}`, + role: 'assistant', + timestamp: now(), + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: `m-${text}`, + delta: text, + timestamp: now(), + }, + { + type: EventType.TEXT_MESSAGE_END, + messageId: `m-${text}`, + timestamp: now(), + }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'stop' } }, + }, + ] +} + +/** One model call that calls the tool `name` with no input. */ +function toolCallTurn(name: string): Array { + return [ + { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: now(), + }, + { + type: EventType.TOOL_CALL_START, + toolCallId: 'call-0', + toolCallName: name, + timestamp: now(), + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call-0', + delta: '{}', + timestamp: now(), + }, + { type: EventType.TOOL_CALL_END, toolCallId: 'call-0', timestamp: now() }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'tool_calls' } }, + }, + ] +} + +/** A model that plays `turns` in order and keeps the messages of each call. */ +function scripted(turns: Array> = []) { + const calls: Array> = [] + const adapter: AnyTextAdapter = { + kind: 'text', + name: 'mock', + model: 'mock', + // `~types` holds types only. Its values are never read, so they are casts. + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + chatStream: (options) => { + const turn = turns[calls.length] ?? textTurn('ok') + calls.push(options.messages) + return (async function* () { + yield* turn + })() + }, + } + return { adapter, calls } +} + +/** An image model that returns `bytes` as one image. */ +function imageModel(bytes: Uint8Array): ImageAdapter { + return { + kind: 'image', + name: 'fake-image', + model: 'fake-image-model', + '~types': { + providerOptions: {}, + modelProviderOptionsByName: {}, + modelSizeByName: {}, + modelInputModalitiesByName: {}, + }, + generateImages: async () => ({ + id: 'image-1', + model: 'fake-image-model', + images: [{ b64Json: Buffer.from(bytes).toString('base64') }], + }), + } +} + +/** An agent that paints one image with `bytes`. */ +function painter(bytes: Uint8Array) { + return defineAgent({ + name: 'painter', + description: 'Paints', + run: async (ctx) => { + await ctx.generateImage({ adapter: imageModel(bytes), prompt: 'a cat' }) + return 'painted' + }, + }) +} + +const cleanups: Array<() => Promise> = [] +afterEach(async () => { + for (const cleanup of cleanups.splice(0).reverse()) await cleanup() +}) + +type ServerOptions = Omit + +/** A real MCP client (spec 2026) on a harness MCP server. */ +async function serve(harness: AnyHarness, options: ServerOptions = {}) { + const host = createHarnessHost({ persistence: memoryPersistence() }) + cleanups.push(() => host.close()) + const server = await createHarnessMcpServer({ host, harness, ...options }) + const client = new Client( + { name: 'tester', version: '1.0.0' }, + { versionNegotiation: { mode: { pin: '2026-07-28' } } }, + ) + const transport = new StreamableHTTPClientTransport(serverUrl, { + fetch: async (input, init) => server.fetch(new Request(input, init)), + }) + await client.connect(transport) + cleanups.push(() => client.close()) + return client +} + +/** A chat harness on an MCP server. `calls` has the messages of each model call. */ +async function chatServer(options: ServerOptions = {}) { + const model = scripted() + const harness = defineHarness({ + name: 'test/mcp-media', + adapter: model.adapter, + }) + const client = await serve(harness, options) + return { client, calls: model.calls } +} + +function chat(client: Client, args: Record) { + return client.callTool({ name: 'chat', arguments: args }) +} + +/** The content of the last user message that a model call got. */ +function userContent(messages: Array | undefined) { + return messages?.findLast((message) => message.role === 'user')?.content +} + +/** The MCP result of a tool call that threw an error with `message`. */ +function toolError(message: string) { + return { isError: true, content: [{ type: 'text', text: message }] } +} + +/** The model content for the question "What is this?" with the image "hello". */ +const helloImageQuestion = [ + { type: 'text', content: 'What is this?' }, + { + type: 'image', + source: { type: 'data', value: helloBase64, mimeType: 'image/png' }, + }, +] + +/** + * A temp folder with `allowed/cat.png` ("hello"), `allowed/notes.xyz`, and + * `outside/secret.png`. + */ +async function folders() { + const root = await mkdtemp(join(tmpdir(), 'harness-mcp-media-')) + cleanups.push(() => rm(root, { recursive: true, force: true })) + const allowed = join(root, 'allowed') + const outside = join(root, 'outside') + await mkdir(allowed) + await mkdir(outside) + await writeFile(join(allowed, 'cat.png'), hello) + await writeFile(join(allowed, 'notes.xyz'), 'notes') + await writeFile(join(outside, 'secret.png'), 'secret') + return { allowed, outside } +} + +const outsideRefusal = (path: string) => + `Cannot read ${path}. It must be a file inside the folders this server may read.` + +describe('chat attachments', () => { + it('sends a data attachment to the model', async () => { + const { client, calls } = await chatServer() + + await chat(client, { + message: 'What is this?', + attachments: [ + { data: helloBase64, mimeType: 'image/png', name: 'cat.png' }, + ], + }) + + expect(userContent(calls[0])).toEqual(helloImageQuestion) + }) + + it('sends a path attachment inside filePaths to the model', async () => { + const { allowed } = await folders() + const { client, calls } = await chatServer({ filePaths: [allowed] }) + + await chat(client, { + message: 'What is this?', + attachments: [{ path: join(allowed, 'cat.png') }], + }) + + expect(userContent(calls[0])).toEqual(helloImageQuestion) + }) + + it.each([ + { + given: "a '..' path out of filePaths", + filePaths: true, + // Not `join`: it would remove the '..'. + path: (allowed: string) => + `${allowed}${sep}..${sep}outside${sep}secret.png`, + refusal: outsideRefusal, + }, + { + given: 'a path when the server has no filePaths', + filePaths: false, + path: (allowed: string) => join(allowed, 'cat.png'), + refusal: () => + 'This server does not read files. Send the file as data or url.', + }, + { + given: 'a file type it does not know', + filePaths: true, + path: (allowed: string) => join(allowed, 'notes.xyz'), + refusal: () => + 'Unknown file type: notes.xyz. Send the file as data with its mimeType.', + }, + ])('refuses $given', async ({ filePaths, path, refusal }) => { + const { allowed } = await folders() + const { client, calls } = await chatServer( + filePaths ? { filePaths: [allowed] } : {}, + ) + + const reply = await chat(client, { + message: 'Read this', + attachments: [{ path: path(allowed) }], + }) + + expect(reply).toMatchObject(toolError(refusal(path(allowed)))) + expect(calls).toEqual([]) + }) + + it('refuses a path through a symlink out of filePaths', async (ctx) => { + const { allowed, outside } = await folders() + // A junction needs no admin rights on Windows. Other systems make a + // normal symlink to the folder. + const linked = await symlink(outside, join(allowed, 'link'), 'junction') + .then(() => true) + .catch(() => false) + ctx.skip(!linked, 'This system cannot make a symlink.') + const { client, calls } = await chatServer({ filePaths: [allowed] }) + const path = join(allowed, 'link', 'secret.png') + + const reply = await chat(client, { + message: 'Read this', + attachments: [{ path }], + }) + + expect(reply).toMatchObject(toolError(outsideRefusal(path))) + expect(calls).toEqual([]) + }) + + it('passes a url attachment to the model and does not fetch it', async () => { + // A fetch fails, so the test never reaches the network. + const fetchSpy = vi + .spyOn(globalThis, 'fetch') + .mockRejectedValue(new Error('No network in this test.')) + cleanups.push(async () => fetchSpy.mockRestore()) + const { client, calls } = await chatServer() + + await chat(client, { + message: 'Look', + attachments: [{ url: 'https://cdn.example.com/cat.png' }], + }) + + expect(userContent(calls[0])).toEqual([ + { type: 'text', content: 'Look' }, + { + type: 'image', + source: { + type: 'url', + value: 'https://cdn.example.com/cat.png', + mimeType: 'image/png', + }, + }, + ]) + expect(fetchSpy).not.toHaveBeenCalled() + }) + + it('refuses a url when neither mimeType nor the path gives a kind', async () => { + const { client, calls } = await chatServer() + + const reply = await chat(client, { + message: 'Look', + attachments: [{ url: 'https://example.com/download' }], + }) + + expect(reply).toMatchObject( + toolError( + 'Cannot tell the kind of https://example.com/download. Add its mimeType.', + ), + ) + expect(calls).toEqual([]) + }) + + it('returns the refusal of the media store as a tool error', async () => { + const { client, calls } = await chatServer() + + const reply = await chat(client, { + message: 'Open this', + attachments: [{ data: helloBase64, mimeType: 'application/zip' }], + }) + + expect(reply).toMatchObject( + toolError('Files of type application/zip are not supported.'), + ) + expect(calls).toEqual([]) + }) +}) + +describe('media in results', () => { + it('chat returns a small image of the turn inline and lists it', async () => { + const lead = scripted([toolCallTurn('painter'), textTurn('Here it is.')]) + const client = await serve( + defineHarness({ + name: 'test/mcp-media', + adapter: lead.adapter, + subagents: { agents: [painter(hello)] }, + }), + ) + + const reply = await chat(client, { message: 'Paint a cat' }) + + expect(reply.structuredContent).toEqual({ + status: 'completed', + text: 'Here it is.', + interrupts: [], + questions: [], + media: [ + { + id: expect.any(String), + kind: 'image', + name: expect.any(String), + mimeType: 'image/png', + size: 5, + uri: expect.any(String), + }, + ], + }) + const [media] = mediaListShape.parse(reply.structuredContent).media + expect(media?.uri).toBe(`harness-media://main/${media?.id}`) + expect(reply.content).toEqual([ + { type: 'text', text: JSON.stringify(reply.structuredContent) }, + { type: 'image', data: helloBase64, mimeType: 'image/png' }, + ]) + }) + + /** Runs `agent_painter` with an image over 5 MB. `uri` is its media link. */ + async function paintBig() { + const client = await serve( + defineHarness({ + name: 'test/mcp-media', + adapter: scripted().adapter, + agents: [painter(bigImage)], + expose: { agents: ['painter'] }, + }), + ) + const reply = await client.callTool({ + name: 'agent_painter', + arguments: {}, + }) + const link = z.array(z.unknown()).parse(reply.content)[1] + return { client, reply, uri: linkShape.parse(link).uri } + } + + it('returns an image over 5 MB as a resource link', async () => { + const { reply } = await paintBig() + + expect(reply.content).toEqual([ + { type: 'text', text: 'painted' }, + { + type: 'resource_link', + uri: expect.stringMatching(/^harness-media:\/\/main\/[^/]+$/), + name: expect.any(String), + mimeType: 'image/png', + }, + ]) + }) + + it('reads the bytes of a media link with resources/read', async () => { + const { client, uri } = await paintBig() + + const read = await client.readResource({ uri }) + + const [contents] = read.contents + expect(contents).toMatchObject({ uri, mimeType: 'image/png' }) + const blob = z.object({ blob: z.string() }).parse(contents).blob + expect(Buffer.from(blob, 'base64').equals(Buffer.from(bigImage))).toBe(true) + }) + + it('refuses to read a media link with the thread of another conversation', async () => { + const { client, uri } = await paintBig() + const id = uri.slice('harness-media://main/'.length) + + const read = client.readResource({ uri: `harness-media://other/${id}` }) + + await expect(read).rejects.toThrow( + `Media ${id} was not found in thread other.`, + ) + }) +}) diff --git a/packages/ai-mcp/tests/server/create-server.test.ts b/packages/ai-mcp/tests/server/create-server.test.ts index 98bdefdf50..792e066d75 100644 --- a/packages/ai-mcp/tests/server/create-server.test.ts +++ b/packages/ai-mcp/tests/server/create-server.test.ts @@ -1149,6 +1149,33 @@ describe('createMCPServer', () => { }) }) + it('reads a uri template with its variables and the MIME type of the body', async () => { + const server = createMCPServer({ + name: 'files', + version: '1.0.0', + resources: [ + resourceDefinition({ + uriTemplate: 'media://{folder}/{id}', + name: 'media', + mimeType: 'application/octet-stream', + argsSchema: z.object({ folder: z.string(), id: z.string() }), + }).read(async ({ folder, id }) => ({ + blob: btoa(`${folder}/${id}`), + mimeType: 'image/png', + })), + ], + }) + + await withClient(server, { era: '2026' }, async (client) => { + const read = await client.readResource({ uri: 'media://cats/7' }) + expect(read.contents[0]).toEqual({ + uri: 'media://cats/7', + mimeType: 'image/png', + blob: btoa('cats/7'), + }) + }) + }) + it('keeps assistant prompt messages and drops invalid ones', async () => { const server = createMCPServer({ name: 'prompts', diff --git a/packages/ai-mcp/tests/server/definitions.test.ts b/packages/ai-mcp/tests/server/definitions.test.ts index 5fee912bc5..7f6de3d383 100644 --- a/packages/ai-mcp/tests/server/definitions.test.ts +++ b/packages/ai-mcp/tests/server/definitions.test.ts @@ -27,6 +27,19 @@ describe('resourceDefinition', () => { expect(await resource.read()).toEqual({ text: 'body' }) }) + it('gives the read function the variables parsed by argsSchema', async () => { + const resource = resourceDefinition({ + name: 'user', + mimeType: 'text/plain', + uriTemplate: 'users://{id}', + argsSchema: z.object({ id: z.string().transform(Number) }), + }).read(async ({ id }, uri) => ({ text: `${id + 1} ${uri?.href}` })) + + expect(await resource.read({ id: '41' }, new URL('users://41'))).toEqual({ + text: '42 users://41', + }) + }) + it('throws when uri and uriTemplate are missing', () => { expect(() => resourceDefinition({ From f10e6f51bfbd46aca648f3ec62c3f829ba806496 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 29 Sep 2026 12:55:06 +0200 Subject: [PATCH 062/286] feat(ai-harness-cli): attach files with @path, save media to a folder, and an E2E test for web media Line mode, -p, piped input, and a custom ui attach @path files. Generated media is saved to ./-media (--media-dir to change it): line mode prints a saved line, and -p --output ndjson adds the path to the harness.media event. --mcp lets path attachments read the working folder. The MIME lookup is shared as mimeTypeOf. The E2E test uploads an image, sends it, and loads the generated image from its signed URL. Claude-Session: https://claude.ai/code/session_01APYv1qshKyjPPpkFyRZhfZ --- packages/ai-harness-cli/src/args.ts | 7 + packages/ai-harness-cli/src/attach.ts | 155 +++++++++ packages/ai-harness-cli/src/commands.ts | 21 +- packages/ai-harness-cli/src/index.ts | 51 ++- packages/ai-harness-cli/src/lines.ts | 14 +- packages/ai-harness-cli/src/print.ts | 39 ++- packages/ai-harness-cli/src/printer.ts | 24 +- packages/ai-harness-cli/tests/attach.test.ts | 208 +++++++++++ .../ai-harness-cli/tests/child-view.test.ts | 12 +- .../ai-harness-cli/tests/cli-media.test.ts | 323 ++++++++++++++++++ packages/ai-harness/src/index.ts | 1 + packages/ai-harness/src/media-ref.ts | 38 +++ packages/ai-mcp/src/harness.ts | 31 +- testing/e2e/fixtures/harness/media.json | 37 ++ .../e2e/src/routes/api.harness-protocol.$.ts | 30 +- testing/e2e/tests/harness-media.spec.ts | 174 ++++++++++ 16 files changed, 1109 insertions(+), 56 deletions(-) create mode 100644 packages/ai-harness-cli/src/attach.ts create mode 100644 packages/ai-harness-cli/tests/attach.test.ts create mode 100644 packages/ai-harness-cli/tests/cli-media.test.ts create mode 100644 testing/e2e/fixtures/harness/media.json create mode 100644 testing/e2e/tests/harness-media.spec.ts diff --git a/packages/ai-harness-cli/src/args.ts b/packages/ai-harness-cli/src/args.ts index aa9e0d06f8..c5b89c797f 100644 --- a/packages/ai-harness-cli/src/args.ts +++ b/packages/ai-harness-cli/src/args.ts @@ -13,12 +13,14 @@ export interface CliArgs { token?: string thread: string dashboard?: string + mediaDir?: string } export const USAGE = `Usage: [options] With no options in a terminal, starts the UI of this CLI, if it has one. Otherwise, and for piped input, reads one message or command per line. +In a message, @path sends a file with it (@"my cat.png" for a path with spaces). Options: -p, --print Run one prompt, print the answer, and exit @@ -32,6 +34,7 @@ Options: --token Bearer token for --serve (default: HARNESS_TOKEN, or a new random token) --thread Conversation id (default: main) --dashboard Connect to a dashboard (pairs on first use; set HARNESS_DASHBOARD_TOKEN to skip) + --media-dir Save the media a turn makes here (default: ./-media) -h, --help Show this help Exit codes for --print: 0 done, 1 failed, 2 waiting for approval, 130 cancelled.` @@ -53,6 +56,7 @@ export function parseCliArgs(argv: ReadonlyArray): CliArgs { token: { type: 'string' }, thread: { type: 'string' }, dashboard: { type: 'string' }, + 'media-dir': { type: 'string' }, }, strict: true, }) @@ -77,5 +81,8 @@ export function parseCliArgs(argv: ReadonlyArray): CliArgs { ...(values.token !== undefined ? { token: values.token } : {}), thread: values.thread ?? 'main', ...(values.dashboard !== undefined ? { dashboard: values.dashboard } : {}), + ...(values['media-dir'] !== undefined + ? { mediaDir: values['media-dir'] } + : {}), } } diff --git a/packages/ai-harness-cli/src/attach.ts b/packages/ai-harness-cli/src/attach.ts new file mode 100644 index 0000000000..880689a407 --- /dev/null +++ b/packages/ai-harness-cli/src/attach.ts @@ -0,0 +1,155 @@ +import { mkdir, readFile, stat, writeFile } from 'node:fs/promises' +import { basename, join, resolve } from 'node:path' +import { mediaPart, mimeTypeOf } from '@tanstack/ai-harness' +import type { ContentPart } from '@tanstack/ai' +import type { + HarnessSession, + MediaRecord, + UserInput, +} from '@tanstack/ai-harness' + +/** + * An `@path` token: an `@` at the start of the text or after white space, + * then a quoted path (`@"my cat.png"`) or the text up to the next white + * space. The white space before the `@` is part of the match, so it goes + * out of the text with the token. + */ +const TOKEN = /(^|\s)@(?:"([^"]+)"|(\S+))/g + +/** The path of a token. One of the two groups always matches. */ +function pathOf(token: RegExpExecArray) { + return token[2] ?? token[3] ?? '' +} + +/** + * The file that `path` names, when it is a readable regular file, else + * `undefined`. A file with an extension that `mimeTypeOf` does not know + * throws. + */ +async function readToken(path: string, cwd: string) { + const file = resolve(cwd, path) + const info = await stat(file).catch(() => undefined) + if (!info?.isFile()) return undefined + const name = basename(file) + const mimeType = mimeTypeOf(name) + if (mimeType === undefined) { + throw new Error( + `Cannot attach ${name}: the file type is not known. To send the path as text, remove the @.`, + ) + } + // ponytail: reads the whole file, and the store refuses it after that when + // it is over `media.maxBytes`. Stream it when big local files matter. + const bytes = await readFile(file).catch(() => undefined) + return bytes === undefined ? undefined : { bytes, mimeType, name } +} + +/** + * The message for `text`, with each `@path` token that names a file sent as + * a stored media part. The token goes out of the text. Any other `@` stays + * as typed: `me@acme.dev`, `@types/node`, or a path with no file. Returns + * `text` as it is when no token names a file. + * + * Throws for a file of a type it does not know, before it stores a file. + * `putMedia` throws a `MediaError` for a file over `media.maxBytes`, or of a + * kind that the harness does not take. The message does not go then. + */ +export async function attach( + session: HarnessSession, + text: string, + options: { cwd: string }, +) { + const tokens = [...text.matchAll(TOKEN)] + // Read every file first, so a bad file stops the send before a store. + const files = await Promise.all( + tokens.map((token) => readToken(pathOf(token), options.cwd)), + ) + if (files.every((file) => file === undefined)) return text + + const media: Array = [] + let rest = '' + let end = 0 + for (const [index, token] of tokens.entries()) { + const file = files[index] + if (file === undefined) continue + rest += text.slice(end, token.index) + end = token.index + token[0].length + const record = await session.putMedia(file.bytes, { + mimeType: file.mimeType, + name: file.name, + }) + media.push(mediaPart(record)) + } + const content = `${rest}${text.slice(end)}`.trim() + const input: UserInput = + content === '' ? media : [{ type: 'text', content }, ...media] + return input +} + +function isTaken(error: unknown) { + return error instanceof Error && 'code' in error && error.code === 'EEXIST' +} + +/** + * Write the bytes of a media file into `dir`, and return the path. The file + * gets the name of the record, with `-1`, `-2`, ... before the extension + * when that name is taken. It never overwrites a file. A name with a folder + * in it (`/`, `\`, a drive `:`), `.`, or `..` is refused, because a name can + * come from an upload. + */ +export async function saveMedia( + session: HarnessSession, + record: Pick, + dir: string, +) { + const { name } = record + const isPlainName = + name !== '' && name !== '.' && name !== '..' && !/[\\/:]/.test(name) + if (!isPlainName) { + throw new Error( + `Cannot save ${JSON.stringify(name)}: it is not a plain file name.`, + ) + } + const bytes = await session.loadMedia(record.id) + await mkdir(dir, { recursive: true }) + const dot = name.lastIndexOf('.') + const stem = dot > 0 ? name.slice(0, dot) : name + const extension = dot > 0 ? name.slice(dot) : '' + for (let count = 0; ; count += 1) { + const path = join(dir, count === 0 ? name : `${stem}-${count}${extension}`) + try { + // `wx` fails when the file is there, so a file is never overwritten. + await writeFile(path, bytes, { flag: 'wx' }) + return path + } catch (error) { + if (!isTaken(error)) throw error + } + } +} + +/** + * Save a media file with `saveMedia`, and get the line that says where it + * went, or why it was not saved. It does not throw. + */ +export async function saveMediaLine( + session: HarnessSession, + record: Pick, + dir: string, +) { + try { + const path = await saveMedia(session, record, dir) + return { path, text: `[${record.kind} saved: ${path}]` } + } catch (error) { + const reason = error instanceof Error ? error.message : String(error) + return { text: `[${record.kind} not saved: ${reason}]` } + } +} + +/** + * The default media folder of a harness, relative to the working folder: + * the name as a safe folder name, then `-media`. `acme/coder` saves into + * `acme-coder-media`. + */ +export function defaultMediaDir(harnessName: string) { + const safe = harnessName.replace(/[^\w.-]+/g, '-').replace(/^[-.]+|-+$/g, '') + return `${safe || 'harness'}-media` +} diff --git a/packages/ai-harness-cli/src/commands.ts b/packages/ai-harness-cli/src/commands.ts index e78eb69852..34a7390945 100644 --- a/packages/ai-harness-cli/src/commands.ts +++ b/packages/ai-harness-cli/src/commands.ts @@ -1,4 +1,5 @@ -import type { HarnessSession } from '@tanstack/ai-harness' +import { attach } from './attach' +import type { HarnessSession, UserInput } from '@tanstack/ai-harness' export const HELP_TEXT = [ 'Type a message and press Enter. While the agent works, a new message steers it.', @@ -83,7 +84,9 @@ export type LineResult = /** * Handle one line of user input: a slash command, a steer while a turn runs, - * or a new prompt. Used by the line mode, and by any UI that calls it. + * or a new prompt. Used by the line mode, and by any UI that calls it. A + * prompt or a steer sends each `@path` file (relative to the working folder) + * with it. A command and an answer do not. */ export async function handleLine( session: HarnessSession, @@ -106,8 +109,18 @@ export async function handleLine( } } if (!text.startsWith('/')) { - if (session.snapshot().status === 'running') await session.steer(text) - else void session.prompt(text) + let input: UserInput + try { + input = await attach(session, text, { cwd: process.cwd() }) + } catch (error) { + // A file that cannot go (a type it does not know, too big) stops the send. + return { + type: 'notice', + text: error instanceof Error ? error.message : String(error), + } + } + if (session.snapshot().status === 'running') await session.steer(input) + else void session.prompt(input) return { type: 'sent' } } const [name = '', ...rest] = text.slice(1).split(' ') diff --git a/packages/ai-harness-cli/src/index.ts b/packages/ai-harness-cli/src/index.ts index b2ef38b820..ac057e0a99 100644 --- a/packages/ai-harness-cli/src/index.ts +++ b/packages/ai-harness-cli/src/index.ts @@ -1,11 +1,17 @@ import { createHarnessHost } from '@tanstack/ai-harness' import { createSessionView } from '@tanstack/ai-harness/view' import { USAGE, parseCliArgs } from './args' +import { attach, defaultMediaDir } from './attach' import { runLines } from './lines' import { EXIT, runPrint } from './print' import { createToken, serve } from './serve' -import type { AnyHarness, HarnessPersistence } from '@tanstack/ai-harness' -import type { SessionView } from '@tanstack/ai-harness/view' +import type { + AnyHarness, + HarnessPersistence, + HarnessSession, + UserInput, +} from '@tanstack/ai-harness' +import type { SessionView, SessionViewSource } from '@tanstack/ai-harness/view' export interface RunCliOptions { /** Where sessions keep state. Default: in memory. */ @@ -20,10 +26,43 @@ export interface RunCliOptions { * Your own screen for an interactive terminal, with any UI library. It gets * a ready session view and resolves when the user quits. Piped input and * the other modes (`--print`, `--acp`, `--mcp`, `--serve`, `--dashboard`) do not use it. + * `view.send(text)` sends each `@path` file in `text`, as line mode does. */ ui?: (view: SessionView) => Promise | void } +/** + * The session for a `ui` view: a typed message or steer sends its `@path` + * files, as in line mode. + */ +function attachingSource(session: HarnessSession) { + const withFiles = (message: UserInput) => + typeof message === 'string' + ? attach(session, message, { cwd: process.cwd() }) + : Promise.resolve(message) + const source: SessionViewSource = { + // The view shows a failed operation from its events, so only a file that + // cannot go rejects here. It becomes an error notice in the view. + prompt: (message) => + withFiles(message).then((input) => { + session.prompt(input).then(undefined, () => {}) + }), + steer: async (message) => session.steer(await withFiles(message)), + resolve: (resume) => session.resolve(resume), + cancel: (operationId) => session.cancel(operationId), + answer: (questionId, value) => session.answer(questionId, value), + command: (name, input) => session.command(name, input), + setConfig: (key, value) => session.setConfig(key, value), + events: (options) => session.events(options), + snapshot: () => session.snapshot(), + transcript: () => session.transcript(), + describe: () => session.describe(), + mediaUrl: (id) => session.mediaUrl(id), + loadMedia: (id) => session.loadMedia(id), + } + return source +} + /** * Run a harness from the terminal. Resolves to the process exit code. * @@ -88,6 +127,8 @@ export async function runCli( harness, threadId: args.thread, approvals, + // A local client can attach the files of the working folder by path. + filePaths: [process.cwd()], }) // stdout carries only MCP messages. The server stops when stdin ends. const ended = new Promise((resolve) => { @@ -153,15 +194,18 @@ export async function runCli( } const session = await host.open(harness, { threadId: args.thread }) + // Relative to the working folder, as a relative --media-dir is. + const mediaDir = args.mediaDir ?? defaultMediaDir(harness.name) if (args.print !== undefined) { return await runPrint(session, args.print, { output: args.output, stdout, stderr, + mediaDir, }) } if (stdin.isTTY && options.ui) { - const view = createSessionView(session) + const view = createSessionView(attachingSource(session)) try { await view.ready await options.ui(view) @@ -171,6 +215,7 @@ export async function runCli( } else { await runLines(session, stdin, stdout, { openSignIns: Boolean(stdin.isTTY), + mediaDir, }) } return EXIT.ok diff --git a/packages/ai-harness-cli/src/lines.ts b/packages/ai-harness-cli/src/lines.ts index a245db0c0b..a4d940dc87 100644 --- a/packages/ai-harness-cli/src/lines.ts +++ b/packages/ai-harness-cli/src/lines.ts @@ -1,5 +1,6 @@ import { createInterface } from 'node:readline' import { createSessionView } from '@tanstack/ai-harness/view' +import { saveMediaLine } from './attach' import { handleLine } from './commands' import { createPrinter } from './printer' import { approvalQuestion, openUrl, resolveAll, waitIdle } from './session-view' @@ -15,15 +16,23 @@ interface Output { * next line answers it (`y` approves, anything else rejects). * * `openSignIns`: open sign-in links in the browser (for an interactive terminal). + * `mediaDir`: the folder for the media a turn makes. Each new file is saved + * there, with one `[image saved: ]` line. */ export async function runLines( session: HarnessSession, input: NodeJS.ReadableStream, stdout: Output, - options: { openSignIns?: boolean } = {}, + options: { openSignIns?: boolean; mediaDir: string }, ) { const view = createSessionView(session) - const printer = createPrinter((text) => stdout.write(text)) + const saves: Array> = [] + const printer = createPrinter((text) => stdout.write(text), { + onMedia: (part) => { + const saved = saveMediaLine(session, part, options.mediaDir) + saves.push(saved.then(({ text }) => printer.line(text))) + }, + }) await view.ready // Line mode prints only what happens from now on. printer.mark(view.store.get()) @@ -65,6 +74,7 @@ export async function runLines( stopOpening() printing.unsubscribe() printer.print(view.store.get()) + await Promise.allSettled(saves) view.dispose() printer.end() } diff --git a/packages/ai-harness-cli/src/print.ts b/packages/ai-harness-cli/src/print.ts index fbb8ac271d..c5c06d2b16 100644 --- a/packages/ai-harness-cli/src/print.ts +++ b/packages/ai-harness-cli/src/print.ts @@ -1,4 +1,6 @@ import { EventType } from '@tanstack/ai' +import { HARNESS_EVENTS, isMediaRecord } from '@tanstack/ai-harness' +import { attach, saveMediaLine } from './attach' import type { HarnessSession } from '@tanstack/ai-harness' export const EXIT = { @@ -14,17 +16,44 @@ interface Output { /** * Run one prompt and print it. `text` prints the answer as it streams. - * `ndjson` prints every AG-UI event as one JSON line. + * `ndjson` prints every AG-UI event as one JSON line. Each `@path` file in + * the prompt goes with it. + * + * Each media file the turn makes is saved into `mediaDir`. `ndjson` adds the + * saved `path` to the value of its `harness.media` event. `text` prints the + * saved line to stderr, so stdout keeps only the answer. */ export async function runPrint( session: HarnessSession, prompt: string, - options: { output: 'text' | 'ndjson'; stdout: Output; stderr: Output }, -): Promise { + options: { + output: 'text' | 'ndjson' + stdout: Output + stderr: Output + mediaDir: string + }, +) { const { stdout, stderr } = options - const operation = session.prompt(prompt) + const input = await attach(session, prompt, { cwd: process.cwd() }) + const operation = session.prompt(input) for await (const chunk of operation.stream()) { - if (options.output === 'ndjson') { + const media = + chunk.type === EventType.CUSTOM && + chunk.name === HARNESS_EVENTS.media && + isMediaRecord(chunk.value) + ? chunk.value + : undefined + if (media !== undefined) { + const saved = await saveMediaLine(session, media, options.mediaDir) + if (options.output === 'text' || saved.path === undefined) { + stderr.write(`${saved.text}\n`) + } + if (options.output === 'ndjson') { + // A file that was not saved has no `path`: JSON drops `undefined`. + const value = { ...media, path: saved.path } + stdout.write(`${JSON.stringify({ ...chunk, value })}\n`) + } + } else if (options.output === 'ndjson') { stdout.write(`${JSON.stringify(chunk)}\n`) } else if ( chunk.type === EventType.TEXT_MESSAGE_CONTENT && diff --git a/packages/ai-harness-cli/src/printer.ts b/packages/ai-harness-cli/src/printer.ts index 8b4c892721..ead061f8ad 100644 --- a/packages/ai-harness-cli/src/printer.ts +++ b/packages/ai-harness-cli/src/printer.ts @@ -1,13 +1,24 @@ import { shortAnswer, signInText } from './session-view' -import type { AgentPart, SessionViewState } from '@tanstack/ai-harness/view' +import type { + AgentPart, + MediaPart, + SessionViewState, + ViewPart, +} from '@tanstack/ai-harness/view' /** * Prints a session view as lines: streamed text as it grows, and one * bracketed line per tool call, child agent, notice, question, and sign-in. * Notices that line mode prints itself (command results, rejected inputs, * and UI notes) are skipped. + * + * Media parts print nothing here. `onMedia` gets each new media part of an + * assistant message once, from the lead or from a child at any depth. */ -export function createPrinter(write: (text: string) => unknown) { +export function createPrinter( + write: (text: string) => unknown, + options: { onMedia?: (part: MediaPart) => void } = {}, +) { const done = new Set() const written = new Map() // True while streamed text has no line break at its end yet. @@ -57,6 +68,14 @@ export function createPrinter(write: (text: string) => unknown) { if (part.status === 'failed') once(`${key}:end`, `agent ${part.name} failed: ${part.error ?? ''}`) } + const media = (parts: ReadonlyArray) => { + for (const part of parts) { + if (part.type === 'agent') media(part.parts) + if (part.type !== 'media' || done.has(`media:${part.id}`)) continue + done.add(`media:${part.id}`) + if (!marking) options.onMedia?.(part) + } + } const print = (state: SessionViewState) => { for (const message of state.messages) { if (message.role === 'notice') { @@ -71,6 +90,7 @@ export function createPrinter(write: (text: string) => unknown) { if (part.type === 'tool-call') once(key, `tool ${part.name}`) if (part.type === 'agent') agent(key, part) }) + media(message.parts) } for (const question of state.questions) once(`question:${question.id}`, `? ${question.message}`) diff --git a/packages/ai-harness-cli/tests/attach.test.ts b/packages/ai-harness-cli/tests/attach.test.ts new file mode 100644 index 0000000000..1f0aaea15b --- /dev/null +++ b/packages/ai-harness-cli/tests/attach.test.ts @@ -0,0 +1,208 @@ +import { mkdir, mkdtemp, readFile, rm, stat, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { + afterAll, + beforeAll, + beforeEach, + describe, + expect, + it, + vi, +} from 'vitest' +import { + createHarnessHost, + defineHarness, + mediaIdOf, +} from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { attach, defaultMediaDir, saveMedia } from '../src/attach' +import type { AnyTextAdapter } from '@tanstack/ai' +import type { + HarnessHost, + HarnessSession, + UserInput, +} from '@tanstack/ai-harness' + +// The real file system, with spies, so a test can see which files are read. +vi.mock('node:fs/promises', { spy: true }) + +/** A model that is never called: these tests only store and save files. */ +function idleModel(): AnyTextAdapter { + return { + kind: 'text', + name: 'mock', + model: 'test-model', + // `~types` holds types only. Its values are never read, so they are casts. + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + chatStream: () => { + throw new Error('The model must not be called.') + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } +} + +/** The part that `attach` makes for a stored PNG. */ +const pngPart = { + type: 'image', + source: { + type: 'url', + value: expect.stringMatching(/^harness-media:./), + mimeType: 'image/png', + }, +} + +let dir = '' +let host: HarnessHost +let session: HarnessSession + +beforeAll(async () => { + dir = await mkdtemp(join(tmpdir(), 'harness-attach-')) + await writeFile(join(dir, 'cat.png'), 'hello') + await writeFile(join(dir, 'my cat.png'), 'meow') + await writeFile(join(dir, 'notes.xyz'), 'plain') + await mkdir(join(dir, 'folder.png')) + host = createHarnessHost({ persistence: memoryPersistence() }) + session = await host.open( + defineHarness({ name: 'test/attach', adapter: idleModel() }), + { threadId: 't' }, + ) +}) + +afterAll(async () => { + await host.close() + await rm(dir, { recursive: true, force: true }) +}) + +beforeEach(() => { + vi.clearAllMocks() +}) + +/** The name and text of each stored file that `input` sends. */ +async function storedFiles(input: UserInput) { + const parts = typeof input === 'string' ? [] : input + const ids = parts.flatMap((part) => mediaIdOf(part) ?? []) + return Promise.all( + ids.map(async (id) => ({ + name: (await session.getMedia(id))?.name, + text: new TextDecoder().decode(await session.loadMedia(id)), + })), + ) +} + +describe('attach', () => { + it('sends a file token as a media part and takes the token out of the text', async () => { + const input = await attach(session, 'what is wrong in @cat.png today', { + cwd: dir, + }) + + expect(input).toEqual([ + { type: 'text', content: 'what is wrong in today' }, + pngPart, + ]) + expect(await storedFiles(input)).toEqual([ + { name: 'cat.png', text: 'hello' }, + ]) + }) + + it('reads a quoted path with spaces', async () => { + const input = await attach(session, 'look at @"my cat.png"', { cwd: dir }) + + expect(input).toEqual([{ type: 'text', content: 'look at' }, pngPart]) + expect(await storedFiles(input)).toEqual([ + { name: 'my cat.png', text: 'meow' }, + ]) + }) + + it('sends no text part for a message that is only a file', async () => { + expect(await attach(session, '@cat.png', { cwd: dir })).toEqual([pngPart]) + }) + + it('keeps an email and a package name as text, and reads no file for the email', async () => { + const email = 'write to me@acme.dev' + const packageName = 'install @types/node first' + + expect(await attach(session, email, { cwd: dir })).toBe(email) + expect(stat).not.toHaveBeenCalled() + expect(readFile).not.toHaveBeenCalled() + expect(await attach(session, packageName, { cwd: dir })).toBe(packageName) + // A token that can be a path is looked up, so the spy above sees reads. + expect(stat).toHaveBeenCalled() + }) + + it.each([ + ['a missing file', 'see @missing.png now'], + ['a folder', 'see @folder.png now'], + ])('keeps %s as text', async (_label, text) => { + expect(await attach(session, text, { cwd: dir })).toBe(text) + }) + + it('refuses a file of a type it does not know', async () => { + await expect( + attach(session, 'read @notes.xyz', { cwd: dir }), + ).rejects.toThrow( + 'Cannot attach notes.xyz: the file type is not known. To send the path as text, remove the @.', + ) + }) +}) + +describe('saveMedia', () => { + it('makes the folder and adds a number instead of overwriting a file', async () => { + const record = await session.putMedia(new TextEncoder().encode('hello'), { + mimeType: 'image/png', + name: 'cat.png', + }) + const out = join(dir, 'saved', 'media') + + const paths = [ + await saveMedia(session, record, out), + await saveMedia(session, record, out), + await saveMedia(session, record, out), + ] + + expect(paths).toEqual([ + join(out, 'cat.png'), + join(out, 'cat-1.png'), + join(out, 'cat-2.png'), + ]) + expect(await readFile(join(out, 'cat-2.png'), 'utf8')).toBe('hello') + }) + + it.each(['../escaped.png', 'sub/escaped.png', 'sub\\escaped.png', '..'])( + 'refuses the name %s', + async (name) => { + const record = await session.putMedia(new TextEncoder().encode('x'), { + mimeType: 'image/png', + name, + }) + const out = join(dir, 'refused') + + await expect(saveMedia(session, record, out)).rejects.toThrow( + 'it is not a plain file name', + ) + await expect(stat(join(dir, 'escaped.png'))).rejects.toThrow() + }, + ) +}) + +describe('defaultMediaDir', () => { + it.each([ + ['acme/coder', 'acme-coder-media'], + ['@acme/coder', 'acme-coder-media'], + ])('saves the media of %s into %s', (name, folder) => { + expect(defaultMediaDir(name)).toBe(folder) + }) +}) diff --git a/packages/ai-harness-cli/tests/child-view.test.ts b/packages/ai-harness-cli/tests/child-view.test.ts index 046e355427..f06771b9fa 100644 --- a/packages/ai-harness-cli/tests/child-view.test.ts +++ b/packages/ai-harness-cli/tests/child-view.test.ts @@ -1,3 +1,5 @@ +import { tmpdir } from 'node:os' +import { join } from 'node:path' import { PassThrough } from 'node:stream' import { describe, expect, it } from 'vitest' import { z } from 'zod' @@ -119,9 +121,13 @@ describe('line mode with a child agent', () => { ) const input = new PassThrough() let out = '' - const running = runLines(session, input, { - write: (text: string) => (out += text), - }) + // This turn makes no media, so nothing is written to the media folder. + const running = runLines( + session, + input, + { write: (text: string) => (out += text) }, + { mediaDir: join(tmpdir(), 'child-view-media') }, + ) input.end('go\n') await running await host.close() diff --git a/packages/ai-harness-cli/tests/cli-media.test.ts b/packages/ai-harness-cli/tests/cli-media.test.ts new file mode 100644 index 0000000000..5cce9b5142 --- /dev/null +++ b/packages/ai-harness-cli/tests/cli-media.test.ts @@ -0,0 +1,323 @@ +import { mkdtemp, readFile, readdir, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { Readable } from 'node:stream' +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' +import { EventType, defineAgent } from '@tanstack/ai' +import { defineHarness } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { EXIT, runCli } from '../src' +import type { AnyTextAdapter, ImageAdapter, StreamChunk } from '@tanstack/ai' + +// ponytail: `textTurn`, `toolTurn`, `scripted`, `capture`, and `stdinFrom` +// repeat the ones in cli.test.ts. Move them to a helpers file when a third +// test file needs them. +const now = () => Date.now() +const textTurn = (text: string): Array => [ + { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, + { + type: EventType.TEXT_MESSAGE_START, + messageId: 'm', + role: 'assistant', + timestamp: now(), + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm', + delta: text, + timestamp: now(), + }, + { type: EventType.TEXT_MESSAGE_END, messageId: 'm', timestamp: now() }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'stop' } }, + }, +] +const toolTurn = (name: string, args: string): Array => [ + { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, + { + type: EventType.TOOL_CALL_START, + toolCallId: 'call_1', + toolCallName: name, + timestamp: now(), + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call_1', + delta: args, + timestamp: now(), + }, + { type: EventType.TOOL_CALL_END, toolCallId: 'call_1', timestamp: now() }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'tool_calls' } }, + }, +] + +function scripted(turns: Array>) { + let call = 0 + const seen: Array> = [] + const adapter: AnyTextAdapter = { + kind: 'text', + name: 'mock', + model: 'test-model', + // `~types` holds types only. Its values are never read, so they are casts. + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + chatStream: (options) => { + seen.push(options.messages) + const chunks = turns[call] ?? textTurn('') + call += 1 + return (async function* () { + yield* chunks + })() + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } + return { adapter, seen } +} + +function capture() { + let text = '' + return { + write: (chunk: string) => (text += chunk), + get text() { + return text + }, + } +} + +/** A stdin that reads `lines`, as a terminal or as a pipe. */ +function stdinFrom(lines: Array, { isTTY }: { isTTY: boolean }) { + // NodeJS.ReadStream is a TTY socket. A test cannot make one without a real + // terminal, and runCli only reads the lines and `isTTY`, so this casts. + return Object.assign(Readable.from(lines), { + isTTY, + }) as unknown as NodeJS.ReadStream +} + +/** An image model that paints the bytes "hello". */ +function imageAdapter(): ImageAdapter { + return { + kind: 'image', + name: 'fake-image', + model: 'fake-image-model', + '~types': { + providerOptions: {}, + modelProviderOptionsByName: {}, + modelSizeByName: {}, + modelInputModalitiesByName: {}, + }, + // 'aGVsbG8=' is base64 for "hello". + generateImages: async () => ({ + id: 'image-1', + model: 'fake-image-model', + images: [{ b64Json: 'aGVsbG8=' }], + }), + } +} + +const painter = defineAgent({ + name: 'painter', + description: 'Paints', + run: async (ctx) => { + await ctx.generateImage({ adapter: imageAdapter(), prompt: 'a cat' }) + return 'painted' + }, +}) + +/** + * Run a harness whose lead calls the painter once, then answers, with + * `argv` and the piped `lines`. Gives the exit code and the output. + */ +async function paint(argv: Array, lines: Array = []) { + const { adapter } = scripted([ + toolTurn('painter', '{}'), + textTurn('Here is your cat.'), + ]) + const harness = defineHarness({ + name: 'test/painting', + adapter, + subagents: { agents: [painter] }, + }) + const stdout = capture() + const stderr = capture() + const code = await runCli(harness, { + argv, + stdin: stdinFrom(lines, { isTTY: false }), + stdout, + stderr, + persistence: memoryPersistence(), + }) + return { code, stdout: stdout.text, stderr: stderr.text } +} + +let root = '' + +beforeAll(async () => { + root = await mkdtemp(join(tmpdir(), 'harness-cli-media-')) + await writeFile(join(root, 'cat.png'), 'hello') + await writeFile(join(root, 'notes.xyz'), 'plain') +}) + +afterAll(async () => { + await rm(root, { recursive: true, force: true }) +}) + +/** A prompt that sends cat.png by its full path, quoted for spaces. */ +const aboutCat = () => `what is @"${join(root, 'cat.png')}"` + +/** What the model sees for `aboutCat()`: the text, then the bytes. */ +const seenAboutCat = [ + { type: 'text', content: 'what is' }, + { + type: 'image', + source: { type: 'data', value: 'aGVsbG8=', mimeType: 'image/png' }, + }, +] + +/** The one file in `dir`, as its path and its text. */ +async function onlyFile(dir: string) { + const names = await readdir(dir) + expect(names).toHaveLength(1) + const path = join(dir, names[0] ?? '') + return { path, text: await readFile(path, 'utf8') } +} + +describe('media in line mode', () => { + it('saves a new image into the media folder once and prints where', async () => { + const mediaDir = join(root, 'lines') + + const { code, stdout } = await paint( + ['--media-dir', mediaDir], + ['paint a cat\n'], + ) + + expect(code).toBe(EXIT.ok) + const saved = await onlyFile(mediaDir) + expect(saved.text).toBe('hello') + expect(stdout).toContain(`[image saved: ${saved.path}]\n`) + expect(stdout).toContain('Here is your cat.') + }) + + it('prints why a file cannot go, and does not call the model', async () => { + const { adapter, seen } = scripted([]) + const stdout = capture() + + await runCli(defineHarness({ name: 'test/media-refused', adapter }), { + argv: [], + stdin: stdinFrom([`read @"${join(root, 'notes.xyz')}"\n`], { + isTTY: false, + }), + stdout, + stderr: capture(), + persistence: memoryPersistence(), + }) + + expect(stdout.text).toContain( + 'Cannot attach notes.xyz: the file type is not known.', + ) + expect(seen).toHaveLength(0) + }) +}) + +describe('media in print mode', () => { + it('prints the media event with the saved path in ndjson', async () => { + const mediaDir = join(root, 'ndjson') + + const { stdout } = await paint([ + '-p', + 'paint a cat', + '--output', + 'ndjson', + '--media-dir', + mediaDir, + ]) + + const saved = await onlyFile(mediaDir) + const media = stdout + .trim() + .split('\n') + .map((line) => JSON.parse(line)) + .filter((event) => event.name === 'harness.media') + expect(media).toEqual([ + expect.objectContaining({ + type: 'CUSTOM', + value: expect.objectContaining({ + kind: 'image', + source: 'generated', + size: 5, + path: saved.path, + }), + }), + ]) + }) + + it('prints the saved line to stderr and only the answer to stdout', async () => { + const mediaDir = join(root, 'text') + + const { code, stdout, stderr } = await paint([ + '-p', + 'paint a cat', + '--media-dir', + mediaDir, + ]) + + expect(code).toBe(EXIT.ok) + const saved = await onlyFile(mediaDir) + expect(stdout).toBe('Here is your cat.\n') + expect(stderr).toBe(`[image saved: ${saved.path}]\n`) + }) + + it('sends an @file to the model as data', async () => { + const { adapter, seen } = scripted([textTurn('A cat.')]) + + const code = await runCli(defineHarness({ name: 'test/attach', adapter }), { + argv: ['-p', aboutCat()], + stdout: capture(), + stderr: capture(), + persistence: memoryPersistence(), + }) + + expect(code).toBe(EXIT.ok) + expect(seen[0]?.at(-1)?.content).toEqual(seenAboutCat) + }) +}) + +describe('media in a custom ui', () => { + it('sends an @file of view.send to the model', async () => { + const { adapter, seen } = scripted([textTurn('A cat.')]) + + await runCli(defineHarness({ name: 'test/ui-attach', adapter }), { + argv: [], + stdin: stdinFrom([], { isTTY: true }), + stdout: capture(), + stderr: capture(), + persistence: memoryPersistence(), + ui: async (view) => { + await view.send(aboutCat()) + await vi.waitFor(() => expect(seen).toHaveLength(1)) + }, + }) + + expect(seen[0]?.at(-1)?.content).toEqual(seenAboutCat) + }) +}) diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts index 3db6e68777..6ae1431c6f 100644 --- a/packages/ai-harness/src/index.ts +++ b/packages/ai-harness/src/index.ts @@ -98,6 +98,7 @@ export { mediaIdOf, mediaOfMessage, mediaPart, + mimeTypeOf, } from './media-ref' export { MediaError } from './media' diff --git a/packages/ai-harness/src/media-ref.ts b/packages/ai-harness/src/media-ref.ts index 5e83ac15d2..b3806cecad 100644 --- a/packages/ai-harness/src/media-ref.ts +++ b/packages/ai-harness/src/media-ref.ts @@ -76,6 +76,44 @@ export function kindOf(mimeType: string) { } } +// ponytail: common types only. Send any other file with its MIME type. +const mimeTypes = new Map([ + ['png', 'image/png'], + ['jpg', 'image/jpeg'], + ['jpeg', 'image/jpeg'], + ['gif', 'image/gif'], + ['webp', 'image/webp'], + ['mp3', 'audio/mpeg'], + ['wav', 'audio/wav'], + ['ogg', 'audio/ogg'], + ['m4a', 'audio/mp4'], + ['flac', 'audio/flac'], + ['mp4', 'video/mp4'], + ['webm', 'video/webm'], + ['mov', 'video/quicktime'], + ['pdf', 'application/pdf'], + ['txt', 'text/plain'], + ['md', 'text/markdown'], + ['csv', 'text/csv'], + ['html', 'text/html'], +]) + +/** + * The MIME type of a file name, a path, or a URL path, by its extension. + * Returns `undefined` for no extension or an extension it does not know. + * It knows the common image, audio, video, PDF, and text types. + * + * @example + * mimeTypeOf('cat.PNG') // 'image/png' + * mimeTypeOf('notes') // undefined + */ +export function mimeTypeOf(fileName: string) { + const extension = /\.([^./\\]+)$/.exec(fileName)?.[1] + return extension === undefined + ? undefined + : mimeTypes.get(extension.toLowerCase()) +} + /** * The content part that sends a stored media file to a turn. Its source is * the URL `harness-media:`. The harness swaps it for the bytes only when diff --git a/packages/ai-mcp/src/harness.ts b/packages/ai-mcp/src/harness.ts index eacf38174c..3a048208c4 100644 --- a/packages/ai-mcp/src/harness.ts +++ b/packages/ai-mcp/src/harness.ts @@ -9,6 +9,7 @@ import { isMediaRecord, kindOf, mediaPart, + mimeTypeOf, } from '@tanstack/ai-harness' // The inferred return type names `MCPHandleOptions`. Without an import of // `./server/index`, the .d.ts emit writes `./server.js`, which does not resolve. @@ -79,28 +80,6 @@ const MEDIA_URI_PREFIX = 'harness-media://' /** The biggest image or audio file that a result carries inline. */ const INLINE_MEDIA_MAX_BYTES = 5 * 1024 * 1024 -// ponytail: common types only. Send any other file as data with its mimeType. -const mimeTypes = new Map([ - ['png', 'image/png'], - ['jpg', 'image/jpeg'], - ['jpeg', 'image/jpeg'], - ['gif', 'image/gif'], - ['webp', 'image/webp'], - ['mp3', 'audio/mpeg'], - ['wav', 'audio/wav'], - ['ogg', 'audio/ogg'], - ['m4a', 'audio/mp4'], - ['flac', 'audio/flac'], - ['mp4', 'video/mp4'], - ['webm', 'video/webm'], - ['mov', 'video/quicktime'], - ['pdf', 'application/pdf'], - ['txt', 'text/plain'], - ['md', 'text/markdown'], - ['csv', 'text/csv'], - ['html', 'text/html'], -]) - const threadIdSchema: JSONSchema = { type: 'string', description: @@ -719,14 +698,6 @@ function urlPart(attachment: { url: string; mimeType?: string }) { return part } -/** The MIME type of a file name or URL path by its extension, or `undefined`. */ -function mimeTypeOf(path: string) { - const extension = /\.([^./\\]+)$/.exec(path)?.[1] - return extension === undefined - ? undefined - : mimeTypes.get(extension.toLowerCase()) -} - /** The media records that the settled `operation` published, in order. */ async function mediaOf(operation: Operation) { const records: Array = [] diff --git a/testing/e2e/fixtures/harness/media.json b/testing/e2e/fixtures/harness/media.json new file mode 100644 index 0000000000..8565402a86 --- /dev/null +++ b/testing/e2e/fixtures/harness/media.json @@ -0,0 +1,37 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[harness-media] paint a copy of this dot", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "painter", + "arguments": "{\"prompt\":\"[harness-media-image] a blue dot\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[harness-media] paint a copy of this dot", + "sequenceIndex": 1 + }, + "response": { + "content": "Here is your blue dot." + } + }, + { + "match": { + "userMessage": "[harness-media-image] a blue dot" + }, + "response": { + "image": { + "b64Json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR42mNgYPgPAAEDAQA2dBFAAAAAAElFTkSuQmCC" + } + } + } + ] +} diff --git a/testing/e2e/src/routes/api.harness-protocol.$.ts b/testing/e2e/src/routes/api.harness-protocol.$.ts index 727e618a92..214f6bb3db 100644 --- a/testing/e2e/src/routes/api.harness-protocol.$.ts +++ b/testing/e2e/src/routes/api.harness-protocol.$.ts @@ -11,26 +11,39 @@ import { import { todos } from '@tanstack/ai-harness/plugins' import { memoryPersistence } from '@tanstack/ai-persistence' import { z } from 'zod' +import { createImageAdapter } from '@/lib/media-providers' import { createTextAdapter } from '@/lib/providers' /** * The harness session protocol behind `createHarnessHandler`. The main * model is the OpenAI adapter against aimock. The aimock port and test id - * come from headers, so one handler serves every test. + * come from headers, so one handler serves every test. The model can call + * `painter`, which makes one image with the OpenAI image adapter. The + * harness keeps that image, and the uploads, in memory. */ -const hosts = new Map>() +// One host for every test. Each test uses its own thread ids, and a signed +// media URL comes without the test headers. +const host = createHarnessHost({ persistence: memoryPersistence() }) function handlerFor(request: Request) { const testId = request.headers.get('x-test-id') ?? 'default' const port = Number(request.headers.get('x-aimock-port') ?? '4010') - let host = hosts.get(testId) - if (!host) { - host = createHarnessHost({ persistence: memoryPersistence() }) - hosts.set(testId, host) - } + const painter = defineAgent({ + name: 'painter', + description: 'Paints one image', + inputSchema: z.object({ prompt: z.string() }), + run: async (ctx) => { + await ctx.generateImage({ + adapter: createImageAdapter('openai', port, testId), + prompt: ctx.input.prompt, + }) + return 'Painted one image.' + }, + }) const harness = defineHarness({ name: 'e2e/protocol', adapter: createTextAdapter('openai', undefined, port, testId).adapter, + subagents: { agents: [painter] }, agents: [ defineAgent({ name: 'echo', @@ -68,6 +81,9 @@ function handlerFor(request: Request) { req.headers.get('authorization') === 'Bearer e2e-token' ? { id: 'e2e' } : null, + // Each request makes a new handler, so a fixed secret keeps a signed URL + // working on the next request. + mediaSecret: 'e2e-media-secret', }) } diff --git a/testing/e2e/tests/harness-media.spec.ts b/testing/e2e/tests/harness-media.spec.ts new file mode 100644 index 0000000000..449ba9ef24 --- /dev/null +++ b/testing/e2e/tests/harness-media.spec.ts @@ -0,0 +1,174 @@ +import { test, expect } from './fixtures' +import type { APIRequestContext } from '@playwright/test' +import type { MediaRecord } from '@tanstack/ai-harness' + +const BASE = '/api/harness-protocol' + +// A 1x1 red PNG that the user sends, and the 1x1 blue PNG that the image +// agent makes (the aimock image in fixtures/harness/media.json). +const RED_DOT = Buffer.from( + 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR42mP4z8AAAAMBAQD3A0FDAAAAAElFTkSuQmCC', + 'base64', +) +const BLUE_DOT = Buffer.from( + 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR42mNgYPgPAAEDAQA2dBFAAAAAAElFTkSuQmCC', + 'base64', +) + +/** The media routes of the protocol handler, as the e2e principal. */ +function mediaApi( + request: APIRequestContext, + testId: string, + aimockPort: number, +) { + const headers = { + authorization: 'Bearer e2e-token', + 'x-test-id': testId, + 'x-aimock-port': String(aimockPort), + } + return { + headers, + /** Upload `bytes` as a PNG file of `threadId`. */ + async upload(threadId: string, bytes: Buffer) { + const query = new URLSearchParams({ threadId, name: 'dot.png' }) + const response = await request.post(`${BASE}/media?${query}`, { + headers: { ...headers, 'content-type': 'image/png' }, + data: bytes, + }) + const record: MediaRecord = await response.json() + return record + }, + /** The signed URL of a file, from the media-url route. */ + async signedUrl(record: MediaRecord) { + const query = new URLSearchParams({ + threadId: record.threadId, + id: record.id, + }) + const response = await request.get(`${BASE}/media-url?${query}`, { + headers, + }) + const body: { path: string } = await response.json() + return `${BASE}/${body.path}` + }, + } +} + +test.describe('harness media', () => { + test('serves an uploaded image from its signed URL without auth', async ({ + request, + testId, + aimockPort, + }) => { + const api = mediaApi(request, testId, aimockPort) + const record = await api.upload(`media-${testId}`, RED_DOT) + + const response = await request.get(await api.signedUrl(record)) + + expect(record).toMatchObject({ + threadId: `media-${testId}`, + kind: 'image', + mimeType: 'image/png', + name: 'dot.png', + size: RED_DOT.byteLength, + source: 'user', + }) + expect(response.status()).toBe(200) + expect(await response.body()).toEqual(RED_DOT) + expect(response.headers()).toMatchObject({ + 'content-type': 'image/png', + 'x-content-type-options': 'nosniff', + 'content-security-policy': "default-src 'none'; sandbox", + }) + }) + + test('refuses a signed URL whose file id was changed (403)', async ({ + request, + testId, + aimockPort, + }) => { + const api = mediaApi(request, testId, aimockPort) + const record = await api.upload(`media-sig-${testId}`, RED_DOT) + const other = await api.upload(`media-sig-${testId}`, BLUE_DOT) + const url = new URL(await api.signedUrl(record), 'http://e2e') + + url.searchParams.set('id', other.id) + const response = await request.get(`${url.pathname}${url.search}`) + + expect(response.status()).toBe(403) + }) + + test('a prompt with an uploaded image runs the image agent, and its image loads from a signed URL', async ({ + request, + testId, + aimockPort, + }) => { + const api = mediaApi(request, testId, aimockPort) + const threadId = `media-run-${testId}` + const upload = await api.upload(threadId, RED_DOT) + + // AG-UI text parts carry `text`. The media part has the TanStack shape. + const run = await request.post(`${BASE}/run`, { + headers: { ...api.headers, 'content-type': 'application/json' }, + data: { + threadId, + runId: 'client-run', + messages: [ + { + id: 'u1', + role: 'user', + content: [ + { + type: 'text', + text: '[harness-media] paint a copy of this dot', + }, + { + type: 'image', + source: { + type: 'url', + value: `harness-media:${upload.id}`, + mimeType: 'image/png', + }, + }, + ], + }, + ], + tools: [], + context: [], + state: {}, + forwardedProps: {}, + }, + }) + const events = (await run.text()) + .split('\n') + .filter((line) => line.startsWith('data: ')) + .map((line) => JSON.parse(line.slice(6))) + // The stream also has the painter's own message. The last one is the + // model's answer. + const texts = events.filter( + (event) => event.type === 'TEXT_MESSAGE_CONTENT', + ) + const answer = texts + .filter((event) => event.messageId === texts.at(-1)?.messageId) + .map((event) => event.delta) + .join('') + const painted: MediaRecord = events.find( + (event) => event.type === 'CUSTOM' && event.name === 'harness.media', + )?.value + + expect(answer).toBe('Here is your blue dot.') + expect(painted).toMatchObject({ + threadId, + kind: 'image', + mimeType: 'image/png', + size: BLUE_DOT.byteLength, + source: 'generated', + subagentRunId: expect.stringMatching(/^subagent-/), + }) + + const image = await request.get(await api.signedUrl(painted)) + + expect(image.status()).toBe(200) + expect(image.headers()['content-type']).toBe('image/png') + expect(await image.body()).toEqual(BLUE_DOT) + }) +}) From 44dfbe85db8266e77874a666b259ddc0a162fed8 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 29 Sep 2026 13:17:35 +0200 Subject: [PATCH 063/286] docs(harness): send and show media; example keeps agent media A new page, "Send and show media", covers uploads, signed URLs, media from agents, limits, and the stores. The custom UI, CLI, MCP server, connect, and subagents pages show the media parts, @path attachments, the media folder, MCP attachments, and media links. The example lets the harness keep its images and videos, uses stream: true for video, and shows media lines in the Ink screen. Claude-Session: https://claude.ai/code/session_01APYv1qshKyjPPpkFyRZhfZ --- docs/config.json | 17 +- docs/harness/cli.md | 27 +++ docs/harness/connect.md | 8 +- docs/harness/custom-ui.md | 68 +++++++ docs/harness/mcp-server.md | 29 ++- docs/harness/media.md | 278 ++++++++++++++++++++++++++++ docs/harness/subagents.md | 2 + examples/harness-cli/.gitignore | 2 + examples/harness-cli/README.md | 3 +- examples/harness-cli/src/harness.ts | 11 +- examples/harness-cli/src/media.ts | 104 ++++------- examples/harness-cli/src/tui.tsx | 194 +++++++++++++++---- 12 files changed, 622 insertions(+), 121 deletions(-) create mode 100644 docs/harness/media.md diff --git a/docs/config.json b/docs/config.json index b43bb68369..e8e81c4154 100644 --- a/docs/config.json +++ b/docs/config.json @@ -875,13 +875,19 @@ { "label": "Connect clients", "to": "harness/connect", - "addedAt": "2026-09-26" + "addedAt": "2026-09-26", + "updatedAt": "2026-09-29" + }, + { + "label": "Send and show media", + "to": "harness/media", + "addedAt": "2026-09-29" }, { "label": "Run in the terminal", "to": "harness/cli", "addedAt": "2026-09-26", - "updatedAt": "2026-09-28" + "updatedAt": "2026-09-29" }, { "label": "Durable sessions", @@ -914,7 +920,8 @@ { "label": "Use from any MCP client", "to": "harness/mcp-server", - "addedAt": "2026-09-28" + "addedAt": "2026-09-28", + "updatedAt": "2026-09-29" }, { "label": "Code mode in a harness", @@ -935,13 +942,13 @@ "label": "Build your own UI", "to": "harness/custom-ui", "addedAt": "2026-09-28", - "updatedAt": "2026-09-28" + "updatedAt": "2026-09-29" }, { "label": "Run agents from a harness", "to": "harness/subagents", "addedAt": "2026-09-26", - "updatedAt": "2026-09-28" + "updatedAt": "2026-09-29" }, { "label": "Deploy a harness", diff --git a/docs/harness/cli.md b/docs/harness/cli.md index cc45a4f81b..1aecb74cfa 100644 --- a/docs/harness/cli.md +++ b/docs/harness/cli.md @@ -87,6 +87,32 @@ For the running work: Plugin commands (for example `/model` or `/todos`) show up in `/help`. When a turn stops for an approval or a plugin asks a question, type your answer. For yes-or-no questions, `y` approves and `n` refuses. +## Send files and save media + +To send a file with a message, write `@` and the path of the file: + +```text +what is wrong in @./bug.png +compare @"old logo.png" with @./new-logo.png +``` + +- The path is relative to the working folder. Put quotes around a path with spaces. +- `@path` works in line mode, with `-p`, with piped input, and in your own `ui`. +- A `@word` that is not a file stays in the text, for example `@types/node`. +- If the CLI does not know the type of the file, it does not send the message. To send the path as text, remove the `@`. + +The CLI saves the media that a turn makes in a folder. The default folder is `./-media` in the working folder, for example `./acme-assistant-media`. To use another folder, add `--media-dir `: + +```bash +npx tsx cli.ts -p "Draw a logo for acme" --media-dir ./out +``` + +- Line mode prints `[image saved: ]` for each file. +- `-p` prints the same line on stderr, so stdout keeps only the answer. +- `-p --output ndjson` adds the saved `path` to the value of the `harness.media` event. + +The CLI never replaces a file. If the name is taken, it adds `-1`, `-2`, and so on to the new name. + ## Run your own screen Line mode prints plain lines. For a full screen with your own layout, pass `ui` to `runCli`. It works with any TUI library, for example Ink, OpenTUI, or blessed. @@ -116,5 +142,6 @@ For a full Ink screen with approvals, questions, sign-ins, and child agents, cop - One entry file that runs your harness as a terminal app, a script step, an editor agent, an MCP server, or an HTTP server. - Your own terminal screen on the same session, with any TUI library. +- Files that you send with `@path`, and a folder with the media that your agents make. Next: keep long turns alive through crashes with [durable sessions](./durable-sessions). diff --git a/docs/harness/connect.md b/docs/harness/connect.md index f38633fb99..2173bd9ebc 100644 --- a/docs/harness/connect.md +++ b/docs/harness/connect.md @@ -44,7 +44,7 @@ export const handler = createHarnessHandler({ The handler answers these paths under your route: - `GET capabilities`: the AG-UI capabilities, with the agents in `expose.agents`. -- `POST run`: standard AG-UI. One request runs one prompt and streams it as SSE. Any AG-UI client works. +- `POST run`: standard AG-UI. One request runs one prompt and streams it as SSE. Any AG-UI client works. The prompt keeps every content part of the last user message, for example an image. - `GET events?threadId=`: every event of the session as SSE. Each event id is a cursor, so a reconnect with `Last-Event-ID` continues where it stopped. - `POST control`: send `{ threadId, input }`, for example `{ op: 'prompt', message }`. You get a receipt back. - `GET snapshot?threadId=`: the status, running operations, and waiting approvals. @@ -74,6 +74,8 @@ for await (const entry of client.events()) { `events()` reconnects after a network error and continues from the last cursor. A second tab, a phone, or a reload all see the same session. +To upload files and show the media that your agents make, read [Send and show media](./media). + ## Use a WebSocket For one connection that carries events and inputs, authorize the upgrade, then hand the socket to `handleHarnessSocket`: @@ -99,7 +101,7 @@ import { serveAcp } from '@tanstack/ai-acp/agent' serveAcp({ host, harness: assistant }) ``` -Tool approvals become permission requests in the editor. ACP v2 is still a draft, so this API is experimental. +Tool approvals become permission requests in the editor. Image blocks, audio blocks, and embedded files in a prompt go to the media store of the session. ACP v2 is still a draft, so this API is experimental. ## Use it as the model of another chat @@ -116,7 +118,7 @@ const stream = chat({ }) ``` -Each outer thread gets its own inner session, so the harness keeps its own history. +Each outer thread gets its own inner session, so the harness keeps its own history. The inner session gets every content part of the last user message, so images and files go through too. ## What you have now diff --git a/docs/harness/custom-ui.md b/docs/harness/custom-ui.md index aeedae39de..d243cde990 100644 --- a/docs/harness/custom-ui.md +++ b/docs/harness/custom-ui.md @@ -142,6 +142,7 @@ An assistant message has `parts`. Each part is one of these: - `text` or `reasoning`: the text, which grows while it streams. - `tool-call`: a tool call with its `name`, `args`, and `status` (`running`, `done`, `failed`, or `needs-approval`). - `agent`: a child agent, with its own `parts`. +- `media`: a file that an agent made. See [Show media](#show-media). A notice has a `kind`: @@ -195,6 +196,72 @@ Progress: The view calls your handler after the item is in the state, so the handler can read `view.store.get()` and find the item. `view.on` returns a function that removes the handler. +## Show media + +A user sends a screenshot, and an agent draws an image. The view gives you each file as a media part, with a URL that works in ``, `